Examples › RAG (Gemini embeddings) — vector similarity
RAG example

RAG over a local knowledge base — Gemini embeddings

The same seven-vertex graph as rag-bm25, with the retriever swapped from BM25 to Gemini embeddings + cosine similarity. The retriever calls library.ResolveEmbeddingClient — never reads GEMINI_API_KEY on its own — so a registered EmbeddingClientFactory determines where the credential actually comes from.

What this example teaches

  • A library.Retriever that builds and queries an in-memory vector index at startup
  • Using library.ResolveEmbeddingClient(ctx, provider, model) instead of reading env vars directly
  • How EmbeddingClientFactory decouples the retriever from credential sourcing
  • Cosine similarity ranking over gemini-embedding-001 vectors
  • The same graph shape works for any embedding provider once a factory is registered
View source on GitHub
Same graph

The seven vertices are identical

Swap the retriever, leave the graph alone. The downstream prompt assembly, citation parsing, and citation validation work the same regardless of how documents were ranked. That's the value of the Retriever abstraction.

  1. question_const → question
  2. RetrieveOp (k=3, vector backend) → documents, texts
  3. BuildRAGPromptOp → prompt
  4. RetrievedSourcesOp → retrieved_sources
  5. AIComputeStringToStringOp → raw_answer
  6. ParseCitationsOp → body, sources
  7. ValidateCitationsOp → accepted_sources, rejected_sources
Implementation

Walkthrough

1. A Retriever backed by in-memory vectors

At startup, the example embeds every .txt passage with library.ResolveEmbeddingClient(ctx, "gemini", "gemini-embedding-001") and stores vectors in memory. Each request embeds the query the same way, then ranks the corpus by cosine similarity. The retriever never touches GEMINI_API_KEY directly — credentials flow through the registered factory.

embed_retriever.go (sketch) 
func (r *EmbedRetriever) Retrieve(ctx context.Context, q string, k int) ([]library.Document, error) {
    client, err := library.ResolveEmbeddingClient(ctx, "gemini", "gemini-embedding-001")
    if err != nil {
        return nil, err
    }
    vecs, err := client.Embed(ctx, []string{q})
    if err != nil { return nil, err }

    return r.topKByCosine(vecs[0], k), nil
}

2. Registering the embedding factory

The bundled EnvEmbeddingClientFactory supports provider="gemini" out of the box via GEMINI_API_KEY. To plug in Voyage, OpenAI, Cohere, or anything else, register your own factory:

main.go 
// uses the bundled gemini-only factory — no extra wiring
docs, _ := loadKB("testdata/kb")
retriever, _ := NewEmbedRetriever(context.Background(), docs)
library.SetDefaultRetriever(retriever)

// to use a different embedding provider:
// library.SetDefaultEmbeddingClientFactory(&VoyageFactory{...})
// or library.RegisterEmbeddingClientFactory("voyage-prod", ...)
// and set client_factory_id on the RetrieveOp vertex.

Run it

shell 
export CLAUDE_API_KEY=<your key>
export GEMINI_API_KEY=<your key>  # for embeddings

go run ./examples/rag-gemini-embed --question "How is data encrypted at rest?"

Run it as an MCP server

The same binary is dual-mode. Pass --mcp and instead of running once and exiting it speaks the Model Context Protocol over stdin/stdout, exposing the entire workflow as a single MCP tool, answer_from_kb. The Go SDK derives the tool's input schema from the workflow's UserInput struct and validates every tools/call request against it. The KB is embedded once at server startup, so every subsequent tool call reuses the in-memory Gemini vectors.

shell 
# Speak MCP over stdin/stdout instead of running once
go run ./examples/rag-gemini-embed --mcp

Register it with any MCP client by pointing the client at that command:

.mcp.json 
{
  "mcpServers": {
    "rag-gemini-embed": {
      "command": "go",
      "args": ["run", "./examples/rag-gemini-embed", "--mcp"]
    }
  }
}
← RAG (BM25) Next: WithRepair →