Skip to content

Retrieval lab

Retrieval lab

You build retrieval_lab, a small hybrid retrieval pipeline over Postgres: a Markdown-aware chunker, a deterministic hashing embedder, a pgvector table with a generated full-text column, vector and full-text search fused with reciprocal rank fusion (RRF), a reranker with MMR, retrieval metrics on a golden set, and an answer assembler that cites its sources. The last stages swap in local models (fastembed) and ask Claude for cited answers.

The corpus is fixtures/corpus: 29 Markdown documents (a snapshot of course lessons and interview-angle pages) and three public-domain stories, with golden.yaml, 48 queries of three kinds (factual, conceptual, lexical). A chunk is relevant to a golden query when it belongs to the query’s doc_id and its text contains the query’s anchor, case-insensitive and whitespace-normalised. See fixtures/corpus/README.md.

retrieval_lab/types.py is given and shared by every stage (Document, Chunk, Hit, Ranked, GoldenQuery, Metrics, Draft, Faithfulness, Answer). Write the stage modules against it; do not change it.

How to run

Run every command below in a terminal at the root of your clone of this repository, after a one-time uv sync there; uv run miubais doctor lists what this machine has.

uv run miubais start m4-retrieval-lab

This copies course/m4-retrieval-lab/starter/ to workspace/m4-retrieval-lab/ at the repo root. That copy is yours: the stubs under workspace/m4-retrieval-lab/src/retrieval_lab/ raise NotImplementedError until you fill them in, and the judgment stage asks for a file in the workspace root, named in its description below. The tests stay in course/m4-retrieval-lab/tests/ and import your package from the workspace. start --force replaces the workspace with a fresh copy.

uv run miubais check m4-retrieval-lab

Runs the stages in order and stops at the first failure, printing its pytest output, so stage 02 waits for stage 01. --stage <stage id> runs one stage and --all keeps going past a failure. A stage whose requires (in stages.yaml) are missing on your machine is skipped, not failed. Results are kept in ~/.miubais/progress.json; uv run miubais reset m4-retrieval-lab forgets them (--all for every exercise), and start --force does the same along with the fresh copy.

uv run miubais hint m4-retrieval-lab

Takes the last stage check saw fail, its test output and your diff against the starter, and asks the tutor on your Claude Code login for a nudge of under 90 words that is checked not to hand you the solution. If the model cannot be reached, you get a pointer to the failing test instead.

The Run check and Hint buttons on the site’s exercise page do the same while uv run miubais serve is running, and share the progress file with the terminal.

