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 Genkit Python 0.8.1 on PyPI, there are no built-in define_retriever / define_indexer abstractions yet. What does exist is a first-class embedder API, plus a docs= argument on generate if you already have Document objects to pass through. 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:
- Indexes documents by chunking and embedding them
- Accepts user queries, embeds the query, and finds the top-K matching chunks
- Builds an augmented prompt and generates a grounded answer
- Returns the answer with source citations
- 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]==0.8.1"
uv pip install numpy # only needed for cosine similarity helperSet 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_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 dimsgoogleai/gemini-embedding-001 produces 3072-dimensional vectors. The plugin also lists googleai/text-embedding-004 (768 dims), but that model ID currently 404s on the Google AI (generativelanguage) embed endpoint I tested against with genkit-google-genai==0.8.1. Stick with gemini-embedding-001 (or gemini-embedding-2) unless you have confirmed Vertex AI access for the smaller model.
Step 1: Chunking Strategy
Chunking is where most RAG implementations quietly fail. Too large a chunk leads to context dilution. Too small leads to 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 is roughly 80 to 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 chunksThe 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
# Alternative: pass Document objects via docs= instead of stuffing text into the prompt.
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)A note on docs=: ai.generate(..., docs=[...]) accepts a list of Genkit Document objects. If your retriever already returns Document instances, that path is cleaner than concatenating chunk text into the prompt string. The rest of this post keeps the string form so the full retrieval loop stays visible.
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 to 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 [email protected].
""",
}
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 googleai/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:
googleai/gemini-embedding-001(3072 dims): works on Google AI today; good default for precision-critical retrievalgoogleai/gemini-embedding-2(also 3072 dims on the endpoint I tested): newer listing in the 0.8.1 plugin enum- The plugin still exposes
googleai/text-embedding-004, but that ID 404’d on Google AI embed for me on 0.8.1. Do not plan a 768-dim rollout on it until you re-verify - 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 to 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, which helps when 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 under 500ms, cache embeddings aggressively and use gemini-2.0-flash over gemini-2.5-pro.
What’s Next
- Swap
InMemoryVectorStorefor 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_indexerin the Python SDK. They’ll wrap exactly this pattern with Genkit DevUI observability built in.
The full working file is about 150 lines. Run it with your own docs and see where recall lands.