Vector Search using Bigquery

Vector search provided by Google Cloud services 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 GCP. These are suitable for production use cases.

Installation

npm install @genkit-ai/vertexai

Configuration

Create a vector search index in GCP. Details on creating vector search index can be found at Create your Vector Search Index
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 the GCP 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,
    };
}