Lab 10 · Contextual Enrichment Upgrade

Anthropic's contextual-retrieval work reported that contextual embeddings alone cut top-20 retrieval failures by roughly 35%; in this lab you will reproduce the mechanism on a corpus engineered to show it starkly — answer-bearing chunk retrieval jumps from 15% to 60% in the offline configuration — and you will build the cost model that tells you what enrichment costs per 1,000 chunks, with and without prompt caching. The skill being trained is the data engineer's version of the technique: enrichment is an indexing-time pipeline stage with a per-chunk unit cost, not a query-time trick.

Time & tier

~90 minutes hands-on · Core tier (maps to module M6).

Setup

This lab builds directly on Lab 09. Work in a new folder but reuse the corpus generator:

mkdir -p ~/labs/de-lab-10 && cd ~/labs/de-lab-10
python3 -m venv .venv && source .venv/bin/activate
pip install sentence-transformers chromadb rank_bm25 numpy anthropic
cp ~/labs/de-lab-09/make_corpus.py .
python make_corpus.py   # writes corpus/docs.jsonl + corpus/golden.jsonl (120 docs, 20 queries)

If you skipped Lab 09, copy make_corpus.py from its Setup section — it is self-contained and deterministic.

The lab has two modes controlled by a --dry-run flag. In dry-run mode (the default, $0, fully offline) the "LLM-written context" is produced by a deterministic template so you can build and measure the whole pipeline without a key. With an ANTHROPIC_API_KEY set you can rerun the enrichment stage with real model-written context, which is richer than the template; the corpus is small enough that the run costs well under a dollar.

Guided steps

1. 1.

   Chunk the corpus the way a naive pipeline would. Create enrich.py and start with a sentence-group splitter that cuts each article into four chunks: the shared safety boilerplate, the applies-to/symptom sentences, the diagnosis/resolution sentences, and the cross-references/escalation tail.

   Copy

   import json, pathlib, re, sys
   import chromadb
   from rank_bm25 import BM25Okapi
   from sentence_transformers import SentenceTransformer
   
   docs = [json.loads(l) for l in pathlib.Path("corpus/docs.jsonl").read_text().splitlines()]
   golden = [json.loads(l) for l in pathlib.Path("corpus/golden.jsonl").read_text().splitlines()]
   
   def chunk_doc(d):
       sents = re.split(r"(?<=\.) ", d["text"])
       return [
           " ".join(sents[0:3]),   # safety boilerplate (identical across docs)
           " ".join(sents[3:5]),   # applies-to + symptom
           " ".join(sents[5:8]),   # diagnosis + resolution + acceptance test
           " ".join(sents[8:]),    # cross-references + escalation
       ]
   

   Print the four chunks of kb-000 and stare at chunk 2 (index 2): it contains the fix — the thing every golden query is really asking for — but never mentions the machine type or the part code. That is the retrieval hole you are about to measure.

2. 2.

   Write the context generator, real and dry-run. The real version asks the model to write one sentence situating the chunk within its document; the dry-run version derives an equivalent line from the document title deterministically. Note the prompt-caching structure in the real version: the full document goes in a cached system block, so the four chunk-calls per document pay for the document tokens once.

   Copy

   def make_context_dry(d, chunk):
       """Deterministic stand-in for the LLM-written chunk context."""
       return (f"Context: from the service bulletin for the "
               f"{d['title'].split(': ')[1].split(' (part')[0]} on part {d['part_code']}. ")
   
   def make_context_llm(d, chunk):
       import anthropic
       client = anthropic.Anthropic()
       msg = client.messages.create(
           model="claude-opus-4-8",
           max_tokens=100,
           thinking={"type": "adaptive"},
           system=[{"type": "text",
                    "text": "You situate document chunks for a retrieval index. "
                            "Reply with one sentence of context and nothing else.\n\n"
                            "<document>\n" + d["title"] + "\n" + d["text"] + "\n</document>",
                    "cache_control": {"type": "ephemeral"}}],
           messages=[{"role": "user", "content":
                      "Write one short sentence situating this chunk within the document, "
                      "naming the equipment and part number if relevant.\n\n<chunk>\n"
                      + chunk + "\n</chunk>"}],
       )
       return "".join(b.text for b in msg.content if b.type == "text").strip() + " "
   

