RAG over a local knowledge base — BM25

1. Register a Retriever before Run

Implement library.Retriever, build it from the loaded knowledge base, and install it as the process default. Each Document is tagged with library.MetadataSource = "filename.txt" so the prompt builder and the validator can both reference the same identifier.

docs, _ := loadKB("testdata/kb") // each Document has Metadata[MetadataSource]
library.SetDefaultRetriever(NewBM25Retriever(docs))

2. Wire RetrieveOp into the graph

The retrieval vertex exposes two parallel outputs: Documents (full records) and Texts (just the content strings). Two downstream custom ops read Documents — one to build the LLM prompt, the other to derive the citation allow-list.

Vertex("retrieve").Op("RetrieveOp").
    Params(map[string]string{"k": "3"}).
    Input("Query", "question").
    Output("Documents", "documents").
    Output("Texts", "texts").

Vertex("format_prompt").Op("BuildRAGPromptOp").
    Input("Question", "question").
    Input("Documents", "documents").
    Output("Prompt", "prompt").

Vertex("retrieved_sources").Op("RetrievedSourcesOp").
    Input("Documents", "documents").
    Output("Sources", "retrieved_sources").

3. Prompt-injection-resistant prompt assembly

BuildRAGPromptOp wraps each retrieved passage in <passage source="...">...</passage>. Passage content is XML-escaped so an attacker-controlled KB document cannot close its own tag and inject instructions. The surrounding prose tells the model to treat anything inside <passage> as untrusted data. The prompt instructs the model to end with "Sources: file1.txt, file2.txt".

4. Validate citations at the boundary

After ParseCitationsOp extracts the source list from the model's response, ValidateCitationsOp filters them against the allow-list of identifiers actually present in retrieved documents — built from retrieved_sources, not the full loaded corpus. A model that hallucinates the name of a real-but-unretrieved KB file would otherwise slip past the check.

Vertex("validate_citations").Op("ValidateCitationsOp").
    Input("Raw",     "sources").
    Input("Allowed", "retrieved_sources").
    Output("Accepted", "accepted_sources").
    Output("Rejected", "rejected_sources").

The driver surfaces Accepted to the user and slog-warns each entry in Rejected — observability, not a graph failure.

Run it

export CLAUDE_API_KEY=<your key>

go run ./examples/rag-bm25 --question "What does the deployment guide say about rollouts?"
go run ./examples/rag-bm25 --question "How is data encrypted at rest?" --kb /path/to/your/kb

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, so an agent runs the whole retrieve → prompt → answer → cite pipeline with one tool call.

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

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

{
  "mcpServers": {
    "rag-bm25": {
      "command": "go",
      "args": ["run", "./examples/rag-bm25", "--mcp"]
    }
  }
}