Production RAG — hybrid retrieval + cited answers

Production RAG — hybrid retrieval + cited answers#

Track 08 - Production · Notebook 05 · Runtime: ~30s LIVE, <1s replay

A 12-passage corpus, BM25 + dense + RRF fusion, then an Anthropic call that returns a cited answer. Citation accuracy is measured against ground truth.

from llm_systems_cookbook.nb import bootstrap
from llm_systems_cookbook._utils import repo_root

import json
import math
import os
import re
from collections import Counter
from pathlib import Path

s = bootstrap("08_production_05_hybrid_rag_production")

LIVE = bool(os.environ.get("ANTHROPIC_API_KEY"))
MODEL = os.environ.get("MODEL_ANTHROPIC", "claude-sonnet-4-6")
FIXTURE = json.loads((Path(repo_root()) / "notebooks/08_production/_fixtures/05_rag.json").read_text())
PASSAGES = FIXTURE["passages"]
QUERIES  = FIXTURE["queries"]
TRUTH    = FIXTURE["expected_passage_ids"]  # 1-indexed
print(f"mode={'LIVE' if LIVE else 'REPLAY'}  passages={len(PASSAGES)}  queries={len(QUERIES)}")

Corpus and queries#

12 paragraphs about a fictional code-search product, 5 factual queries with one ground-truth passage each.

for i, p in enumerate(PASSAGES[:3], start=1):
    print(f"  [{i:>2}] {p[:90]}{'...' if len(p) > 90 else ''}")
print("  ...")
for q, t in zip(QUERIES, TRUTH):
    print(f"  Q (truth=p{t}): {q}")

BM25 — sparse score (live, pure Python)#

_TOK = re.compile(r"[a-z0-9]+")
def tokenize(s: str) -> list[str]:
    return _TOK.findall(s.lower())

class BM25:
    def __init__(self, corpus: list[str], k1: float = 1.5, b: float = 0.75) -> None:
        self.k1, self.b = k1, b
        self.docs = [tokenize(d) for d in corpus]
        self.dl = [len(d) for d in self.docs]
        self.avgdl = sum(self.dl) / len(self.dl)
        self.df: dict[str, int] = Counter()
        for d in self.docs:
            for w in set(d):
                self.df[w] += 1
        self.N = len(self.docs)

    def score(self, query: str) -> list[float]:
        qtoks = tokenize(query)
        out = [0.0] * self.N
        for w in qtoks:
            if w not in self.df:
                continue
            idf = math.log(1 + (self.N - self.df[w] + 0.5) / (self.df[w] + 0.5))
            for i, d in enumerate(self.docs):
                tf = d.count(w)
                if tf == 0: continue
                out[i] += idf * (tf * (self.k1 + 1)) / (tf + self.k1 * (1 - self.b + self.b * self.dl[i] / self.avgdl))
        return out


bm25 = BM25(PASSAGES)
bm25_scores = [bm25.score(q) for q in QUERIES]
print("BM25 ranks per query (1-indexed top-3):")
for q, scs in zip(QUERIES, bm25_scores):
    top = sorted(range(len(scs)), key=lambda i: -scs[i])[:3]
    print(f"  {q[:50]:<52} → {[i + 1 for i in top]}")

Dense scores (LIVE: sentence-transformers · REPLAY: fixture)#

Cosine similarity between query and passage embeddings. The fixture carries the matrix from a real run of BAAI/bge-small-en-v1.5 so the notebook is self-contained on Colab.

def compute_dense_scores() -> list[list[float]]:
    if not LIVE:
        return FIXTURE["dense_scores"]
    from sentence_transformers import SentenceTransformer  # noqa: PLC0415
    enc = SentenceTransformer("BAAI/bge-small-en-v1.5")
    p_emb = enc.encode(PASSAGES, normalize_embeddings=True)
    q_emb = enc.encode(QUERIES, normalize_embeddings=True)
    return (q_emb @ p_emb.T).tolist()


dense_scores = compute_dense_scores()
print("dense top-3 per query:")
for q, scs in zip(QUERIES, dense_scores):
    top = sorted(range(len(scs)), key=lambda i: -scs[i])[:3]
    print(f"  {q[:50]:<52} → {[i + 1 for i in top]}")

Reciprocal Rank Fusion (RRF)#

RRF_score(d) = Σ 1 / (k + rank_i(d)) over each ranker. k=60 is the published default. RRF needs no calibration between sparse and dense scores, which is why it has eaten the hybrid-retrieval world.

def rrf_fuse(score_lists: list[list[float]], k: int = 60) -> list[float]:
    n_docs = len(score_lists[0])
    fused = [0.0] * n_docs
    for scs in score_lists:
        ranks = sorted(range(n_docs), key=lambda i: -scs[i])
        for r, idx in enumerate(ranks):
            fused[idx] += 1.0 / (k + r + 1)
    return fused


fused = [rrf_fuse([bm25_scores[i], dense_scores[i]]) for i in range(len(QUERIES))]

top_ks = []
for q, scs in zip(QUERIES, fused):
    top = sorted(range(len(scs)), key=lambda i: -scs[i])[:3]
    top_ks.append([i + 1 for i in top])
    print(f"  {q[:50]:<52} → {top_ks[-1]}")

