~learn.
Code lab · World 11 · Designing AI Systems

RAG chunker and context packer.

Split a policy section on its strongest boundaries into chunks under a token limit, then pack reranked chunks into a token budget, best first, without duplicates or weak evidence.

Medium~40 min3 example + 9 hidden tests+90 XPJump to the code ↓
Saving also stores the Python runtime (~13.5 MB, once for every lab) so this lab runs with no internet.

MegaMart Assist answers from its policy pages. Two pieces of Day 3 decide what the model actually reads:

  • Chunking (section 2.2): a chunk is what you embed, retrieve and paste. Split on the strongest boundary first, paragraphs, then lines, then words, keeping pieces together while they fit (LangChain's recursive splitter).
  • Packing (section 6.3): the reranker scored the candidates; fit the best into the retrieved-context budget (≤ 3,200 tokens in section 6.2), best first, and stop padding with weak evidence.

Tokens here are whitespace-separated words: count_tokens is given, so the tests are exact.

split_section(doc_id, title, text, max_tokens) returns a list of Chunks:

  • If the whole text fits in max_tokens, it's one chunk.
  • Otherwise split on the first separator in SEPARATORS ("\n\n", then "\n", then " ") that occurs in the text. Drop empty pieces (after strip()). Walk the pieces in order, joining consecutive pieces with that separator while the result fits; start a new chunk when the next piece won't fit.
  • A single piece that is itself over max_tokens is split recursively with the next separator, and its chunks take their place in order.
  • Every chunk: id = f"{doc_id}:{n}" numbered from 0 in order, doc_id, title, the chunk's text (stripped), tokens = count_tokens(text), content_hash = hash_text(text) and score = 0.0.

assemble(chunks, budget, min_score=0.35) returns None or (context, ids):

  • Go through the chunks by score, highest first (ties: id alphabetical).
  • Stop at the first chunk below min_score: weak evidence doesn't get in, however much room is left.
  • Skip a chunk whose content_hash was already picked (the same text from another doc).
  • Skip a chunk that would take the total over budget, and keep going: a smaller one may still fit.
  • Nothing picked: None (the caller abstains or hands off). Otherwise context = "\n".join(render(c) for c in picked) and ids in the same order.
Python 3.14 loads on first run (~6.5 MB once)14
Loading the editor…
not run yet · 8 s limit
press Run tests

Write your code, then run 3 example tests and 9 hidden tests. Each run is limited to 8 seconds: an infinite loop gets stopped and Python restarts, your code stays.