Writing

Building a RAG Pipeline with Genkit Python (End-to-End)

A knowledge Q&A system over a small corpus of internal documents. It:

Every AI project eventually hits the same wall: the model doesn’t know your data. RAG (Retrieval-Augmented Generation) is the standard fix — you embed your documents, find the relevant ones at query time, and stuff them into the prompt. Simple idea, surprisingly sharp edges in practice.

This article walks you through a complete RAG pipeline using Genkit Python: real embeddings via the Google AI API, an in-memory vector store you can swap for anything, source citations in the output, and a retrieval quality sanity check. No hand-waving.

One honest caveat up front: as of the current Genkit Python SDK (main branch, June 2026), there are no built-in define_retriever / define_indexer abstractions yet — those are coming and already exist in the JS SDK. What does exist is a first-class embedder API. I’ll show you the manual pattern for retrieval; the plumbing is thin enough that you won’t miss the abstraction.

What You’re Building

A knowledge Q&A system over a small corpus of internal documents. It:

  1. Indexes documents by chunking and embedding them
  2. Accepts user queries, embeds the query, and finds the top-K matching chunks
  3. Builds an augmented prompt and generates a grounded answer
  4. Returns the answer with source citations
  5. Runs a lightweight overlap metric to sanity-check retrieval quality

Setup

uv venv --python 3.12 .venv && source .venv/bin/activate
uv pip install "genkit[google-genai] @ git+https://github.com/firebase/genkit.git#subdirectory=py/packages/genkit"
uv pip install numpy  # only needed for cosine similarity helper

Set your API key:

export GEMINI_API_KEY="your-key-here"

The Embedder API

The SDK exposes two methods: ai.embed() for a single document and ai.embed_many() for batching. Both return list[Embedding] where each Embedding has a .embedding: list[float] field.

from genkit import Genkit
from genkit.plugins.google_genai import GoogleAI

ai = Genkit(
    plugins=[GoogleAI()],
    model='googleai/gemini-2.0-flash',
)

async def get_embedding(text: str) -> list[float]:
    """Embed a single string and return the raw vector."""
    results = await ai.embed(
        embedder='googleai/gemini-embedding-001',
        content=text,
    )
    return results[0].embedding  # list[float], 3072 dims

gemini-embedding-001 produces 3072-dimensional vectors. If you need lower dimensionality, text-embedding-004 gives you 768 dims — useful when you’re storing millions of chunks and RAM matters.

Step 1: Chunking Strategy

Chunking is where most RAG implementations quietly fail. Too large a chunk → context dilution. Too small → loss of coherent meaning.

from dataclasses import dataclass

@dataclass
class Chunk:
    text: str
    source: str          # document name / URL
    chunk_index: int
    embedding: list[float] | None = None


def chunk_document(text: str, source: str, size: int = 400, overlap: int = 80) -> list[Chunk]:
    """
    Fixed-size character chunking with overlap.

    size=400 chars ≈ 80–100 tokens for English text, which sits comfortably
    inside gemini-embedding-001's 2048-token context window while leaving
    headroom for the query when you later build the RAG prompt.
    """
    chunks = []
    start = 0
    idx = 0
    while start < len(text):
        end = start + size
        chunk_text = text[start:end].strip()
        if chunk_text:
            chunks.append(Chunk(text=chunk_text, source=source, chunk_index=idx))
        start += size - overlap
        idx += 1
    return chunks

The overlap is load-bearing. Without it, a sentence split across a chunk boundary loses coherence and the embedding for both halves degrades.

Step 2: Indexing — Embed and Store

import asyncio
import math

class InMemoryVectorStore:
    """
    Minimal in-memory vector store. Replace the inner list with
    Chroma, Pinecone, pgvector, or AlloyDB for production.
    """
    def __init__(self):
        self._chunks: list[Chunk] = []

    def add(self, chunk: Chunk) -> None:
        assert chunk.embedding is not None
        self._chunks.append(chunk)

    def search(self, query_embedding: list[float], top_k: int = 3) -> list[tuple[Chunk, float]]:
        """Return top-K chunks by cosine similarity."""
        scored = [
            (chunk, _cosine_similarity(query_embedding, chunk.embedding))
            for chunk in self._chunks
            if chunk.embedding is not None
        ]
        scored.sort(key=lambda x: x[1], reverse=True)
        return scored[:top_k]


def _cosine_similarity(a: list[float], b: list[float]) -> float:
    dot = sum(x * y for x, y in zip(a, b))
    norm_a = math.sqrt(sum(x * x for x in a))
    norm_b = math.sqrt(sum(x * x for x in b))
    return dot / (norm_a * norm_b) if norm_a and norm_b else 0.0


