AlloyDB for PostgreSQL

The AlloyDB plugin provides provides the retriever implementation to search a AlloyDB database using the pgvector extension.

Configuration

To use this plugin, follow these steps:

Import the plugin

import "github.com/firebase/genkit/go/plugins/alloydb"

Create a PostgresEngine instance:
- Using basic authentication

pEngine, err := alloydb.NewPostgresEngine(ctx,
        alloydb.WithUser("user"),
        alloydb.WithPassword("password"),
        alloydb.WithAlloyDBInstance("my-project", "us-central1", "my-cluster", "my-instance"),
        alloydb.WithDatabase("my-database"))

Using email authentication

pEngine, err := alloydb.NewPostgresEngine(ctx,
    alloydb.WithAlloyDBInstance("my-project", "us-central1", "my-cluster", "my-instance"),
    alloydb.WithDatabase("my-database"),
    alloydb.WithIAMAccountEmail("mail@company.com"))

Using custom pool

pool, err := pgxpool.New(ctx, "add_your_connection_string")
if err != nil {
    return err
}

pEngine, err := alloydb.NewPostgresEngine(ctx,
    alloydb.WithDatabase("db_test"),
    alloydb.WithPool(pool))

Create the Postgres plugin
- Using the genkit method init

postgres := &alloydb.Postgres{
    engine: pEngine,
}

g := genkit.Init(ctx, genkit.WithPlugins(postgres))

Usage

To add documents to a AlloyDB index, first create an index definition that specifies the features of the table:

cfg := &alloydb.Config{
    TableName:             "documents",
    SchemaName:            "public",
    ContentColumn:         "content",
    EmbeddingColumn:       "embedding",
    MetadataColumns:       []string{"source", "category"},
    IDColumn:              "custom_id",
    MetadataJSONColumn:    "custom_metadata",
    Embedder:              embedder,
    EmbedderOptions:       nil,
}

doc, retriever, err := alloydb.DefineRetriever(ctx, g, postgres, cfg)
if err != nil {
  return err
}

docs := []*ai.Document{{
        Content: []*ai.Part{{
        Kind:        ai.PartText,
        ContentType: "text/plain",
        Text:        "The product features include...",
        }},
      Metadata: map[string]any{"source": "website", "category": "product-docs", "custom_id": "doc-123"},
  }}

if err := doc.Index(ctx, docs); err != nil {
    return err
}

Similarly, to retrieve documents from an index, use the retriever method:

doc,  retriever, err := alloydb.DefineRetriever(ctx, g, postgres, cfg)
if err != nil {
  return err
}

d2 := ai.DocumentFromText("The product features include...", nil)

resp, err := retriever.Retrieve(ctx, &ai.RetrieverRequest{
    Query: d2,
    Options: &alloydb.RetrieverOptions{
        K: 5,
        Filter: "source='website' AND category='product-docs'",
    },
})

if err != nil {
    return err
}

It’s also possible to use the Retrieve method from genkit

d2 := ai.DocumentFromText("The product features include...", nil)

retrieverOptions := &alloydb.RetrieverOptions{
    K:      5,
    Filter: "source='website' AND category='product-docs'",
}

resp, err := genkit.Retrieve(ctx, g, ai.WithRetriever(retriever), ai.WithDocs(d2), ai.WithConfig(retrieverOptions))
if err != nil {
    return err
}

See the Retrieval-augmented generation page for a general discussion on using retrievers for RAG.