Vertex AI Vector Search with Firestore

Vertex AI Vector Search allows you to index and retrieve documents. The documents are stored in Firestore 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 vector search index in Vertex AI. Details on creating vector search index can be found at Create your Vector Search Index
2. Create a Firestore Dataset and a Collection within that dataset to store the documents that will be indexed. More information to create Firestore datasets is available here

What the Firestore helpers write

GetFirestoreDocumentIndexer writes one auto-ID document per ai.Document into the named collection, with two fields: content (the document’s Content parts) and metadata. It commits them in a single batch and returns the generated document IDs, which are the IDs the vector search index stores as datapoints. GetFirestoreDocumentRetriever reads those documents back by ID.

The collection needs no vector index. The embeddings live in the Vertex AI index, not in Firestore; Firestore only holds the document bodies.

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

To use the GCP vector search with Firestore, 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 Firestore dataset:

import (
    "context"
    "log"

    "cloud.google.com/go/firestore"
    "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"
)

ctx := context.Background()

vectorsearchPlugin := &vectorsearch.VertexAIVectorSearch{
    ProjectID: "${GOOGLE_CLOUD_PROJECT_ID}",
    Location:  "${GOOGLE_CLOUD_PROJECT_LOCATION}",
}

g := genkit.Init(ctx, genkit.WithPlugins(
    &googlegenai.VertexAI{},
    vectorsearchPlugin,
))

databaseId := "${FIRESTORE_DATABASE_ID}"         // Replace with your Firestore database ID
collectionName := "${FIRESTORE_COLLECTION_NAME}" // Replace with your Firestore collection name

firestoreClient, err := firestore.NewClientWithDatabase(ctx, vectorsearchPlugin.ProjectID, databaseId)
if err != nil {
    log.Fatalf("failed to create Firestore client: %v", err)
}
defer firestoreClient.Close()

documentIndexer := vectorsearch.GetFirestoreDocumentIndexer(firestoreClient, collectionName)
documentRetriever := vectorsearch.GetFirestoreDocumentRetriever(firestoreClient, collectionName)

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
}

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 Firestore. 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 Firestore. 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 Firestore
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 Firestore. This function retrieves the docs corresponding to the Neighbor IDs
// found using vector search index.

resp, err := retriever.Retrieve(ctx, &ai.RetrieverRequest{
    Query: ai.DocumentFromText("How do I make a perfect espresso?", 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.