async def index_documents(docs: dict[str, str], store: InMemoryVectorStore) -> None:
    """
    Chunk all documents, batch-embed, and store.
    We embed in batches of 20 to stay under API rate limits.
    """
    all_chunks: list[Chunk] = []
    for source, text in docs.items():
        all_chunks.extend(chunk_document(text, source=source))

    print(f"Indexing {len(all_chunks)} chunks from {len(docs)} documents...")

    # Batch embed — ai.embed_many() sends a single API call per batch
    BATCH_SIZE = 20
    for i in range(0, len(all_chunks), BATCH_SIZE):
        batch = all_chunks[i:i + BATCH_SIZE]
        embeddings = await ai.embed_many(
            embedder='googleai/gemini-embedding-001',
            content=[c.text for c in batch],
        )
        for chunk, emb in zip(batch, embeddings):
            chunk.embedding = emb.embedding
            store.add(chunk)
        # Polite pause between batches; remove in production with proper retry logic
        if i + BATCH_SIZE < len(all_chunks):
            await asyncio.sleep(0.5)

    print(f"Done. {len(store._chunks)} chunks indexed.")

Step 3: The Retriever Interface

Since Genkit Python doesn’t yet have a define_retriever abstraction, we wrap retrieval in a plain async function. When the abstraction lands, the swap will be mechanical.

from pydantic import BaseModel

class RetrievedContext(BaseModel):
    text: str
    source: str
    score: float


async def retrieve(
    query: str,
    store: InMemoryVectorStore,
    top_k: int = 3,
) -> list[RetrievedContext]:
    """Embed query, search store, return ranked context with provenance."""
    query_embedding = await get_embedding(query)
    results = store.search(query_embedding, top_k=top_k)
    return [
        RetrievedContext(text=chunk.text, source=chunk.source, score=score)
        for chunk, score in results
    ]

Step 4: The Complete RAG Flow

class RAGInput(BaseModel):
    query: str
    top_k: int = 3


class RAGOutput(BaseModel):
    answer: str
    sources: list[str]


def build_rag_prompt(query: str, context_chunks: list[RetrievedContext]) -> str:
    context_text = "\n\n---\n\n".join(
        f"[Source: {c.source}]\n{c.text}" for c in context_chunks
    )
    return f"""Answer the question below using ONLY the provided context.
If the context doesn't contain enough information, say so clearly.
At the end, list the sources you used.

CONTEXT:
{context_text}

QUESTION: {query}

Answer concisely and cite your sources inline like [Source: name]."""


@ai.flow()
async def rag_flow(input: RAGInput) -> RAGOutput:
    # Step 1: Retrieve relevant chunks
    context_chunks = await retrieve(input.query, store, top_k=input.top_k)

    if not context_chunks:
        return RAGOutput(answer="No relevant documents found.", sources=[])

    # Step 2: Build augmented prompt
    prompt = build_rag_prompt(input.query, context_chunks)

    # Step 3: Generate answer
    response = await ai.generate(prompt=prompt)

    # Step 4: Extract cited sources (deduped, preserving order)
    seen = set()
    sources = []
    for c in context_chunks:
        if c.source not in seen:
            seen.add(c.source)
            sources.append(c.source)

    return RAGOutput(answer=response.text, sources=sources)

Step 5: Wiring It All Together

# Sample corpus — swap for your real documents
DOCUMENTS = {
    "incident-runbook-v2.md": """
        When an alert fires, first check the dashboard at /metrics.
        Acknowledge the alert within 5 minutes to prevent escalation.
        For P0 incidents, page the on-call engineer immediately.
        All incidents must be logged in the incident tracker with a severity level.
        Post-mortems are required for P0 and P1 incidents within 48 hours.
    """,
    "deployment-guide.md": """
        Deployments run every Tuesday and Thursday at 2 PM UTC.
        All deployments require a passing CI pipeline before proceeding.
        Use the deploy script: ./scripts/deploy.sh --env production.
        Rollbacks are triggered automatically if error rate exceeds 5% within 10 minutes.
        Never deploy on Fridays. The deploy freeze period is December 23–January 2.
    """,
    "api-rate-limits.md": """
        The public API enforces 1000 requests per minute per API key.
        Burst allowance is 200 requests in any 10-second window.
        Rate limit headers: X-RateLimit-Remaining and X-RateLimit-Reset.
        Exceeding limits returns HTTP 429. Implement exponential backoff.
        Enterprise plans have configurable limits — contact sales@example.com.
    """,
}

store = InMemoryVectorStore()