3. 3.

   Build both indexes and a hybrid search over chunks. Reuse Lab 09's hybrid pattern (dense + BM25 + RRF), but over 480 chunks instead of 120 documents:

   Copy

   DRY = "--llm" not in sys.argv
   model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
   make_context = make_context_dry if DRY else make_context_llm
   
   def build_index(enriched: bool):
       ids, texts, parents = [], [], {}
       for d in docs:
           for ci, chunk in enumerate(chunk_doc(d)):
               cid = f"{d['id']}::c{ci}"
               text = (make_context(d, chunk) + chunk) if enriched else chunk
               ids.append(cid); texts.append(text); parents[cid] = d["id"]
       emb = model.encode(texts, normalize_embeddings=True, show_progress_bar=False)
       client = chromadb.Client()
       name = "enriched" if enriched else "naive"
       try: client.delete_collection(name)
       except Exception: pass
       col = client.create_collection(name, metadata={"hnsw:space": "cosine"})
       col.add(ids=ids, embeddings=emb.tolist(), documents=texts)
       tokenized = [t.lower().replace("(", " ").replace(")", " ").replace(",", " ")
                     .replace(".", " ").split() for t in texts]
       return col, BM25Okapi(tokenized), ids, parents
   
   def rrf(a, b, k=60, top=20):
       s = {}
       for r, doc in enumerate(a): s[doc] = s.get(doc, 0) + 1 / (k + r + 1)
       for r, doc in enumerate(b): s[doc] = s.get(doc, 0) + 1 / (k + r + 1)
       return [d for d, _ in sorted(s.items(), key=lambda x: -x[1])][:top]
   

4. 4.

   Evaluate with an honest metric. Parent-document recall is too forgiving — with five fused chunks mapping to up to five parents, both indexes score 100% and teach you nothing (try it if you doubt this). The metric that matches what a RAG system needs is answer-bearing chunk recall: is the ::c2 chunk (diagnosis + resolution) of a relevant document in the top 5?

   Copy

   def evaluate(col, bm25, ids, parents, label):
       hits = rr_sum = 0
       for q in golden:
           qv = model.encode([q["query"]], normalize_embeddings=True)
           dense = col.query(query_embeddings=qv.tolist(), n_results=20)["ids"][0]
           scores = bm25.get_scores(q["query"].lower().split())
           order = sorted(range(len(scores)), key=lambda i: -scores[i])[:20]
           sparse = [ids[i] for i in order]
           fused = rrf(dense, sparse, top=5)
           targets = {f"{r}::c2" for r in q["relevant_ids"]}
           rank = next((i for i, c in enumerate(fused) if c in targets), None)
           if rank is not None:
               hits += 1; rr_sum += 1 / (rank + 1)
       print(f"{label:10s} answer-chunk recall@5 = {hits}/20 ({100*hits/20:.0f}%)  MRR={rr_sum/20:.2f}")
       return hits
   
   n1 = evaluate(*build_index(False), "naive")
   n2 = evaluate(*build_index(True), "enriched")
   print(f"retrieval failure rate: {100*(1-n1/20):.0f}% -> {100*(1-n2/20):.0f}%")
   

5. 5.

   Run and record. python enrich.py (dry-run). Expected: naive 3/20 (15%), enriched 12/20 (60%) — failure rate 85% → 40%, MRR roughly 0.06 → 0.52. Even a one-line templated context more than triples answer-chunk retrieval, because the resolution chunk now carries the equipment name and part code that queries actually mention. With --llm and an API key, the model writes richer context (symptom paraphrases, section role) and typically closes more of the remaining gap; record both rows if you run it.

6. 6.

   Build the cost model. Create cost_model.md with this arithmetic, using current list prices for the model you used (for claude-opus-4-8: $5 per million input tokens, $25 per million output tokens; cache writes ~1.25x input, cache reads ~0.1x input). Per document: ~260 input tokens of document context, 4 chunk calls of ~80 tokens each, ~30 output tokens per call. Without caching, the document tokens are re-sent 4 times: 4 × (260 + 80) = 1,360 input + 120 output tokens per document. With prompt caching, the document is written to cache once (260 × 1.25 = 325 token-equivalents) and read 3 times (3 × 26), plus 4 × 80 chunk tokens: ~723 input-equivalents + 120 output. Scale both to 1,000 chunks (250 documents) and to a realistic 1M-chunk corpus. You should land near $3.20 per 1K chunks uncached vs ~$1.85 per 1K chunks cached at Opus-tier prices — and note that a batch-processing tier (50% discount, since enrichment is never latency-sensitive) roughly halves both numbers again. Add one line stating when enrichment is not worth it: chunks that already carry their context (titles repeated per chunk, self-contained FAQ entries) buy little for the spend — measure first, as you just did.

