Disclaimer: The opinions expressed in this article are my own and do not represent the views of Google. This content is based solely on publicly available information.
Note: the runnable walkthrough below uses SQLite + TF-IDF (Term Frequency–Inverse Document Frequency) for portability, so no FastAPI server or pgvector setup is required. We then show how to swap in FastAPI and pgvector at the end (production stage).
Most RAG tutorials show you how to call an embedding API and do a similarity search. That is the easy 20%. The hard 80% is the architecture around it: how you chunk documents, how you prevent retrieved chunks from being redundant, how you expose this as a reliable HTTP API, and how you reason about the latency budget at each stage.
This article builds the complete stack from scratch in pure Python. No FastAPI or pgvector is required to run the code, and the architecture maps directly to what you deploy with those tools. Every component is a drop-in replacement: swap the TF-IDF vectors for sentence-transformers, the SQLite store for pgvector, and the mock LLM for Claude.
Architecture Overview
A production RAG API has five components in a sequential pipeline:
flowchart LR
Q[User Query] --> E[Embedding]
E --> V[Vector Search]
V --> M[MMR Rerank]
M --> L[LLM Call]
L --> R[Response]
Each component is independently configurable. The critical design decisions happen at the seams: how big to make chunks, how many to retrieve, how aggressively to rerank for diversity.
Why Chunking Strategy Matters
The single biggest influence on retrieval quality is chunk size. Too small and individual chunks lack enough context: retrieved passages are fragments that the LLM cannot interpret in isolation. Too large and chunks contain multiple topics, so a chunk matching on paragraph 1 also drags in paragraph 4 which is about something different, diluting precision.
The standard heuristic is 150–200 tokens with 20–30 tokens of overlap. Overlap prevents information loss at chunk boundaries: a sentence spanning two adjacent chunks is fully visible in at least one chunk instead of silently split across both.
Complete Implementation
The implementation below builds all five pipeline components in one self-contained file with no external dependencies; swap each piece for its production equivalent when you deploy.
import re
import math
import json
import sqlite3
import time
from dataclasses import dataclass, field
from typing import Optional
# ---------------------------------------------------------------------------
# Data model
# ---------------------------------------------------------------------------
@dataclass
class Chunk:
chunk_id: str
doc_id: str
text: str
metadata: dict = field(default_factory=dict)
@dataclass
class RAGResult:
query: str
answer: str
chunks_used: list
retrieval_ms: float
total_ms: float
# ---------------------------------------------------------------------------
# Utilities
# ---------------------------------------------------------------------------
def _tokenize(text: str) -> list[str]:
return re.findall(r'\b\w+\b', text.lower())
def _tfidf_vector(text: str, idf: dict[str, float]) -> dict[str, float]:
tokens = _tokenize(text)
if not tokens:
return {}
tf: dict[str, float] = {}
for t in tokens:
tf[t] = tf.get(t, 0) + 1
n = len(tokens)
# Default IDF of 1.0 for unknown terms (raw TF when idf is empty)
return {t: (c / n) * idf.get(t, 1.0) for t, c in tf.items()}
def _cosine(a: dict[str, float], b: dict[str, float]) -> float:
if not a or not b:
return 0.0
dot = sum(a.get(k, 0) * v for k, v in b.items())
na = math.sqrt(sum(v * v for v in a.values()))
nb = math.sqrt(sum(v * v for v in b.values()))
if na == 0 or nb == 0:
return 0.0
return dot / (na * nb)
# ---------------------------------------------------------------------------
# Chunker
# ---------------------------------------------------------------------------
class Chunker:
"""Splits documents into overlapping fixed-size windows of whitespace-split words.
Note: chunk_size/overlap are counted in whitespace-split words (text.split()),
not tokenizer tokens. Substitute a tokenizer-aware count in production.
"""
def __init__(self, chunk_size: int = 150, overlap: int = 30) -> None:
self.chunk_size = chunk_size
self.overlap = overlap
def chunk(self, doc_id: str, text: str,
metadata: Optional[dict] = None) -> list[Chunk]:
metadata = metadata or {}
words = text.split()
step = max(1, self.chunk_size - self.overlap)
chunks: list[Chunk] = []
i = 0
seq = 0
while i < len(words):
end = min(i + self.chunk_size, len(words))
chunk_text = " ".join(words[i:end])
chunks.append(Chunk(
chunk_id=f"{doc_id}::{seq}",
doc_id=doc_id,
text=chunk_text,
metadata=metadata,
))
if end >= len(words):
break
i += step
seq += 1
return chunks
# ---------------------------------------------------------------------------
# VectorStore
# ---------------------------------------------------------------------------
class VectorStore:
"""SQLite-backed chunk store with in-memory TF-IDF vectors."""
def __init__(self, db_path: str = ":memory:") -> None:
self.conn = sqlite3.connect(db_path)
self._idf: dict[str, float] = {}
self._vectors: dict[str, dict[str, float]] = {}
self._texts: dict[str, str] = {}
self._chunks: dict[str, Chunk] = {}
self._init_schema()
def _init_schema(self) -> None:
self.conn.execute("""
CREATE TABLE IF NOT EXISTS chunks (
chunk_id TEXT PRIMARY KEY,
doc_id TEXT,
text TEXT,
metadata TEXT
)
""")
self.conn.commit()
def upsert(self, chunks: list[Chunk]) -> None:
rows = [(c.chunk_id, c.doc_id, c.text, json.dumps(c.metadata))
for c in chunks]
self.conn.executemany(
"INSERT OR REPLACE INTO chunks VALUES (?, ?, ?, ?)", rows
)
self.conn.commit()
for c in chunks:
self._chunks[c.chunk_id] = c
self._texts[c.chunk_id] = c.text
self._rebuild_idf()
def _rebuild_idf(self) -> None:
df: dict[str, int] = {}
N = len(self._texts)
for text in self._texts.values():
for t in set(_tokenize(text)):
df[t] = df.get(t, 0) + 1
self._idf = {t: math.log((N + 1) / (n + 1)) + 1
for t, n in df.items()}
self._vectors = {
cid: _tfidf_vector(text, self._idf)
for cid, text in self._texts.items()
}
def search(self, query: str,
top_k: int = 10) -> list[tuple[Chunk, float]]:
q_vec = _tfidf_vector(query, self._idf)
scored = [
(self._chunks[cid], _cosine(q_vec, vec))
for cid, vec in self._vectors.items()
]
return sorted(scored, key=lambda x: -x[1])[:top_k]
def stats(self) -> dict:
cur = self.conn.execute(
"SELECT COUNT(*), COUNT(DISTINCT doc_id) FROM chunks")
n_chunks, n_docs = cur.fetchone()
return {"chunks": n_chunks, "docs": n_docs}
# ---------------------------------------------------------------------------
# MMR Reranker
# ---------------------------------------------------------------------------
def mmr_rerank(
query_vec: dict[str, float],
candidates: list[tuple[Chunk, float]],
idf: dict[str, float],
top_k: int = 5,
lambda_mult: float = 0.5,
) -> list[tuple[Chunk, float]]:
"""Maximal Marginal Relevance reranking for diversity.
MMR_score(d) = lambda * sim(query, d) -
(1 - lambda) * max_{s in selected} sim(s, d)
"""
selected: list[tuple[Chunk, float]] = []
remaining: list[tuple[Chunk, float]] = list(candidates)
# Pre-compute corpus-IDF-weighted vectors for the candidate chunks
cand_vecs = {c.chunk_id: _tfidf_vector(c.text, idf) for c, _ in candidates}
while remaining and len(selected) < top_k:
scores = []
for chunk, rel_score in remaining:
c_vec = cand_vecs[chunk.chunk_id]
if selected:
redundancy = max(
_cosine(cand_vecs[s.chunk_id], c_vec)
for s, _ in selected
)
else:
redundancy = 0.0
mmr = lambda_mult * rel_score - (1 - lambda_mult) * redundancy
scores.append((chunk, mmr))
best = max(scores, key=lambda x: x[1])
selected.append(best)
remaining = [(c, s) for c, s in remaining
if c.chunk_id != best[0].chunk_id]
return selected
# ---------------------------------------------------------------------------
# Mock LLM (swap for real Claude API call)
# ---------------------------------------------------------------------------
def mock_llm(prompt: str) -> str:
"""Placeholder - replace with Anthropic client call in production."""
lines = prompt.split("\n")
context_lines = [l for l in lines if l.strip().startswith("[")]
sources = sorted({l.split("doc=")[1].split(")")[0]
for l in context_lines if "doc=" in l})
return (f"Based on the provided context, the answer draws on information "
f"from sources: {', '.join(sources)}.")
# ---------------------------------------------------------------------------
# RAG Pipeline
# ---------------------------------------------------------------------------
class RAGPipeline:
def __init__(
self,
chunk_size: int = 150,
overlap: int = 30,
top_k: int = 5,
mmr_lambda: float = 0.5,
db_path: str = ":memory:",
) -> None:
self.chunker = Chunker(chunk_size, overlap)
self.store = VectorStore(db_path)
self.top_k = top_k
self.mmr_lambda = mmr_lambda
def ingest(self, doc_id: str, text: str,
metadata: Optional[dict] = None) -> int:
chunks = self.chunker.chunk(doc_id, text, metadata)
self.store.upsert(chunks)
return len(chunks)
def query(self, question: str) -> RAGResult:
t0 = time.perf_counter()
# 1. Retrieve more candidates than needed for MMR
candidates = self.store.search(question, top_k=self.top_k * 3)
retrieval_ms = (time.perf_counter() - t0) * 1000
# 2. MMR rerank
q_vec = _tfidf_vector(question, self.store._idf)
reranked = mmr_rerank(q_vec, candidates, self.store._idf,
top_k=self.top_k, lambda_mult=self.mmr_lambda)
# 3. Build augmented prompt
context = "\n".join(
f" [{i+1}] (doc={c.doc_id}) {c.text}"
for i, (c, _) in enumerate(reranked)
)
prompt = (
"Answer using only the context below.\n\n"
f"Context:\n{context}\n\n"
f"Question: {question}\nAnswer:"
)
# 4. Generate (replace mock_llm with real Claude call)
answer = mock_llm(prompt)
total_ms = (time.perf_counter() - t0) * 1000
return RAGResult(
query=question,
answer=answer,
chunks_used=[c for c, _ in reranked],
retrieval_ms=retrieval_ms,
total_ms=total_ms,
)
def stats(self) -> dict:
return self.store.stats()
# ---------------------------------------------------------------------------
# Demo driver
# ---------------------------------------------------------------------------
DOCS = {
"attention_paper":
"Attention mechanisms have revolutionized natural language processing. "
"The Transformer architecture introduced by Vaswani et al. (2017) uses "
"self-attention to compute representations of sequences. Unlike recurrent "
"neural networks, transformers process all tokens in parallel, leading to "
"significant speedups. The key operation is scaled dot-product attention: "
"given queries Q, keys K, and values V, attention is computed as "
"softmax(QK^T / sqrt(d_k)) * V. Multi-head attention runs this operation h "
"times with different learned projections, allowing the model to attend "
"to different aspects of the input simultaneously. Positional encodings "
"are added to token embeddings to preserve sequence order. BERT uses "
"bidirectional attention for encoding, while GPT uses causal (masked) "
"attention for autoregressive generation. The KV cache stores previously "
"computed key and value tensors during inference to avoid redundant "
"computation. FlashAttention improves memory efficiency by tiling the "
"computation to fit in SRAM.",
"rag_overview":
"Retrieval-Augmented Generation (RAG) enhances language model responses "
"by retrieving relevant documents before generation. A RAG system has "
"three main components: a document store (vector database), a retriever, "
"and a generator. During ingestion, documents are split into chunks and "
"converted to embeddings using a model like sentence-transformers or "
"OpenAI's text-embedding-ada-002. Embeddings are stored in a vector "
"database such as pgvector, Pinecone, or Weaviate. At query time, the "
"user's question is embedded and the top-k most similar chunks are "
"retrieved using approximate nearest neighbor search (HNSW or IVF). The "
"retrieved chunks are inserted into the LLM prompt as context. RAG "
"reduces hallucination by grounding the model in retrieved facts. Hybrid "
"search combines dense vector search with sparse BM25 for better recall. "
"Maximal Marginal Relevance (MMR) reranking reduces redundancy in the "
"retrieved set. Prefix caching can cache the context portion of the "
"prompt to reduce latency.",
"vector_databases":
"Vector databases store high-dimensional embeddings and support efficient "
"similarity search. HNSW (Hierarchical Navigable Small World) is a "
"graph-based index that achieves O(log n) query time with high recall. "
"IVF (Inverted File Index) clusters vectors into Voronoi cells and "
"searches only the nearest clusters. DiskANN enables billion-scale search "
"from disk without requiring all vectors in RAM. pgvector is a PostgreSQL "
"extension that adds vector similarity operators and supports both exact "
"and approximate search. The ivfflat index type partitions vectors into "
"lists; the hnsw index type builds a multi-layer graph. Filtered search "
"(WHERE clause + vector search) requires special index support to avoid "
"full scans. Multi-tenancy can be implemented via namespace columns with "
"partial indexes or separate tables per tenant. Memory usage scales as "
"O(n * d * bytes_per_float) where n is number of vectors and d is "
"dimensionality.",
}
if __name__ == "__main__":
pipeline = RAGPipeline(chunk_size=100, overlap=20, top_k=3, mmr_lambda=0.5)
print("--- Ingesting documents ---")
for doc_id, text in DOCS.items():
n = pipeline.ingest(doc_id, text)
print(f" {doc_id}: {n} chunks")
print(f" Total chunks stored: {pipeline.stats()['chunks']}")
queries = [
"How does the KV cache work in transformers?",
"What vector database index should I use for billion-scale search?",
"How does MMR reranking improve retrieval quality?",
]
print("\n--- Querying ---")
for q in queries:
result = pipeline.query(q)
ordered = [c.doc_id for c in result.chunks_used]
print(f" Q: {q}")
print(f" Sources (ranked): {ordered}")
print(f" Retrieval: {result.retrieval_ms:.1f}ms "
f"Total: {result.total_ms:.1f}ms")
print()
print("--- Chunk size sensitivity ---")
for cs in [50, 100, 200]:
p = RAGPipeline(chunk_size=cs, overlap=cs // 5)
for doc_id, text in DOCS.items():
p.ingest(doc_id, text)
n = p.stats()["chunks"]
print(f" chunk_size={cs:4d}: {n} chunks total")
Running this produces:
--- Ingesting documents ---
attention_paper: 2 chunks
rag_overview: 2 chunks
vector_databases: 2 chunks
Total chunks stored: 6
--- Querying ---
Q: How does the KV cache work in transformers?
Sources (ranked): ['attention_paper', 'vector_databases', 'rag_overview']
Retrieval: 0.0ms Total: 0.3ms
Q: What vector database index should I use for billion-scale search?
Sources (ranked): ['vector_databases', 'attention_paper', 'rag_overview']
Retrieval: 0.0ms Total: 0.3ms
Q: How does MMR reranking improve retrieval quality?
Sources (ranked): ['rag_overview', 'attention_paper', 'vector_databases']
Retrieval: 0.0ms Total: 0.3ms
--- Chunk size sensitivity ---
chunk_size= 50: 12 chunks total
chunk_size= 100: 6 chunks total
chunk_size= 200: 3 chunks total
All three queries retrieve the expected primary source first, and MMR diversifies the remaining slots across the other documents. Retrieval is sub-millisecond at 6 chunks; the latency budget is dominated by LLM generation in production.
The first sensitivity to inspect is chunk size itself, since it controls both the number of vectors you store and how specific each retrieved fragment is.
Figure 1: Chunk count, average chunk length, and top-1 retrieval similarity across a sweep of chunk sizes. (Figure generated by a companion plotting script not shown inline; values are illustrative.)
Chunk count falls roughly as 1/chunk_size (the inline demo’s sweep produces 12 → 6 → 3 chunks across 50/100/200 tokens), so the storage cost of small chunks grows roughly as 1/chunk_size: each doubling of chunk size halves the chunk count (12 → 6 → 3 across 50/100/200 tokens, with overlap scaled proportionally as chunk_size/5). The average chunk length saturates once chunk_size exceeds the document length itself: beyond that, raising chunk_size cannot grow chunks further, only thin out the index. Top-1 retrieval similarity is non-monotone, peaking in the small-to-medium chunk range before falling at the longer chunk sizes. Smaller chunks score higher on cosine because the query terms make up a larger fraction of the chunk’s token count, but precision wins on this metric come at the cost of producing single-fragment passages that lack supporting context. That tradeoff is why the 150–200 token band is the production sweet spot: it sits past the steepest similarity decay but still preserves enough context for the LLM to interpret each passage in isolation.
The next decision is how to balance relevance against diversity when assembling the top-k window, which is what MMR’s lambda parameter controls.
Figure 2: MMR lambda swept from 0 to 1, with diversity (left) and per-document chunk counts (right) at top-k=5. (Figure generated by a companion plotting script not shown inline; values are illustrative.)
The left panel sweeps lambda from 0 (max diversity) to 1 (max relevance) and tracks average pairwise similarity (red) and average relevance to the query (blue). Both lines are nearly flat across the range because the corpus only contains six candidate chunks across three documents, so the candidate pool is already exhausted before MMR has room to substitute different chunks: every lambda ends up selecting the same five candidates, just in slightly different orders. The right panel makes the same point visually: the stacked bars are nearly identical from lambda=0 through lambda=1 because the doc-id distribution barely changes when the candidate pool is small relative to top_k. In a production corpus with thousands of chunks per query, the same plot separates cleanly: lambda=0.0 spreads chunks across many source documents and lambda=1.0 concentrates on the single highest-similarity document. The behavior to remember is the principle, not the magnitude on this toy corpus: MMR’s effect grows with candidate count, which is why retrieving 3× the desired top_k matters.
The third diagnostic is latency, where the practical question is which stage to optimize first.
Figure 3: Per-stage latency as corpus size grows, plus a stage breakdown at the largest corpus. (Figure generated by a companion plotting script not shown inline; values are illustrative.)
At small corpus sizes (<250 docs) MMR reranking adds a roughly constant overhead that dominates retrieval; the two cross over near 250 documents (illustrative threshold; the exact crossover depends on your embedding dimensionality and top_k), after which vector search latency grows linearly with corpus size while reranking stays roughly flat. The crossover is where search overtakes rerank as the dominant cost, and where a pgvector HNSW (Hierarchical Navigable Small World) index becomes worthwhile because it flattens the search line to something near-constant. The breakdown at the largest corpus shows the entire retrieval pipeline still costs only a few milliseconds before any LLM call. Compared to the 500–3000ms a real Claude generation costs, every millisecond of retrieval tuning is a rounding error. Optimize generation first; the retrieval layer is well under any reasonable latency budget.
The VectorStore in Detail
The TF-IDF store is a direct simulation of what pgvector does. The IDF weights (inverse document frequency) are recomputed on every upsert; in production you maintain these incrementally. The cosine similarity function normalises by L2 norm, which is equivalent to the <=> (cosine distance) operator in pgvector.
The key production mapping:
| This implementation | Production equivalent |
|---|---|
| TF-IDF sparse vectors | sentence-transformers or text-embedding-ada-002 |
| In-memory vector cache | pgvector HNSW (Hierarchical Navigable Small World) or ivfflat (pgvector’s inverted-file flat index) |
| SQLite row scan | ANN (Approximate Nearest Neighbor) query via <=> operator |
_rebuild_idf() on insert | Batch embedding + COPY to Postgres |
The interface contract (upsert(chunks) and search(query, top_k)) stays identical.
MMR: Why It Matters
Pure cosine similarity retrieves the most relevant chunks, but the top-5 results often contain nearly identical passages. If your document corpus has three paragraphs about the KV cache (introduction, technical detail, and code example), a KV cache query without MMR returns all three even though they overlap heavily in content.
MMR’s per-step selection criterion:
MMR_score(d) = λ × similarity(query, d) − (1−λ) × max_{s ∈ selected} similarity(s, d)
The second term penalises any candidate whose content resembles an already-selected chunk. At λ=1.0, MMR reduces to pure cosine ranking. At λ=0.0, it selects for maximum diversity regardless of relevance. The production default of λ=0.5 balances both goals.
The implementation above uses a greedy forward pass: at each step it selects the best remaining candidate under the MMR criterion, adds it to the selected set, and removes it from candidates. This greedy algorithm is suboptimal compared to exact MMR (which would require exponential search) but is fast and achieves near-optimal diversity in practice.
Production Deployment: FastAPI + pgvector
Replacing this demo with production components requires three changes; the chunking, MMR, and prompt-building logic stay unchanged.
Switch to real embeddings:
from sentence_transformers import SentenceTransformer
model = SentenceTransformer("BAAI/bge-base-en-v1.5")
def embed(text: str) -> list[float]:
return model.encode(text, normalize_embeddings=True).tolist()
Store in pgvector:
CREATE EXTENSION IF NOT EXISTS vector;
CREATE TABLE chunks (
chunk_id TEXT PRIMARY KEY,
doc_id TEXT NOT NULL,
text TEXT NOT NULL,
metadata JSONB,
embedding vector(768)
);
CREATE INDEX ON chunks USING hnsw (embedding vector_cosine_ops)
WITH (m = 16, ef_construction = 64);
Query with FastAPI + pgvector:
from fastapi import FastAPI, Depends
from sqlalchemy import text
from sqlalchemy.ext.asyncio import (
AsyncSession,
create_async_engine,
async_sessionmaker,
)
DATABASE_URL = "postgresql+asyncpg://user:password@localhost/db" # replace with your connection string
engine = create_async_engine(DATABASE_URL)
AsyncSessionLocal = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
app = FastAPI()
# Module-scope config + helpers used by the routes below.
TOP_K = 5
chunker = Chunker(chunk_size=150, overlap=30)
def build_prompt(query: str, chunks: list[tuple[Chunk, float]]) -> str:
context = "\n".join(c.text for c, _ in chunks)
return f"Context:\n{context}\n\nQuestion: {query}"
# Standard FastAPI dependency that yields an AsyncSession per request.
async def get_db():
async with AsyncSessionLocal() as session:
yield session
# mmr_rerank_dense is the dense (embedding-based) version of mmr_rerank shown
# earlier; same shape, takes embeddings instead of TF-IDF vectors.
def mmr_rerank_dense(
query_emb: list[float],
candidates: list[tuple[Chunk, float]],
top_k: int = 5,
lambda_mult: float = 0.5,
) -> list[tuple[Chunk, float]]:
def cos(a: list[float], b: list[float]) -> float:
dot = sum(x * y for x, y in zip(a, b))
na = math.sqrt(sum(x * x for x in a))
nb = math.sqrt(sum(x * x for x in b))
return 0.0 if na == 0 or nb == 0 else dot / (na * nb)
# Each candidate carries (Chunk, relevance_score); we also need the chunk
# embedding. In a real deployment, fetch chunk embeddings alongside the row
# in the SQL query; here we assume c.metadata["embedding"] holds it.
cand_embs = {c.chunk_id: c.metadata["embedding"] for c, _ in candidates}
selected: list[tuple[Chunk, float]] = []
remaining: list[tuple[Chunk, float]] = list(candidates)
while remaining and len(selected) < top_k:
scored = []
for chunk, rel in remaining:
c_emb = cand_embs[chunk.chunk_id]
redundancy = (
max(cos(cand_embs[s.chunk_id], c_emb) for s, _ in selected)
if selected else 0.0
)
mmr = lambda_mult * rel - (1 - lambda_mult) * redundancy
scored.append((chunk, mmr))
best = max(scored, key=lambda x: x[1])
selected.append(best)
remaining = [(c, s) for c, s in remaining
if c.chunk_id != best[0].chunk_id]
return selected
async def call_claude(prompt: str, model: str = "claude-opus-4-7") -> str: # substitute with a snapshot-suffixed model from the Anthropic model catalog (e.g., claude-opus-4-7-YYYYMMDD)
from anthropic import AsyncAnthropic
client = AsyncAnthropic()
msg = await client.messages.create(model=model, max_tokens=1024,
messages=[{"role": "user", "content": prompt}])
return msg.content[0].text
@app.post("/ingest")
async def ingest(doc_id: str, body: str,
db: AsyncSession = Depends(get_db)):
chunks = chunker.chunk(doc_id, body)
for c in chunks:
emb = embed(c.text)
# Using raw text() queries for clarity; in production prefer SQLAlchemy Core / ORM for type safety and automatic param binding
await db.execute(
text(
"INSERT INTO chunks VALUES (:chunk_id, :doc_id, :text, :meta, :emb) "
"ON CONFLICT (chunk_id) DO UPDATE SET embedding = EXCLUDED.embedding"
),
{"chunk_id": c.chunk_id, "doc_id": c.doc_id,
"text": c.text, "meta": json.dumps(c.metadata),
"emb": emb},
)
await db.commit()
return {"doc_id": doc_id, "chunks_created": len(chunks)}
@app.get("/query")
async def query(q: str, db: AsyncSession = Depends(get_db)):
q_emb = embed(q)
rows = await db.execute(
text(
"SELECT chunk_id, doc_id, text, embedding, "
"1 - (embedding <=> :q) AS similarity "
"FROM chunks ORDER BY embedding <=> :q LIMIT :k"
),
{"q": q_emb, "k": TOP_K * 3},
)
candidates = []
for r in rows.mappings().all():
# pgvector returns the embedding as a string '[0.1,0.2,...]' under raw asyncpg/SQLAlchemy; parse to list[float]. The pgvector.asyncpg.register_vector helper handles this transparently in production.
import json
raw_embed = r["embedding"]
embed_list = raw_embed if isinstance(raw_embed, list) else json.loads(raw_embed)
candidates.append(
(
Chunk(
r["chunk_id"], r["doc_id"], r["text"],
metadata={"embedding": embed_list},
),
r["similarity"],
)
)
# We already use `mmr_rerank_dense` here; in the demo above we used TF-IDF
# for the same MMR formula since SQLite has no native vector type.
reranked = mmr_rerank_dense(q_emb, candidates, top_k=TOP_K)
answer = await call_claude(build_prompt(q, reranked))
return {"answer": answer, "chunks": [c.chunk_id for c, _ in reranked]}
The Chunker, mmr_rerank, and prompt construction logic remain unchanged; only the embedding function and similarity backend change. The async call_claude defined above is the version wired into the FastAPI route; a sync variant using Anthropic() is straightforward to write if you need it outside a request context.
Latency Budget
In a production RAG system the latency splits roughly as follows:
| Stage | Typical latency | Notes |
|---|---|---|
| Embedding the query | 5–20ms | GPU: ~5ms; CPU: ~20ms |
| ANN search (pgvector HNSW) | 1–5ms | O(log n) with HNSW index |
| MMR rerank | 1–3ms | O(k²) in candidate count |
| LLM generation (Claude) | 500–3000ms | Dominant cost |
| Total | ~520–3030ms | Retrieval is 1–5% of total |
The implication: optimising retrieval latency only matters after the generation step has been profiled and tuned. A 10% improvement in retrieval latency saves ~0.5ms on a 500ms response, which is invisible to users.
Practical Recommendations
Five concrete decisions determine retrieval quality in production, ordered by impact:
-
Start with chunk_size=150, overlap=30, and measure recall@5 on a held-out query set. Adjust chunk_size based on document density: technical documentation with dense information prefers smaller chunks (100–150 tokens), whereas narrative text works better with larger chunks (200–300 tokens).
-
Use λ=0.5 for MMR as the default. Lower λ (0.3–0.4) for queries that span multiple topics. Higher λ (0.7–0.8) for narrow factual questions where you want the most relevant passages regardless of redundancy.
-
Retrieve 3× the desired top_k before MMR. MMR needs a broad candidate set to find diverse chunks. Retrieving exactly top_k and then reranking defeats the purpose because you need candidates that were initially ranked lower to provide diversity.
-
Monitor retrieval latency per percentile (P50, P95, P99), not just mean. A query against an HNSW index has rare worst-case slow paths when the entry node is far from the query. Log individual retrieval times.
-
Add metadata filtering before vector search when possible. If your chunks carry document type, date, or source metadata, filter at the SQL layer before the ANN pass (
WHERE doc_type = 'policy' AND (embedding <=> $1) < 0.5) instead of post-filtering retrieved results.
Summary
Five components, one pipeline:
| Component | Key parameter | Production tool |
|---|---|---|
| Chunker | chunk_size, overlap | 150–200 tokens, 20–30 token overlap |
| VectorStore | Index type | pgvector HNSW (Hierarchical Navigable Small World) for recall; ivfflat (pgvector’s inverted-file flat index) for memory |
| Retriever | top_k, embedding model | BGE (BAAI General Embedding) / E5 (Microsoft Research’s E5 embedding family) for English; domain-specific where recall matters |
| MMR Reranker | lambda_mult | 0.5 default; lower for multi-document queries |
| Generator | Model, max_tokens | Claude for reasoning quality; faster models for latency |
The latency budget: retrieval takes less than 10ms with ANN indexing; generation takes 500–3000ms. Optimise retrieval only after profiling generation, since generation is almost always the bottleneck.