Retrieval that actually works

import math
import re
from collections import Counter

def tokenize(text):
    return re.findall(r"[a-z0-9_]+", text.lower())

class BM25:
    def __init__(self, texts, k1=1.5, b=0.75):
        self.k1, self.b = k1, b
        self.docs = [Counter(tokenize(t)) for t in texts]
        self.lengths = [sum(d.values()) for d in self.docs]
        self.avg = sum(self.lengths) / len(self.docs)
        df = Counter(term for d in self.docs for term in d)
        n = len(self.docs)
        self.idf = {t: math.log(1 + (n - f + 0.5) / (f + 0.5)) for t, f in df.items()}

    def score(self, query, i):
        d, s = self.docs[i], 0.0
        for t in set(tokenize(query)):
            tf = d.get(t, 0)
            if tf:
                norm = 1 - self.b + self.b * self.lengths[i] / self.avg
                s += self.idf[t] * tf * (self.k1 + 1) / (tf + self.k1 * norm)
        return s

docs = ["deadlock deadlock deadlock: InnoDB error 1213 rolls back one transaction",
        "a deadlock happens when two transactions wait for each other",
        "indexes speed up lookups; a transaction groups statements"]
bm = BM25(docs)
print("idf:", {t: round(bm.idf[t], 2) for t in ("deadlock", "transaction", "1213")})
for q in ["deadlock", "error 1213", "transactions waiting for each other"]:
    print(f"{q!r:40}", [round(bm.score(q, i), 2) for i in range(len(docs))])

from mini_rag import GOLDEN, HybridIndex, lessons_in
from site_docs import load_sections

docs = load_sections()
index = HybridIndex(docs)

def evaluate(mode, k=5):
    hits, rr, found = 0, 0.0, {}
    for question, answers in GOLDEN:
        lessons = lessons_in(index.rank(question, mode), docs, k=10)
        found[question] = bool(answers & set(lessons[:k]))
        hits += found[question]
        rank = next((i for i, url in enumerate(lessons, 1) if url in answers), None)
        rr += 1 / rank if rank else 0
    return hits / len(GOLDEN), rr / len(GOLDEN), found

results = {mode: evaluate(mode) for mode in ("bm25", "dense", "hybrid")}
for mode, (recall, mrr, _) in results.items():
    print(f"{mode:7} recall@5={recall:.2f}  MRR@10={mrr:.2f}")

bm, de = results["bm25"][2], results["dense"][2]
print("\nonly BM25 found: ", [q for q in bm if bm[q] and not de[q]])
print("only dense found:", [q for q in bm if de[q] and not bm[q]])
print("both missed:     ", [q for q in bm if not bm[q] and not de[q]])

from sentence_transformers import CrossEncoder

reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")      # small, runs on CPU

def rerank(question, hits, keep=5):
    scores = reranker.predict([(question, h.doc["text"]) for h in hits])
    ranked = sorted(zip(scores, hits), key=lambda pair: -pair[0])
    return [h for _, h in ranked[:keep]]

candidates = index.search(question, k=50)          # stage 1: fast, wide
top5 = rerank(question, candidates)                # stage 2: slow, precise

from fake_openai import fake_client
from mini_rag import HybridIndex, lessons_in
from site_docs import load_sections

docs = load_sections()
index = HybridIndex(docs)
client = fake_client(script=["stream the response token by token (streaming events, output text deltas)"])

REWRITE = ("Rewrite the user's question as a search query for technical documentation. "
           "Use the standard technical terms. Reply with the query only.")

question = "show tokens to the user as they are generated"
rewritten = client.responses.create(model="gpt-5.6-luna", instructions=REWRITE, input=question).output_text

for label, q in [("original ", question), ("rewritten", rewritten)]:
    print(label, lessons_in(index.rank(q), docs, k=3))

from mini_rag import HybridIndex
from site_docs import load_sections

docs = load_sections()
index = HybridIndex(docs)
q = "how to read EXPLAIN output"

print("no filter:      ", [h.doc["url"] for h in index.search(q, k=4)])
print("course = mysql: ", [h.doc["url"] for h in index.search(q, k=4, where={"course": "mysql"})])

def one_per_lesson(hits, k):
    seen, out = set(), []
    for h in hits:
        lesson = h.doc["url"].split("#")[0]
        if lesson not in seen:
            seen.add(lesson)
            out.append(h)
    return out[:k]

print("one per lesson: ", [h.doc["url"] for h in one_per_lesson(index.search(q, k=20), 4)])