recall_at_3 = sum(t in top_ks[i] for i, t in enumerate(TRUTH)) / len(TRUTH)
mrr = sum(1 / (top_ks[i].index(t) + 1) if t in top_ks[i] else 0 for i, t in enumerate(TRUTH)) / len(TRUTH)
print(f"\nRRF recall@3 = {recall_at_3:.0%}   MRR = {mrr:.2f}")

Cited answer (LIVE: Anthropic · REPLAY: fixture)#

Top-3 fused passages go into the prompt. The system message asks for inline [id] citations. We then verify that every cited id appears in the retrieved set and was the ground-truth passage.

SYSTEM = (
    "Answer the question using ONLY the provided passages. "
    "Cite passages inline like [3] for passage 3. If no passage answers, "
    "say 'no answer in context'."
)


def answer(query: str, top: list[int]) -> dict:
    if not LIVE:
        return FIXTURE["answers"][query]

    import anthropic  # noqa: PLC0415
    context = "\n".join(f"[{i}] {PASSAGES[i - 1]}" for i in top)
    resp = anthropic.Anthropic().messages.create(
        model=MODEL, max_tokens=200,
        system=SYSTEM,
        messages=[{"role": "user", "content": f"Passages:\n{context}\n\nQuestion: {query}"}],
    )
    text = "".join(b.text for b in resp.content if b.type == "text")
    return {"answer": text, "citations": [int(c) for c in re.findall(r"\[(\d+)\]", text)]}


answers = [answer(q, top) for q, top in zip(QUERIES, top_ks)]
for q, a in zip(QUERIES, answers):
    print(f"Q: {q}")
    print(f"A: {a['answer']}\n")

Citation accuracy#

A citation is grounded if the cited passage was actually retrieved; correct if it points at the ground-truth passage.

grounded = []
correct  = []
for a, top, t in zip(answers, top_ks, TRUTH):
    cs = a["citations"] or [0]
    grounded.append(all(c in top for c in cs))
    correct.append(t in cs)

g_rate = sum(grounded) / len(grounded)
c_rate = sum(correct)  / len(correct)
print(f"grounded citations: {g_rate:.0%}    truth-match citations: {c_rate:.0%}")

Visualisation#

import matplotlib.pyplot as plt

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 3.4))

methods = ["BM25", "Dense", "RRF"]
def recall_at_k(score_lists, k=3):
    out = []
    for scs, t in zip(score_lists, TRUTH):
        top = sorted(range(len(scs)), key=lambda i: -scs[i])[:k]
        out.append((t - 1) in top)
    return sum(out) / len(out)

recalls = [recall_at_k(bm25_scores), recall_at_k(dense_scores), recall_at_k(fused)]
ax1.bar(methods, [r * 100 for r in recalls], color=["tab:blue", "tab:orange", "tab:green"])
ax1.set_ylim(0, 105); ax1.set_ylabel("recall@3 (%)")
ax1.set_title("retrieval — single vs hybrid")
for i, v in enumerate(recalls):
    ax1.text(i, v * 100 + 1, f"{v:.0%}", ha="center")

metrics = ["grounded", "truth-match"]
ax2.bar(metrics, [g_rate * 100, c_rate * 100], color=["tab:purple", "tab:green"])
ax2.set_ylim(0, 105); ax2.set_ylabel("rate (%)")
ax2.set_title("citation quality")
for i, v in enumerate([g_rate, c_rate]):
    ax2.text(i, v * 100 + 1, f"{v:.0%}", ha="center")

fig.tight_layout(); plt.show()

Checks#

s.check(
    "rrf_recall_at_least_match_best_single",
    lambda: recall_at_k(fused) >= max(recall_at_k(bm25_scores), recall_at_k(dense_scores)),
    msg=f"RRF={recall_at_k(fused):.0%} BM25={recall_at_k(bm25_scores):.0%} Dense={recall_at_k(dense_scores):.0%}",
)
s.check(
    "every_query_has_cited_answer",
    lambda: all(len(a["citations"]) >= 1 for a in answers),
    msg=f"citation counts = {[len(a['citations']) for a in answers]}",
)
s.check(
    "all_citations_grounded_in_retrieved",
    lambda: g_rate == 1.0,
    msg=f"grounded = {g_rate:.0%}",
)
s.check(
    "truth_citation_rate_at_least_80pct",
    lambda: c_rate >= 0.80,
    msg=f"truth-match = {c_rate:.0%}",
)
s.check(
    "rrf_strictly_beats_bm25_or_dense",
    lambda: recall_at_k(fused) >= recall_at_k(bm25_scores),
    msg="hybrid should never lose to BM25 alone",
)

Notes for production#

Reranker after RRF: a cross-encoder like BAAI/bge-reranker-v2-m3 on the top-20 RRF results then taking top-5 typically lifts MRR by 8–12 points on long-tail queries. Worth one extra forward pass per request.
BGE-M3 multi-vector: replaces dense+sparse+colbert with a single model that emits all three. One encoder, three retrievals, RRF. State of the art on BEIR as of late 2025.
Citation enforcement: when the model omits citations, fall back to attributing each sentence by max-cosine to the retrieved set. For high-stakes systems run a second LLM as a faithfulness judge (this is what RAGAS does internally).
Chunking matters more than the embedding model. Sentence-window or semantic-chunk a long doc; never chunk on a fixed character count without overlap.

s.summary()
s.save()