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

The vector search functionality is built into Genkit Go. You need to import the vectorsearch package:

import "github.com/firebase/genkit/go/plugins/vertexai/vectorsearch"

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

The BigQuery table schema

GetBigQueryDocumentIndexer and GetBigQueryDocumentRetriever are hard-coded to three STRING columns named id, content and metadata. Create the table with exactly that shape:

bq mk --table your-project-id:your-dataset-id.your-table-id \
  id:STRING,content:STRING,metadata:STRING

content and metadata hold the JSON encoding of ai.Document.Content and ai.Document.Metadata as strings, not BigQuery JSON values. The indexer generates a random hex id per document and returns the ids it wrote; the retriever reads the rows back with SELECT id, content, metadata FROM ... WHERE id IN UNNEST(@ids) using the neighbor IDs the vector index returns. A column named or typed differently fails at insert or at read.

Write your own vectorsearch.DocumentIndexer and vectorsearch.DocumentRetriever pair if you need a different schema.

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 (
    "context"
    "log"

    "cloud.google.com/go/bigquery"
    "github.com/firebase/genkit/go/ai"
    "github.com/firebase/genkit/go/genkit"
    "github.com/firebase/genkit/go/plugins/googlegenai"
    "github.com/firebase/genkit/go/plugins/vertexai/vectorsearch"
)

// Define your own config struct to hold all parameters
type VectorsearchConfig struct {
    ProjectID         string
    Location          string
    IndexID           string
    IndexEndpointID   string
    DeployedIndexID   string
    ProjectNumber     string
    PublicDomainName  string
    Embedder          ai.Embedder
    NeighborsCount    int
    DocumentIndexer   vectorsearch.DocumentIndexer
    DocumentRetriever vectorsearch.DocumentRetriever
}

ctx := context.Background()

// Initialize the Vector Search plugin
vectorsearchPlugin := &vectorsearch.VertexAIVectorSearch{
    ProjectID: "your-project-id",
    Location:  "us-central1",
}

// Initialize Genkit with both plugins
g := genkit.Init(ctx, genkit.WithPlugins(
    &googlegenai.VertexAI{},
    vectorsearchPlugin,
))

bqClient, err := bigquery.NewClient(ctx, "your-project-id")
if err != nil {
    log.Fatalf("Failed to create BigQuery client: %v", err)
}

documentIndexer := vectorsearch.GetBigQueryDocumentIndexer(bqClient, "your-dataset-id", "your-table-id")
documentRetriever := vectorsearch.GetBigQueryDocumentRetriever(bqClient, "your-dataset-id", "your-table-id")

vectorsearchParams := &VectorsearchConfig{
    ProjectID:         vectorsearchPlugin.ProjectID,
    Location:          vectorsearchPlugin.Location,
    IndexID:           "${VECTOR_SEARCH_INDEX_ID}",                           // Replace with your index ID
    IndexEndpointID:   "${VECTOR_SEARCH_INDEX_ENDPOINT_ID}",                  // Replace with your index endpoint ID
    DeployedIndexID:   "${VECTOR_SEARCH_DEPLOYED_INDEX_ID}",                  // Replace with your deployed index ID
    ProjectNumber:     "${GOOGLE_CLOUD_PROJECT_NUMBER}",                      // Replace with your Google Cloud project number
    PublicDomainName:  "${VECTOR_SEARCH_PUBLIC_DOMAIN_NAME}",                 // Replace with your public domain name
    Embedder:          googlegenai.VertexAIEmbedder(g, "text-embedding-004"), // Replace with your desired embedder
    NeighborsCount:    10,                                                    // Number of neighbors to retrieve
    DocumentIndexer:   documentIndexer,
    DocumentRetriever: documentRetriever,
}

Values you need to collect

VectorsearchConfig above is a plain local struct, not a Genkit type. It exists only to carry these values from configuration to the call sites below; the plugin never sees it. Name it whatever you like.

• 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.
• ProjectNumber (string): the numeric project ID, not the project ID string. Get it with gcloud projects describe $PROJECT_ID --format='value(projectNumber)', or read it from the Cloud console project picker.
• PublicDomainName (string): the publicEndpointDomainName of the deployed index endpoint. Get it with gcloud ai index-endpoints describe $INDEX_ENDPOINT_ID --region=$LOCATION --format='value(publicEndpointDomainName)'.
• Embedder (ai.Embedder): The embedding model to use. Must be a configured embedder in your Genkit project.
• NeighborsCount (int): Number of neighbors to set in the vector search
• 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.