Stages

  1. 01-chunker: chunker.load_documents(corpus_dir) returns one Document(doc_id, title, text) per file of docs/*.md then stories/*.txt, each sorted by name; doc_id is the file stem and title the first H1 (or the first non-blank line of a story). chunker.chunk(doc, *, max_chars=1200, overlap=1) follows these rules:

    • A Markdown heading (# to ###### followed by a space, outside a fence) starts a new section; the heading line is the first block of its section and heading_path joins the open headings with > (Title > Section A > Sub A.1; text before the first heading has ""). A chunk never spans two sections.
    • Inside a section, blocks are separated by blank lines. A fenced code block (a line starting with three or more backticks or tildes, up to the closing line of the same character) is one block, blank lines and # lines included; an unclosed fence runs to the end of the document.
    • Blocks are packed greedily: a chunk grows while the slice from its first block’s start to the next block’s end stays within max_chars. A block longer than max_chars (a long fence, a long list) becomes one oversized chunk; it is never split.
    • The next chunk starts overlap blocks before the previous one ended, but only when the carried blocks still leave room for at least one new block; otherwise it starts with the new block.
    • start_char/end_char are offsets into doc.text and doc.text[start_char:end_char] == text; chunk_id is f"{doc_id}#{index:03d}", numbered from 0 across the document.

    The tests check heading paths, packing with and without overlap on invented documents, that a 2,000-character fence is one chunk, offsets, ids and determinism on the corpus, that every non-blank character is covered, that no chunk splits a fence, and that every golden anchor lies inside at least one chunk of its document (so recall can reach 1.0).

  2. 02-embed-and-store (Docker): embed.Embedder is a protocol with dims, name and embed(texts) -> np.ndarray of shape (len(texts), dims). embed.HashEmbedder(dims=256) is pinned exactly: tokens are the lowercase words of re.findall(r"\w+", text.lower()), then the word bigrams f"{a} {b}", then every three-character slice of every word; each token is hashed with hashlib.blake2b(token.encode(), digest_size=8), read as a big-endian integer h, and adds +1 at index h % dims (or -1 when the top bit h >> 63 is set); the vector is L2-normalised (float32). It must give the same vector in every process, so no built-in hash(). embed.FastEmbedder(model="BAAI/bge-small-en-v1.5", cache_dir=None) has dims 384 and a name that includes the model; it imports fastembed only when it embeds (written in stage 08). store.Schema.create(conn, dims) creates, if missing, the vector extension and

    CREATE TABLE chunks (chunk_id text PRIMARY KEY, doc_id text, heading_path text,
    start_char int, end_char int, text text,
    tsv tsvector GENERATED ALWAYS AS (to_tsvector('english', text)) STORED,
    embedding vector(<dims>));
    CREATE INDEX ON chunks USING hnsw (embedding vector_cosine_ops);
    CREATE INDEX ON chunks USING gin (tsv);

    and raises ValueError when chunks already exists with a different dimension. store.upsert(conn, chunks, embedder, batch=64) embeds and writes in batches with INSERT ... ON CONFLICT (chunk_id) DO UPDATE, so a second run changes nothing and an edited chunk is replaced; it raises ValueError before writing when embedder.dims differs from the table. store.count(conn) returns the row count. The tests check the pinned recipe against an independent implementation, determinism across processes with different PYTHONHASHSEED, unit norm, the columns and both indexes, idempotent upserts, the stored vectors, and both mismatches.

  3. 03-hybrid-retrieval (Docker): every search returns Hit(chunk_id, score, rank) with ranks from

    1. hybrid.vector_search(conn, embedder, query, n=50) orders by cosine distance <=> and scores 1 - distance, and returns [] for a query without words (a zero vector has no cosine); an HNSW scan returns at most hnsw.ef_search rows (default 40), so raise it to n for the query (SET LOCAL or set_config(..., true) inside a transaction). hybrid.text_search(conn, query, n=50) parses the query with websearch_to_tsquery('english', query), which never raises on user input, then splits its text form at the top-level & and |: the positive terms (words and quoted phrases) are ORed into one query q, and each negated term (-word, -"a phrase", printed as !...) stays an exclusion, tsv @@ q AND NOT tsv @@ (neg1 | neg2 ...). websearch_to_tsquery ANDs every word, and a question in natural language rarely has all its words in one chunk; turning every & into | would instead turn !term into “any chunk without the term”. A query with no positive term returns []. It ranks with ts_rank_cd(tsv, q, 32), ties by chunk_id. hybrid.rrf(lists, k=60) fuses by rank, score(d) = Σ 1 / (k + rank_r(d)) over the lists that contain d, sorted by score then chunk_id, and returns Ranked(chunk_id, score, rank, ranks) with ranks holding the rank in each input list (None when absent). hybrid.search(conn, embedder, query, n=50, top_k=10) fuses the vector and full-text top-n and returns the first top_k. The tests check cosine scores against the embedder, full-text ranks, that queries with operators and quotes ("exact phrase" -term OR x, a stray ', &|!<->) never raise and still return vector results, that a query without words ("", -) returns [] everywhere, that -term and -"a phrase" exclude every chunk containing them while still returning results, the RRF formula and tie-break, that a chunk first in one list beats a chunk 30th in the other list however large that list’s raw score, that a chunk 30th in both lists (2/90) outranks one first in a single list (1/61), that the lexical golden query q40 is found by full text and missed by vectors, and that fused recall@10 on the golden set is at least that of each side and above one of them.
  4. 04-reranker (Docker): rerank.Reranker is a protocol with rerank(query, chunks, top_k) -> list[Ranked], best first, ranks from 1, ties kept in input order. rerank.LexicalReranker(k1=1.2, b=0.75) scores BM25 over the candidate set it is given: query terms are the distinct lowercase words of re.findall(r"\w+", query.lower()) (chunk text is tokenised the same way); for a term t with document frequency df among the N candidates, idf = ln(1 + (N - df + 0.5) / (df + 0.5)), and a chunk of L words (average avg over the candidates) scores Σ idf · tf · (k1 + 1) / (tf + k1 · (1 - b + b · L / avg)) over the query terms it contains. rerank.CrossEncoderReranker(model="Xenova/ms-marco-MiniLM-L-6-v2", cache_dir=None) wraps fastembed’s TextCrossEncoder and imports fastembed only when it ranks a non-empty list (used in stage 08). rerank.mmr(query_vec, candidates, embeddings, *, lam=0.7, top_k) picks greedily the candidate with the highest lam · relevance - (1 - lam) · max cosine to the chunks already picked (the second term is 0 for the first pick; ties go to the earlier candidate). Relevance is the reranker score scaled into [0, 1]: score / max when no score is negative, min-max otherwise, and the cosine to query_vec when every score is equal. It keeps each candidate’s reranker score and renumbers the ranks. rerank.fetch(conn, chunk_ids) returns {chunk_id: (Chunk, embedding)} from the table, and rerank.pipeline(conn, embedder, reranker, query, *, n=50, top_k=10, mmr_lambda=0.7) takes the fused top-n of hybrid.search, reranks all of it and returns the first top_k of MMR. The tests check BM25 against hand-computed scores, that a rare term outranks a common one, that MMR drops a near-duplicate of the first pick (two chunks sharing two of three paragraphs) where lam=1.0 keeps it, the cosine fallback, that the cross encoder is lazy, that the pipeline stays inside the fused pool with a stable prefix, and that reranking raises MRR on the golden set and MMR keeps it at or above hybrid alone.

  5. 05-metrics (Docker): metrics.relevant(chunks, q) is the set of chunk ids of q.doc_id whose text contains q.anchor (case-insensitive, whitespace-normalised). For a ranked list of chunk ids: recall_at_k = relevant in the top k / relevant total; precision_at_k = relevant in the top k / k; mrr = 1 / rank of the first relevant (0 when none); ndcg_at_k has binary gains and 1 / log2(rank + 1) discounts, divided by the ideal DCG over min(k, relevant total) ranks, so one relevant chunk at rank 1 scores 1.0. metrics.evaluate(conn, embedder, reranker, golden, *, k=10, n=50, mmr_lambda=0.7) -> Metrics runs rerank.pipeline per query (reranker=None means hybrid.search alone), averages over the queries (MRR over the top k) and fills by_kind with the same means plus queries per kind. metrics.report(m) prints one row for all queries and one per kind. python -m retrieval_lab.metrics --corpus fixtures/corpus --k 10 indexes the corpus with HashEmbedder into its own schema (--dsn, default $RETRIEVAL_LAB_DSN or the compose database on port 5433; --schema) and prints the report; --failures prints instead query_id, kind, best_rank for every golden query whose first relevant chunk is outside the top k (best_rank is its rank in the full reranked pool, or none when the pool missed it). The tests check every metric against hand-computed values, the nDCG ideal, the recall@50 of the fused pool, the reference numbers below within 0.02 and the floors, the split by kind, and both CLI outputs.

  6. 06-answer-assembly (Docker): answer.sources(conn, embedder, reranker, query, *, top_k=5) returns the pipeline’s top chunks and their reranker scores. answer.assemble(query, chunks, *, scores=None, min_support=1, floor=SUPPORT_FLOOR) -> Draft refuses (Draft("", (), refused=True)) when fewer than min_support chunks score at least floor (with scores=None every chunk counts); otherwise its prompt tells the model to answer only from the numbered sources, to cite every sentence as [n] and to reply exactly I cannot answer from the sources. when none supports an answer, then gives the question and the sources as [i] <chunk_id> followed by the chunk text, and citations holds the chunk ids in that order. answer.check(text, citations) -> Faithfulness(cited, uncited_claims, refused) splits the text into sentences (after ., !, ? followed by whitespace, or a line break, but never after the abbreviations Mr., Mrs., Ms., Mme., Dr., St., vs., e.g., i.e., etc.; a marker that opens a sentence belongs to the one before), cited are the valid source numbers used, uncited_claims counts the sentences without a valid marker (the refusal sentence excepted), and refused is whether the text contains the refusal. A number outside 1..n is not a citation: its sentence counts as uncited and answer.invalid_citations(text, citations) lists it. answer.run(query, draft, *, query_fn=None, model=None) -> Answer(text, faithfulness, cost_usd) returns the refusal without a model call for a refused draft and otherwise sends the prompt through claude_agent_sdk.query with tools=[], max_turns=1, setting_sources=[], max_budget_usd=0.05, strict_mcp_config=True (your own MCP servers stay out of the call) and model claude-opus-5-5. The tests use the scripted SDK fake: the prompt, refusal below the floor and with min_support=2, no model call on refusal, [1] and [3] cited, uncited sentences, abbreviations that do not split, [7] among five sources, the pinned options, that six off-corpus questions refuse on the real index and that 47 of the 48 golden queries clear the floor.

  7. 07-retrieval-error-analysis (Docker, judgment): write RETRIEVAL_ANALYSIS.md in your workspace root. Run python -m retrieval_lab.metrics --corpus fixtures/corpus --k 10 --failures to list the golden queries the pipeline misses, then read each one: the query, the chunk that holds its anchor, the top 10, and where the anchor chunk sits in the vector list, the full-text list, the fused pool and the reranked pool. Give each miss one row of a table | query_id | failure | fix |, with failure one of vocabulary-gap (query and passage share no words, neither list reaches the chunk), chunk-boundary (the cut separated the anchor from the words or heading that would match), ranking (the chunk is in the fused pool but fusion ranks it low), reranker-demotion (fusion had it in the top 10 and the reranker pushed it out) or anchor-ambiguity (another chunk in the top 10 answers the question, but without the anchor), and fix one of hybrid-weights, smaller-chunks, contextual-chunk-headers, query-rewrite, better-embedder. Then write one paragraph (25 words or more) per failure kind you used, naming the kind or sitting under a heading that names it. The test computes the misses itself from your pipeline (HashEmbedder, LexicalReranker, MMR, k = 10) and checks that every miss has a row, that both columns use the vocabulary, that at least 70% of the failure kinds agree with tests/KEY.yaml (some queries accept two readings), and that the paragraphs exist. A table with the same kind on every row, or rows for other query ids, does not pass; four disagreements out of thirteen fail.

  8. 08-live-local-models (live, Docker and the embedding models): write FastEmbedder.embed (fastembed’s TextEmbedding(model_name=..., cache_dir=...), imported on the first non-empty call, float32 rows of 384) and use CrossEncoderReranker from stage 04. The offline tests check the option shapes and the cache directory, that nothing imports fastembed before it is needed, and that both classes wrap a fake fastembed correctly. The live test indexes the corpus with FastEmbedder into a fresh 384-dimension table, evaluates hybrid alone and the full pipeline (cross-encoder rerank of the fused top 50, then MMR), writes the aggregate to runs/live-<date>.json in your workspace, and asserts recall@10 ≥ 0.75 and MRR ≥ 0.60 (the course’s original target was recall@10 ≥ 0.85; see the measured numbers below). It runs for about 2 minutes on a laptop CPU (indexing 653 chunks and cross-encoding 48 × 50 pairs).

  9. 09-live-answer (live, Docker and Claude Code): answer.run without a fake. Three golden queries (one per kind) go through answer.sources, answer.assemble and claude-opus-5-5; the test asserts that each answer either refuses or cites a chunk that holds the anchor, that no answer cites a source number it was not given, and that the three cost at most $0.20 together, and writes runs/live-<date>.json. The offline tests replay the same path with the scripted SDK.

Measured numbers (reference, HashEmbedder + full text + LexicalReranker + MMR)

On the 48 golden queries the fused top-50 pool reaches recall@50 = 0.792, which bounds every reranker over that pool. At k = 10:

pipeline recall@10 precision@10 MRR nDCG@10
hybrid (RRF) 0.656 0.069 0.406 0.468
+ lexical rerank 0.750 0.079 0.545 0.595
+ MMR (λ 0.7) 0.729 0.077 0.536 0.583

By kind with rerank and MMR: lexical recall 1.000 / MRR 0.854, factual 0.938 / 0.656, conceptual 0.250 / 0.098. Every miss but one is a conceptual query: they paraphrase without the document’s vocabulary, so neither hashing vectors nor BM25 can bridge them (stage 07 classifies them; stage 08 measures real embeddings). The course’s original targets were recall@10 ≥ 0.80, MRR ≥ 0.50 and nDCG@10 ≥ 0.55; recall@10 cannot pass 0.792 over a 50-candidate pool, so its floor is the measured value rounded down to 0.05, 0.70 (MRR and nDCG keep their original floors). Stage 05 pins the numbers above within 0.02.

The support floor of stage 06, answer.SUPPORT_FLOOR = 4.0, is a BM25 score from the lexical reranker: off-corpus questions top out between 0 and 3.2 (their stopwords still match), the weakest golden query tops out at 3.86 (q41, which therefore refuses) and the median golden query at about 11. A different reranker (the cross-encoder scores logits) needs its own floor.

Measured live numbers (2026-10-23)

fixtures/live_local_models_run.json (stage 08, BAAI/bge-small-en-v1.5 + Xenova/ms-marco- MiniLM-L-6-v2 + MMR, 48 queries, k = 10):

pipeline recall@10 precision@10 MRR nDCG@10
hybrid (RRF), bge-small 0.740 0.077 0.552 0.598
+ cross-encoder + MMR 0.781 0.081 0.648 0.674

By kind with the cross-encoder: lexical recall 0.906 / MRR 0.896, factual 1.000 / 0.805, conceptual 0.438 / 0.243 (the hashing pipeline had 0.250 / 0.098). MRR clears its floor of 0.60; recall@10 does not reach the course’s original target of 0.85: the fused bge pool holds the anchor chunk for 0.875 of the queries at 50, and the cross-encoder lifts only part of it into the top 10, mostly failing on conceptual queries. The stage 08 floor is therefore the measured recall rounded down to 0.05, 0.75, with MRR ≥ 0.60 kept. The download was 66.5 MB (quantised bge-small) and 91 MB (MiniLM cross-encoder), 18 s; the stage ran for 132 s (index 50 s, reranked evaluation 80 s).

fixtures/live_answer_run.json (stage 09): q03 (factual), q40 (lexical) and q46 (conceptual) each answered with citations of a relevant chunk, none refused, no invalid citation, $0.099 in total. The recorded q40 answer counts one uncited “claim”: the splitter breaks after Sofronie., the period of the name inside the quoted sign.

Postgres with pgvector (stages 02 to 09)

Stages 02 to 09 need Docker. The tests start their own throwaway container (pgvector/pgvector:pg17: Postgres 17 with pgvector 0.8) on a free port of 127.0.0.1, give each test module its own schema, and stop the container when the run ends; you do not start anything for them. The first run pulls the image (about 650 MB on disk). The containers carry the label miubais=m4-tests; the next test run stops any left over for longer than an hour (a grader timeout kills pytest before it can stop its own). To stop them by hand:

docker ps -q --filter label=miubais=m4-tests | xargs docker stop

For your own experiments, docker-compose.yml runs the same image on 127.0.0.1:5433 (5432 is often taken by a local Postgres). From the repository root:

docker compose -f course/m4-retrieval-lab/docker-compose.yml up -d
psql postgresql://postgres:miubais@127.0.0.1:5433/postgres

In Python, connect with psycopg.connect(...), run CREATE EXTENSION IF NOT EXISTS vector once, then pgvector.psycopg.register_vector(conn) so numpy arrays go in and out as vector values. The password is a fixed local placeholder. Stop it with docker compose -f course/m4-retrieval-lab/docker-compose.yml down; add -v to delete its volume.

Model cache (stage 08)

FastEmbedder and the cross-encoder download their ONNX models on first use (66.5 MB and 91 MB, about 160 MB) into cache_dir, which defaults to $FASTEMBED_CACHE_PATH when set and otherwise to ~/.cache/fastembed (fastembed’s own default is a temporary directory that the system may clear). uv run miubais doctor reports embed-model as present once both models are in that directory. No test before stage 08 downloads a model.

Requirements by stage

StageTitleRead firstLiveRequires
01-chunkerChunk Markdown and plain text by headings and paragraphsChunking and embeddingsno-
02-embed-and-storeA hashing embedder and a pgvector table with full textChunking and embeddingsnodocker
03-hybrid-retrievalVector and full-text search fused with reciprocal rank fusionHybrid searchnodocker
04-rerankerA lexical reranker and MMR over the fused candidatesReranking and MMRnodocker
05-metricsRecall, precision, MRR and nDCG on the golden setRetrieval metrics and golden setsnodocker
06-answer-assemblyA cited answer prompt and a faithfulness checkCitations, faithfulness and agentic retrievalnodocker
07-retrieval-error-analysisRetrieval error analysis of the failed golden queries (judgment stage)Retrieval metrics and golden setsnodocker
08-live-local-modelsLive: local embeddings and a cross-encoder on the golden setChunking and embeddingsyesdocker, embed-model
09-live-answerLive: cited answers to three golden queriesCitations, faithfulness and agentic retrievalyesdocker, claude-auth

Check your work

Check your work from the clone:

uv run miubais check m4-retrieval-lab
uv run miubais hint m4-retrieval-lab

Opened from pnpm dev on localhost with uv run miubais serve running, this panel runs them for you.