Lab 03 — Embed & the Full Two-Stage RAG Pipeline

Phase 31 · Lab 03 · Phase README · Warmup

The problem

Cohere sells its platform as three composable APIs — Embed for vectors, Rerank for precision, Chat with documents for grounded answers — and every serious deployment wires them into one pipeline. This lab builds that pipeline end to end, plus the two Embed-specific ideas people consistently get wrong:

1. input_type is not decoration. Embed v3 embeds a query differently from a document — same text, different vector — because short keyword-ish queries and long redundant documents come from different distributions, and the model aligns them into one retrieval geometry. Index your corpus as search_document, embed queries as search_query; use the wrong one and nothing crashes — retrieval quality just silently sags. The miniature makes the asymmetry mechanical: document vectors are hashed token counts (frequency is document signal), query vectors are hashed content-term presence (a query term votes once) — different encoders, one shared hash space, so cross-type cosine is meaningful.
2. Compressed embeddings are a storage/recall tradeoff you can measure. int8 (4× smaller, near-identical cosine), binary (32× smaller, coarser), and Matryoshka-style truncation (shorter prefixes of the same vector) all keep the easy retrievals and start losing the hard ones — the lab shows exactly where.

Then the pipeline: retrieve wide (cosine top-k over the index), rerank narrow (Lab 01's cross-encoder idea), generate grounded (Lab 02's citations idea), and prove the two-stage ordering beats retrieve-only on a planted corpus.

What you build

File map

Run it

pytest test_lab.py -v                       # your lab.py (red until you implement)
LAB_MODULE=solution pytest test_lab.py -v   # the reference (green)
python solution.py                          # worked example

Success criteria

• The same text embeds to different vectors under search_query vs search_document (all four input_types pairwise distinct), and an unknown input_type raises — the typo that silently degrades real pipelines fails loudly here.
• First-stage cosine retrieval surfaces both relevant docs, scores descending, deterministic tie-breaks; duplicate ids and bad k are rejected.
• The planted two-stage win: first-stage top-1 is the term-stuffed kb_policy; after rerank, top-1 is the verbatim how-to kb_howto.
• End to end, answer() returns a grounded answer whose citations index the answer text exactly and cite the promoted document; the scripted hallucination never appears and the grounding score reflects it.
• int8 and binary quantizations retrieve the same top-1 as float on the easy case; int8 cosine ≈ float cosine; 48-dim truncation keeps the easy top-1, 16-dim loses it.
• Everything deterministic: same inputs → identical RagResponse, twice.
• All 27 tests pass under both lab and solution.

How this maps to the real stack

• Embed v3 / v4: co.embed(texts=..., model="embed-v4.0" or "embed-english-v3.0", input_type="search_document" | "search_query" | "classification" | "clustering", embedding_types=["float", "int8", "binary", ...]). The input_type values here are Cohere's own enum, and the real models are trained so query- and document-mode embeddings align across the asymmetry — our two hand-built encoders over a shared hash space are the deterministic stand-in for that trained alignment. Embed v4 adds Matryoshka-style selectable output dimensions; truncate_matryoshka mirrors the usage pattern (short vectors for cheap shortlisting, full vectors for precision).
• Compression is a real Cohere feature, not an add-on: Embed v3/v4 natively return int8 and binary embedding types, and Cohere's guidance is exactly this lab's finding — int8 costs almost nothing, binary buys ~32× storage at a recall cost you should measure on your corpus before committing an index format.
• The pipeline shape is the reference architecture: first-stage ANN retrieval (pgvector / OpenSearch / a managed store) over search_document vectors, Cohere Rerank over the top 50–100, Cohere Chat with documents over the top handful, citations returned to the user. This is the exact stack the Phase 05 hybrid retriever feeds and the Phase 11 eval harness scores (recall@k for stage one, faithfulness/citation accuracy for stage three).
• Observability lives in the response: RagResponse carrying first_stage_ids and reranked_ids is the miniature of logging every stage's candidates — the only way to debug "why did the answer cite the wrong doc" is knowing what each stage saw and kept.

Limits of the miniature. The hashing embedder captures token overlap, not semantics — real Embed places paraphrases near each other, ours cannot ("booked" ≠ "book" here, deliberately: no stemming, so the lexical mechanism stays visible). Real Matryoshka truncation works because the model was trained to front-load information; truncating a hash-bucketed vector just drops random buckets, which is why the lab frames truncation as an honest tradeoff demo, not a production recipe. And brute-force kNN is O(N·dim); production uses HNSW/IVF behind the same search contract.

Extensions (your own machine)

• Swap embed for real Cohere Embed calls (or a local sentence-transformers bi-encoder with query/passage prefixes, e.g. the e5 family's "query: "/"passage: " convention — the same asymmetry idea) behind the same signature and re-run the pipeline.
• Implement packed binary embeddings (int.from_bytes, popcount Hamming similarity) and measure the actual 32× memory ratio.
• Add RRF fusion of BM25 + dense as the first stage (Phase 05) and show retrieve→rerank on top of hybrid first-stage beats both.
• Evaluate properly: build 20 labeled queries over your own docs and report recall@5 for first-stage vs two-stage, and citation coverage for the end-to-end answers.

Interview / resume signal

"Built Cohere's full RAG stack as a deterministic miniature: an asymmetric embedder where input_type genuinely changes the encoding (search_query vs search_document), int8/binary/ Matryoshka-compressed vectors with measured recall tradeoffs, cosine first-stage retrieval, a cross-encoder rerank stage that provably corrects the first stage's top-1, and grounded generation that emits Cohere-shaped citations and drops unsupported claims — with per-stage candidate observability in the response."