Caution

The plugin queries the public findNeighbors endpoint at https://{PublicDomainName}/v1/projects/{ProjectNumber}/locations/{Location}/indexEndpoints/{IndexEndpointID}:findNeighbors. publicEndpointDomainName exists only for endpoints deployed with public endpoint access, so index endpoints deployed behind VPC peering or Private Service Connect are not supported.

Genkit types

These are the types the plugin actually defines.

Type	Fields
vectorsearch.VertexAIVectorSearch	ProjectID string, Location string. This is the plugin value you pass to genkit.Init.
vectorsearch.Config	IndexID string. The only field, so a retriever definition is vectorsearch.DefineRetriever(ctx, g, vectorsearch.Config{IndexID: id}, nil).
vectorsearch.IndexParams	Docs []*ai.Document, Embedder ai.Embedder, EmbedderOptions any, ProjectID string, Location string, IndexID string.
vectorsearch.RetrieveParams	Content *ai.Document, Embedder ai.Embedder, EmbedderOptions any, AuthClient *google.Credentials, ProjectNumber string, Location string, IndexEndpointID string, PublicDomainName string, DeployedIndexID string, NeighborCount int, Restricts []Restrict, NumericRestricts []NumericRestrict, DocumentRetriever DocumentRetriever.

The trailing nil in DefineRetriever is an *ai.RetrieverOptions. Pass one to set the retriever’s label or declare what it supports; nil takes the defaults.

RetrieveParams.Location is ignored: the plugin uses the Location you set on VertexAIVectorSearch.

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:

import (
    "github.com/firebase/genkit/go/ai"
)

// Create documents from text
data := []string{
    "This is the first document.",
    "This is the second document.",
    "This is the third document.",
    "This is the fourth document.",
}

var docs []*ai.Document
for _, text := range data {
    docs = append(docs, ai.DocumentFromText(text, nil))
}

// Index the docs.
// Custom Index function can be used which should internally refer the indexer function for Bigquery
if err := vectorsearch.Index(ctx, g, vectorsearch.IndexParams{
    IndexID:         vectorsearchParams.IndexID,
    Embedder:        vectorsearchParams.Embedder,
    EmbedderOptions: nil,
    Docs:            docs,
    ProjectID:       vectorsearchParams.ProjectID,
    Location:        vectorsearchParams.Location,
}, vectorsearchParams.DocumentIndexer); err != nil {
    return nil, err
}

Retrieving Documents

Call Retrieve on the retriever you defined:

// Define the retriever for vector search.
retriever, err := vectorsearch.DefineRetriever(ctx, g, vectorsearch.Config{
    IndexID: vectorsearchParams.IndexID, // Replace with your index ID
}, nil)
if err != nil {
    log.Fatal(err)
}

// The retriever defined above has built in function called Retrieve() which
// corresponds to vector search retriever function defined in vector search plugin.
// The DocumentRetriever passed as argument corresponds to the documentretriever
// for Bigquery. This function retrieves the docs corresponding to the Neighbor IDs
// found using vector search index.

question := "Your search query"

resp, err := retriever.Retrieve(ctx, &ai.RetrieverRequest{
    Query: ai.DocumentFromText(question, nil),
    Options: &vectorsearch.RetrieveParams{
        Embedder:          vectorsearchParams.Embedder,
        NeighborCount:     vectorsearchParams.NeighborsCount,
        IndexEndpointID:   vectorsearchParams.IndexEndpointID,
        DeployedIndexID:   vectorsearchParams.DeployedIndexID,
        PublicDomainName:  vectorsearchParams.PublicDomainName,
        ProjectNumber:     vectorsearchParams.ProjectNumber,
        DocumentRetriever: vectorsearchParams.DocumentRetriever,
    }})
if err != nil {
    return nil, err
}

Which Retrieve to call

Three forms exist and the other vector store pages use a different one, so:

• retriever.Retrieve(ctx, req) is the ai.Retriever interface method. It is used here because the request carries a typed Options payload, *vectorsearch.RetrieveParams, that the vector search retriever needs on every call.
• genkit.Retrieve(ctx, g, ai.WithRetriever(r), ai.WithTextDocs(q)) is the form the other pages use. It is equivalent for retrievers that need no options. To pass options through it, add ai.WithConfig(&vectorsearch.RetrieveParams{...}).
• ai.Retrieve(ctx, reg, opts...) is the same call taking an api.Registry rather than a *genkit.Genkit. Use it from plugin code that has no *genkit.Genkit.