async def main():
    # Index documents (do this once; in production, persist to your vector DB)
    await index_documents(DOCUMENTS, store)

    # Run RAG queries
    queries = [
        "What do I do when a P0 alert fires?",
        "When are deployments scheduled and what triggers a rollback?",
        "What HTTP status code does the API return when I exceed rate limits?",
    ]

    for query in queries:
        print(f"\n{'='*60}")
        print(f"QUERY: {query}")
        result = await rag_flow(RAGInput(query=query))
        print(f"\nANSWER: {result.answer}")
        print(f"SOURCES: {', '.join(result.sources)}")


if __name__ == '__main__':
    ai.run_main(main())

Step 6: Evaluating Retrieval Quality

Before you tune chunking strategy or embedding model, you need a metric. Start with token overlap recall: for each query with a known relevant source, did the retriever surface it in the top-K?

from dataclasses import dataclass

@dataclass
class EvalCase:
    query: str
    expected_source: str  # at least one chunk from this source should appear in top-K


async def evaluate_retrieval(
    cases: list[EvalCase],
    store: InMemoryVectorStore,
    top_k: int = 3,
) -> dict:
    hits = 0
    results = []

    for case in cases:
        retrieved = await retrieve(case.query, store, top_k=top_k)
        retrieved_sources = {r.source for r in retrieved}
        hit = case.expected_source in retrieved_sources
        hits += int(hit)
        results.append({
            "query": case.query,
            "expected": case.expected_source,
            "retrieved": list(retrieved_sources),
            "hit": hit,
            "top_score": retrieved[0].score if retrieved else 0.0,
        })

    recall_at_k = hits / len(cases)
    print(f"\nRetrieval recall@{top_k}: {recall_at_k:.2%} ({hits}/{len(cases)})")
    for r in results:
        status = "✓" if r["hit"] else "✗"
        print(f"  {status} [{r['top_score']:.3f}] {r['query'][:50]}")

    return {"recall_at_k": recall_at_k, "results": results}


# Run eval
EVAL_CASES = [
    EvalCase("What happens when a P0 fires?", "incident-runbook-v2.md"),
    EvalCase("When can I deploy to production?", "deployment-guide.md"),
    EvalCase("How do I handle HTTP 429 errors?", "api-rate-limits.md"),
    EvalCase("What is the rate limit for API keys?", "api-rate-limits.md"),
]

A 3072-dim model like gemini-embedding-001 should give you 100% recall@3 on a corpus this small. If you’re below 80% on your real corpus, the issue is almost always chunking, not the embedding model.

Production Considerations

Chunking strategy:

  • 400-char chunks with 80-char overlap works for prose. For code: chunk by function. For tables: keep rows together.
  • Semantic chunking (splitting at sentence boundaries) consistently outperforms fixed-size at the cost of more preprocessing. Worth it once you’re past prototype.

Embedding model choice:

  • gemini-embedding-001 (3072 dims): best quality, highest cost per call, ideal for precision-critical retrieval
  • text-embedding-004 (768 dims): 4× more storage-efficient, ~10% lower recall on technical corpora, worth it at scale
  • Re-embed your corpus when you switch models — embeddings from different models are not comparable

Retriever caching:

  • Queries to the same system cluster heavily. A 1-hour TTL cache on query embeddings cuts embedding API calls by 40–60% in practice.
  • Don’t cache over index updates — invalidate on write.

When define_retriever lands: The pattern above maps directly to what the retriever abstraction will look like. When it’s available, you’ll register your retriever once and call ai.retrieve() instead of calling your function directly. The Genkit DevUI will also show retrieval spans in the trace — useful for debugging low-recall queries.

The Full Picture

User Query
    │
    ▼
ai.embed(query)          ← single API call, ~50ms
    │
    ▼
cosine_search(top_k=3)   ← in-memory: <1ms; pgvector at 10M chunks: ~5ms
    │
    ▼
build_rag_prompt()       ← 400-1200 tokens injected into context
    │
    ▼
ai.generate()            ← Gemini call, ~800ms median
    │
    ▼
RAGOutput(answer, sources)

End-to-end latency on this pattern: ~900ms for a cold query. The embedding step is cheap; the generation step dominates. If you need <500ms, cache embeddings aggressively and use gemini-2.0-flash over gemini-2.5-pro.

What’s Next

  • Swap InMemoryVectorStore for Cloud Firestore vector search or AlloyDB — both work with the same interface shown here
  • Add metadata filtering to retrieve() to scope retrieval by document category, date, or access level
  • Watch for define_retriever / define_indexer in the Python SDK — they’ll wrap exactly this pattern with Genkit DevUI observability built in

The full working file is ~150 lines. [Gist link placeholder] — run it with your own docs and see where recall lands.