"""Retrieval backends: dense, dense + query rewriting, hybrid, and reranked. Modes: dense - embedding similarity only (the v0 baseline) dense_rw - LLM expands the question with likely code identifiers and file names BEFORE embedding it (query rewriting) dense_rerank - dense candidates re-scored by a cross-encoder hybrid - dense + BM25 keyword search, fused with weighted RRF hybrid_rerank - hybrid candidates re-scored by a cross-encoder Ablation findings on this corpus (42-question eval set): - dense won over hybrid: BM25 pulled in issue chunks that share common query words, displacing correct chunks from the top-5. - Two reranker generations (bge-reranker-base and v2-m3) both LOWERED MRR: cross-encoders demote raw code chunks for "where is X defined" questions because code *is* the answer without discussing it. - dense_rw targets the remaining misses: questions whose words are all too common to retrieve well ("Query, Path, and Body parameter functions") get expanded with identifiers an LLM can guess. """ import json import re from pathlib import Path import chromadb from rank_bm25 import BM25Okapi from sentence_transformers import CrossEncoder, SentenceTransformer DATA_DIR = Path("data") EMBED_MODEL = "BAAI/bge-small-en-v1.5" RERANK_MODEL = "BAAI/bge-reranker-v2-m3" QUERY_PREFIX = "Represent this sentence for searching relevant passages: " CANDIDATES = 20 # how many candidates hybrid/rerank gather before final cut RRF_K = 60 # standard RRF constant RRF_WEIGHTS = (2.0, 1.0) # dense vs BM25 in the fused ranking # --- query rewriting --- REWRITE_CACHE = DATA_DIR / "rewrite_cache.json" REWRITE_PROMPT = """\ You improve search queries for a search engine over the FastAPI repository (source code, docs, GitHub issues). Rewrite the question as a short search query, adding likely Python identifiers, class/function names, and file names from FastAPI. Output ONLY the query, no explanation. Question: {question}""" # Lazy singletons so models/indexes load once per process, not per query _embedder = None _loaded = False _reranker = None _collection = None _bm25 = None _chunk_ids: list[str] = [] _chunk_by_id: dict[str, dict] = {} def _tokenize(text: str) -> list[str]: return re.findall(r"[a-z0-9_]+", text.lower()) def _load() -> None: global _loaded, _embedder, _collection, _bm25, _chunk_ids, _chunk_by_id if _loaded: return _embedder = SentenceTransformer(EMBED_MODEL) _collection = chromadb.PersistentClient( path=str(DATA_DIR / "chroma")).get_collection("chunks") chunks = [json.loads(line) for line in (DATA_DIR / "chunks.jsonl").read_text().splitlines()] _chunk_ids = [c["id"] for c in chunks] _chunk_by_id = {c["id"]: c for c in chunks} _bm25 = BM25Okapi([_tokenize(c["text"]) for c in chunks]) _loaded = True def _get_reranker() -> CrossEncoder: global _reranker if _reranker is None: _reranker = CrossEncoder(RERANK_MODEL) # first run downloads ~2.3GB return _reranker def _rewrite(question: str) -> str: """Expand the question with likely identifiers via one cached LLM call. Falls back to the original question on any failure (no key, rate limit, network) so retrieval never breaks because of the rewriter. """ import os cache = (json.loads(REWRITE_CACHE.read_text()) if REWRITE_CACHE.exists() else {}) if question in cache: return cache[question] try: from groq import Groq client = Groq(api_key=os.environ["GROQ_API_KEY"]) model = os.environ.get("GROQ_MODEL", "openai/gpt-oss-120b") out = client.chat.completions.create( model=model, messages=[{"role": "user", "content": REWRITE_PROMPT.format(question=question)}], temperature=0.0).choices[0].message.content.strip() expanded = f"{question} {out}" except Exception as e: print(f" [rewrite failed ({e.__class__.__name__}); using original]") return question cache[question] = expanded REWRITE_CACHE.write_text(json.dumps(cache, ensure_ascii=False, indent=2)) return expanded def _dense_ids(question: str, n: int, with_scores: bool = False): emb = _embedder.encode(QUERY_PREFIX + question, normalize_embeddings=True) res = _collection.query(query_embeddings=[emb.tolist()], n_results=n) ids = res["ids"][0] if with_scores: # cosine distance -> similarity (1.0 = identical, 0.0 = unrelated) return ids, [1.0 - d for d in res["distances"][0]] return ids def _bm25_ids(question: str, n: int) -> list[str]: scores = _bm25.get_scores(_tokenize(question)) ranked = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True) return [_chunk_ids[i] for i in ranked[:n]] def _rrf_fuse(rankings: list[list[str]], weights: tuple) -> list[str]: """Weighted RRF: score(id) = sum of w/(RRF_K + rank) over lists.""" scores: dict[str, float] = {} for ranking, w in zip(rankings, weights): for rank, cid in enumerate(ranking, 1): scores[cid] = scores.get(cid, 0.0) + w / (RRF_K + rank) return sorted(scores, key=scores.get, reverse=True) def _rerank(question: str, candidates: list[str]) -> list[str]: pairs = [(question, _chunk_by_id[cid]["text"][:2000]) for cid in candidates] scores = _get_reranker().predict(pairs) return [cid for _, cid in sorted(zip(scores, candidates), reverse=True)] def _to_hit(cid: str) -> dict: c = _chunk_by_id[cid] return { "id": cid, "text": c["text"], "meta": { "source_type": c["source_type"], "path": c["path"], "symbol": c["symbol"] or "", "url": c["url"], }, } def retrieve(question: str, k: int = 5, mode: str = "dense") -> list[dict]: _load() if mode == "dense": ids, scores = _dense_ids(question, k, with_scores=True) hits = [_to_hit(cid) for cid in ids] for h, s in zip(hits, scores): h["score"] = s return hits elif mode == "dense_rw": ids = _dense_ids(_rewrite(question), k) elif mode == "dense_rerank": ids = _rerank(question, _dense_ids(question, CANDIDATES))[:k] elif mode in ("hybrid", "hybrid_rerank"): fused = _rrf_fuse([_dense_ids(question, CANDIDATES), _bm25_ids(question, CANDIDATES)], weights=RRF_WEIGHTS) candidates = fused[:CANDIDATES] if mode == "hybrid_rerank": candidates = _rerank(question, candidates) ids = candidates[:k] else: raise ValueError(f"unknown mode: {mode}") return [_to_hit(cid) for cid in ids]