Vertex AI Vector Search with BigQuery

Vertex AI Vector Search allows you to index and retrieve documents. The documents are stored in Bigquery and the corresponding document IDs are indexed using the vector search index provided by Vertex AI. These are suitable for production use cases.

Installation

npm install @genkit-ai/vertexai

Configuration

1. Create a Vertex AI Vector Search index. Details on creating an index can be found at Create your Vector Search Index
2. Create a Bigquery Dataset and a Table within that dataset to store the documents that will be indexed. More information to create Bigquery datasets is available here

To use Vertex AI Vector Search with Bigquery, initialize it and define a retriever with an embedder. You can also use a custom indexer and retriever for indexing and retrieving documents from the Bigquery dataset:

import { BigQuery } from '@google-cloud/bigquery';

const bq = new BigQuery({
  projectId: PROJECT_ID,
});

const bigQueryDocumentRetriever: DocumentRetriever =
  getBigQueryDocumentRetriever(bq, BIGQUERY_TABLE, BIGQUERY_DATASET);

const bigQueryDocumentIndexer: DocumentIndexer = getBigQueryDocumentIndexer(
  bq,
  BIGQUERY_TABLE,
  BIGQUERY_DATASET,
);

// Configure Genkit with Vertex AI plugin
const ai = genkit({
  plugins: [
    vertexAI({
      projectId: PROJECT_ID,
      location: LOCATION,
      googleAuth: {
        scopes: ['https://www.googleapis.com/auth/cloud-platform'],
      },
    }),
    vertexAIVectorSearch({
      location: LOCATION,
      projectId: PROJECT_ID,
      embedder: textEmbedding004,
      vectorSearchOptions: [
        {
          publicDomainName: VECTOR_SEARCH_PUBLIC_DOMAIN_NAME,
          indexEndpointId: VECTOR_SEARCH_INDEX_ENDPOINT_ID,
          indexId: VECTOR_SEARCH_INDEX_ID,
          deployedIndexId: VECTOR_SEARCH_DEPLOYED_INDEX_ID,
          documentRetriever: bigQueryDocumentRetriever,
          documentIndexer: bigQueryDocumentIndexer,
        },
      ],
    }),
  ],
});

Configuration Options

• projectId (string): GCP Project ID
• location (string): GCP Project location
• indexId (string): Vector search index id
• indexEndpointId (string): Vector search endpoint id corresponding to the vector search index. More details can be found here.
• deployedIndexId (string): Vector search deployed index id corresponding to the vector search endpoint. More details to deploy an index to an index endpoint can be found here.
• publicDomainName (string): Public Domain Name of the vector search index endpoint.
• embedder (ai.Embedder): The embedding model to use. Must be a configured embedder in your Genkit project.
• documentIndexer (func(ctx context.Context, docs []*ai.Document) ([]string, error)): Document indexer used to insert data with unique IDs in Bigquery. This can be a custom document indexer as well depending on the user’s requirement.
• documentRetriever (func(ctx context.Context, neighbors []Neighbor, options any) ([]*ai.Document, error)): Document retriever used to retrieve data with corresponding ID from Bigquery. This can be a custom document retriever as well depending on the user’s requirement.

Usage

Indexing Documents

To populate with data, you need to implement your own indexing logic using the ai.Document format. Genkit provides a sample indexing function as well:

async ({ texts }) => {
  const documents = texts.map((text) => Document.fromText(text));
  await ai.index({
    indexer: vertexAiIndexerRef({
      indexId: VECTOR_SEARCH_INDEX_ID,
      displayName: 'bigquery_index',
    }),
    documents,
  });
  return { result: 'success' };
};

Retrieving Documents

Use ai.Retrieve with the retriever you defined:

async ({ query, k }) => {
  const startTime = performance.now();
  const queryDocument = Document.fromText(query);
  const res = await ai.retrieve({
    retriever: vertexAiRetrieverRef({
      indexId: VECTOR_SEARCH_INDEX_ID,
      displayName: 'bigquery_index',
    }),
    query: queryDocument,
    options: { k },
  });
  const endTime = performance.now();
  return {
    result: res
      .map((doc) => ({
        text: doc.content[0].text!,
        distance: doc.metadata?.distance,
      }))
      .sort((a, b) => b.distance - a.distance),
    length: res.length,
    time: endTime - startTime,
  };
};