Lab 01 — ANN Index + RAG Retrieval Engine

Phase: 11 — Retrieval-Augmented Generation & Vector Search Difficulty: ⭐⭐⭐☆☆ (the algorithms are tractable; the retrieval-quality judgment is ⭐⭐⭐⭐⭐) Time: 4–5 hours

"RAG" is two systems wearing one acronym: a retriever that finds the right chunks, and a generator that writes the answer. Almost every RAG failure is a retrieval failure — and you cannot debug what you cannot measure separately. This lab builds the retrieval half from the metal up: the similarity metrics (cosine / dot / l2_normalize), an exact brute-force kNN index (FlatIndex, the ground truth), an approximate inverted-file index (IVFIndex — cluster the space with seeded k-means, then probe only the nprobe nearest cells), the retrieval-quality metrics that grade the index without a generator (recall_at_k, ndcg_at_k), chunk_text with overlap, MMR re-ranking for diversity, and a budgeted, cited assemble_rag_prompt. Two facts are the soul of the lab and you will prove them: IVF recall climbs to the exact answer as nprobe rises (and equals Flat when nprobe == nlist), and MMR refuses to return five paraphrases of the same sentence.

What you build

• cosine_similarity / dot / l2_normalize — the geometry of "semantically close". Cosine is magnitude-invariant; l2_normalize guards the zero vector (a real corpus has degenerate embeddings and a divide-by-zero in your indexer is a 3 a.m. page).
• FlatIndex — exact kNN: add(id, vector) and search(query, k) scoring the query against every vector. O(n·d), does not scale — and it is the recall oracle every approximate index is graded against.
• IVFIndex(nlist, seed) — inverted-file ANN: train(vectors) builds nlist centroids with deterministic seeded k-means, add drops each vector into its nearest centroid's bucket, search(query, k, nprobe) compares the query only to the nprobe nearest cells. Recall is a single knob: nprobe=1 is fast and lossy; nprobe=nlist probes everything ⇒ exactly equal to Flat.
• recall_at_k / ndcg_at_k — grade retrieval separately from generation: recall = "did we find the right chunks?"; nDCG = "did we rank them well?".
• chunk_text(text, chunk_size, overlap) — overlapping character windows; rejects overlap >= chunk_size (the infinite-loop bug people actually ship).
• mmr_rerank(query_vec, candidates, lambda_, k) — Maximal Marginal Relevance: greedily balance relevance against diversity so near-duplicates don't crowd out the one chunk with the new fact.
• assemble_rag_prompt(question, retrieved_chunks, max_context) — a grounded prompt with numbered [n] citations, truncated to a context budget.

Key concepts

Files

Run

pip install -r requirements.txt
pytest test_lab.py -v                       # against your lab.py
LAB_MODULE=solution pytest test_lab.py -v   # against the reference
python solution.py                          # the worked example (watch recall climb)

Success criteria

• All tests pass against your implementation (and LAB_MODULE=solution).
• You can explain why cosine is preferred over raw dot for embeddings, and why production normalizes once and then uses dot.
• You can explain why test_ivf_recall_increases_monotonically_with_nprobe and test_ivf_equals_flat_when_nprobe_is_nlist are the soul of the ANN lesson — the recall/latency dial and the fact that "search everything" recovers exactness.
• You can explain why test_mmr_picks_diverse_not_near_duplicates is the soul of the reranking lesson — pure top-k returns five paraphrases; MMR trades a sliver of relevance for the chunk that carries a new fact.
• You can state, on a whiteboard, how to evaluate a RAG system: recall@k / nDCG / MRR for the retriever, faithfulness / citation accuracy for the generator, and why you measure them separately.

How this maps to the real stack

Limits of the miniature (say these out loud in the interview): k-means here is plain Lloyd's with a simple empty-cluster policy (FAISS re-seeds and uses better init); IVF stores full vectors (real systems add PQ to compress them — the memory leg of the recall/latency/memory triangle); there is no HNSW graph index here (the other dominant ANN family — navigable small-world graphs with M/ef knobs — covered in the WARMUP); chunking is by character, not token or semantic boundary; and the prompt budget is in characters, not real tokenizer tokens.

Extensions (build these next)

• Add IVF-PQ: quantize each vector to m subspace codebooks and store codes instead of floats; measure the recall and memory drop (the third leg of the triangle).
• Add an HNSW index (layered graph, greedy descent, ef_search knob) and plot its recall/latency curve against IVF on the same corpus.
• Add hybrid search: a BM25 lexical scorer fused with the dense score (reciprocal-rank fusion), and show it rescues exact-match queries dense retrieval misses (IDs, error codes, rare names).
• Add MRR and a tiny labelled query set; compare MMR vs a (mock) cross-encoder reranker on nDCG.
• Swap the character chunker for a token chunker (reuse your Phase-01 tokenizer) and measure the recall change at fixed budget.

Interview / resume

• Talking points: "How does IVF trade latency for recall, and what is nprobe?" "Your RAG answer is wrong — is it a retrieval or a generation failure, and how do you tell?" "Why does naive top-k return near-duplicates, and what is MMR?" "How do you evaluate retrieval separately from generation?" "When does RAG beat fine-tuning?"
• Resume bullet: Built a from-scratch RAG retrieval engine — exact (Flat) and approximate (IVF, seeded k-means + nprobe) vector indexes, recall@k / nDCG retrieval metrics, overlap-aware chunking, MMR diversity re-ranking, and a budgeted, citation-grounded prompt assembler — demonstrating the recall/latency trade and retrieval-vs-generation evaluation.