7. 7.

   Stack it with Lab 09's reranker (optional but recommended). Contextual enrichment and reranking compose — that is how Anthropic's reported numbers reach a ~67% failure-rate reduction when BM25 and reranking join contextual embeddings. Import the cross-encoder/ms-marco-MiniLM-L-6-v2 reranker from Lab 09, rerank the enriched index's top-20 fused candidates down to 5, and re-run the answer-chunk evaluation. Add the row to your report: on this corpus the reranker typically recovers several of the remaining misses because the enriched chunk text now contains the equipment terms the cross-encoder needs to score relevance. Record the added latency too — the composition argument only holds when each layer's cost is on the table next to its gain.

8. 8.

   State the production shape. Close your notes with the pipeline placement: enrichment runs once per chunk at indexing time, inside the same batch job that chunks and embeds; it re-runs only for changed documents (Lab 11 wires up exactly that change detection); and its cost scales with corpus size, not query volume — which is why a technique that would be absurd per-query is routine per-index. One paragraph on failure handling completes the picture: the enrichment call is the only fallible stage, so wrap it with a retry and a fallback to the templated context — a chunk indexed with template context is degraded, a chunk missing from the index is invisible.

Verification

• •
  Dry-run output shows naive answer-chunk recall@5 of 15% and enriched of 60% (±1 query).
• •
  The enriched collection contains 480 chunks whose text begins with Context: from the service bulletin for….
• •
  Parent-document recall (if you measure it) is ~100% for both indexes — confirming you needed the stricter metric to see the effect.
• •
  cost_model.md shows per-1K-chunk costs for uncached, cached, and cached+batch rows, with the arithmetic visible.

Completion artifact

A de-lab-10/ folder containing enrich.py and before-after-report.md: the naive vs enriched recall/MRR table (dry-run row required, --llm row optional), a paragraph interpreting why the resolution chunk drives the gap, and the $ / 1K chunks cost model with caching and batch variants. Self-grading rubric (pass = 8/10):

• Measured uplift (0–3): both indexes evaluated with the answer-chunk metric; enriched beats naive by ≥30 points = 3.
• Metric honesty (0–2): report explains why parent-doc recall was rejected as the metric = 2.
• Cost model (0–3): per-1K-chunk figures for uncached, cached, and batch, with visible arithmetic and a stated price basis = 3.
• Production placement (0–1): enrichment correctly placed at indexing time with incremental re-run noted = 1.
• API-mode readiness (0–1):make_context_llm uses cached system blocks and runs (or is dry-run-validated with the flag) = 1.

Stuck?

• •
  Both indexes score 100%: you are measuring parent-document recall — switch to the ::c2 target set from step 4.
• •
  Naive scores higher than 15%: check that chunks exclude the document title; including titles leaks context into every chunk and hides the effect.
• •
  chromadb says the collection already exists: the delete_collection guard in build_index handles reruns — make sure you kept it.
• •
  --llm mode errors with 401: export ANTHROPIC_API_KEY; dry-run mode needs no key and produces the required artifact rows.
• •
  Cache reads showing zero in --llm mode: the system block must be byte-identical across the four chunk calls of a document — build it once per document, not per call.

Client-safe variant

Dry-run mode is already client-safe: chunking, templated context, embeddings, and evaluation all run locally with no external calls. When client policy allows an approved LLM endpoint (their cloud tenancy or an internal gateway), point make_context_llm at it and keep everything else unchanged — the enrichment prompt sends the client's document to the model, so it must ride the approved channel, and the cost model you built becomes the budget line you show the client before enriching a million chunks of their data. Where no LLM channel exists at all, the templated-context fallback you just measured is the honest recommendation: it captured most of this corpus's headroom for zero spend.