""" Retrieve — given a question, return the most relevant calendar excerpts. This is the "R" in RAG. Phase 2 turns it into a two-stage funnel: stage 1 — RECALL dense (vector) and sparse (BM25) search each nominate ~50 candidates; Reciprocal Rank Fusion merges the two rankings. stage 2 — PRECISION a cross-encoder re-scores the fused shortlist jointly with the query; the final top-k come out the top. Which stages run is a config knob (UBCAL_RETRIEVAL_MODE), so the eval can compare: dense Phase-1 pure vector search (the baseline) hybrid stage 1 only (dense + BM25, RRF-fused) hybrid_rerank both stages (the Phase-2 default) Why RRF for fusion? Cosine similarities and BM25 scores live on totally different scales, so averaging them is meaningless. RRF ignores the scores entirely and combines *ranks* — a document's fused score is sum(1 / (RRF_K + rank)) across the lists it appears in — sidestepping the calibration problem. Exposing the doc `id` on every result is deliberate: the Phase 3 evaluation harness matches retrieved ids against "gold" chunk ids to compute hit@k / recall@k. """ from __future__ import annotations import json import sys from dataclasses import dataclass import numpy as np from . import bm25, config, embed, trace @dataclass class Result: """One retrieved excerpt (chunk record) plus its retrieval score. What `score` means depends on the mode that produced it: cosine similarity (dense), RRF fused score (hybrid), or a cross-encoder logit (hybrid_rerank). Higher is always better within one result list, but values are not comparable across modes. """ id: int score: float recipe: dict # the full metadata record from the index # The index is loaded once and cached for the process lifetime. _embeddings: np.ndarray | None = None _metadata: list[dict] | None = None _texts: list[str] | None = None def _load_index() -> tuple[np.ndarray, list[dict]]: global _embeddings, _metadata if _embeddings is None or _metadata is None: if not config.EMBEDDINGS_PATH.exists() or not config.METADATA_PATH.exists(): sys.exit( "No index found. Build it first with: python -m src.ingest" ) _embeddings = np.load(config.EMBEDDINGS_PATH) with open(config.METADATA_PATH, "r", encoding="utf-8") as f: _metadata = json.load(f) return _embeddings, _metadata def _corpus_texts() -> list[str]: """The raw document texts, id-aligned with the index (what BM25 searches over).""" global _texts if _texts is None: _, metadata = _load_index() _texts = [r["text"] for r in metadata] return _texts def _dense_search(query: str, k: int) -> list[tuple[int, float]]: """Embed the query and rank every chunk by cosine similarity; top-k, best first.""" embeddings, metadata = _load_index() # Embed the query (shape (1, dim)) and reduce to a 1-D vector. query_vec = embed.embed_texts([query])[0] # Cosine similarity == dot product, since every vector is unit-length. # scores[i] is the similarity between the query and chunk i. scores = embeddings @ query_vec # Indices of the top_k highest scores. argpartition finds them without fully # sorting the whole array; we then sort just those k by score, descending. k = min(k, len(metadata)) top_idx = np.argpartition(scores, -k)[-k:] top_idx = top_idx[np.argsort(scores[top_idx])[::-1]] return [(int(i), float(scores[i])) for i in top_idx] def _rrf_fuse(rankings: list[list[int]]) -> list[tuple[int, float]]: """Merge ranked doc-id lists with Reciprocal Rank Fusion; best fused score first.""" fused: dict[int, float] = {} for ranking in rankings: for rank, doc_id in enumerate(ranking, start=1): fused[doc_id] = fused.get(doc_id, 0.0) + 1.0 / (config.RRF_K + rank) return sorted(fused.items(), key=lambda pair: pair[1], reverse=True) def _as_results(ranked: list[tuple[int, float]], limit: int) -> list[Result]: _, metadata = _load_index() return [Result(id=i, score=s, recipe=metadata[i]) for i, s in ranked[:limit]] def _candidates(ranked: list[tuple[int, float]]) -> list[dict]: """Render a ranked (id, score) list for the trace event stream. Just enough to show a candidate column in a UI — the full record follows on the final results, so there is no point carrying 20 chunk bodies per stage. """ _, metadata = _load_index() return [ {"id": i, "score": round(float(s), 6), "title": metadata[i]["title"]} for i, s in ranked ] def warmup() -> int: """Load the index and build the BM25 postings up front; returns the chunk count. Both are lazily cached module globals, so without this the *first* query of a process pays the 22 MB .npy load plus the in-memory BM25 build. Fine for the CLI, bad for the first visitor to a web demo. """ _, metadata = _load_index() bm25.search("warmup", 1, _corpus_texts()) return len(metadata) def retrieve(query: str, top_k: int | None = None, mode: str | None = None) -> list[Result]: """Return the top_k excerpts most relevant to `query`, best first. `mode` defaults to config.RETRIEVAL_MODE (see the module docstring for the three modes). The signature and Result shape are stable across modes — the eval harness and CLI call this one function regardless of pipeline. """ top_k = top_k or config.TOP_K mode = mode or config.RETRIEVAL_MODE if mode not in ("dense", "hybrid", "hybrid_rerank"): sys.exit( f"Unknown retrieval mode {mode!r} — expected dense, hybrid, or hybrid_rerank." ) # The query is echoed so a consumer can attribute this whole block to the right hop of # a multi-hop agent run without having to infer boundaries from the surrounding steps. trace.event("retrieval_start", query=query, mode=mode, top_k=top_k) trace.detail("retrieve mode", mode) trace.detail("embed query (OpenAI API)", repr(query)) if mode == "dense": dense_hits = _dense_search(query, top_k) trace.event("candidates", stage="dense", items=_candidates(dense_hits)) out = _as_results(dense_hits, top_k) trace.results(out) trace.event("retrieval_final", query=query, items=[trace.describe(r) for r in out]) return out # Stage 1 — recall: each retriever nominates candidates, RRF merges the rankings. dense_hits = _dense_search(query, config.DENSE_K) sparse_hits = bm25.search(query, config.SPARSE_K, _corpus_texts()) dense_ids = [i for i, _ in dense_hits] sparse_ids = [i for i, _ in sparse_hits] fused = _rrf_fuse([dense_ids, sparse_ids]) trace.event("candidates", stage="dense", items=_candidates(dense_hits)) trace.event("candidates", stage="bm25", items=_candidates(sparse_hits)) trace.event("candidates", stage="rrf", items=_candidates(fused)) trace.detail("stage 1 recall", f"dense={len(dense_ids)} + bm25={len(sparse_ids)} " f"-> RRF-fused to {len(fused)} unique") if mode == "hybrid": out = _as_results(fused, top_k) trace.results(out) trace.event("retrieval_final", query=query, items=[trace.describe(r) for r in out]) return out # Stage 2 — precision: the cross-encoder re-scores the fused shortlist. # Imported here so dense/hybrid modes never pay the multi-second torch import. from . import rerank shortlist = _as_results(fused, config.RERANK_CANDIDATES) trace.event("shortlist", ids=[r.id for r in shortlist]) trace.detail("stage 2 precision", f"rerank {len(shortlist)} candidates (Cohere API) " f"-> top {top_k}") out = rerank.rerank(query, shortlist, top_k) trace.results(out) trace.event("retrieval_final", query=query, items=[trace.describe(r) for r in out]) return out if __name__ == "__main__": # Quick manual check: python -m src.retrieve "prerequisites for CPSC 210" # (set UBCAL_RETRIEVAL_MODE to compare modes) question = " ".join(sys.argv[1:]) or "prerequisites for CPSC 210" for r in retrieve(question): print(f"[{r.score:.3f}] (id={r.id}) {r.recipe['title']}")