Assemble context

Lesson 10
45 min
Module 4: Construct grounded answers

Lesson 10 of 18 · RAG: Building AI That Uses Your DataAssemble context

Course overview

1. Define the knowledge problem

2. Ingest documents reproducibly

3. Build and compare retrieval

4. Construct grounded answers

5. Evaluate and debug the pipeline

6. Maintain and protect the knowledge base

7. Case studies

What you will be able to do

Decide which passages go into the prompt with the question. Put the best ones first, drop copies, stay within a fixed number of tokens, and keep each passage’s ID and source. Then every answer can cite a passage, and every question costs about the same.

Before you startWhat to know and to have readyAbout 45 minutes

What you need to know

Reranked search with the date and access filters, from Reranking and query changes and Lexical, dense and hybrid search.
Chunk IDs, versions and effective dates, from Preserve metadata.
What a token is, from From text to tokens in LLMs: How They Work and How to Choose One.

What you need on your computer

~/projects/policy-assistant after Module 3, with .venv active (source .venv/bin/activate, or .venv\Scripts\Activate.ps1 in PowerShell). python -m pytest shows 48 passed.
The index index.sqlite, made by python -m policy_assistant ingest in Module 3.
If your project is missing or broken, copy the snapshot end-of-m03 with the steps of Use a snapshot, then run ingest.

No account and no key. Tested with Python 3.12.13 on macOS (Apple silicon). The commands are the same in Windows PowerShell; the Windows steps are not tested.

Models change fast. The numbers in this module are a snapshot. The method is what lasts.

One page of notes for the exam

Some exams allow one page of notes. You cannot copy the whole textbook onto it, so you choose. You write the notes that answer the likely questions, and you do not write the same note twice. You put the most important notes at the top, where you find them fast. And next to each note you write the page of the book, so you can show where it came from.

The assistant has the same task before every answer. Search found five passages. Now someone must decide what goes on the “page” that the model reads.

The problem: why not send every document?

Tomás has a shortcut. “The collection is only 36 documents. Put all of them in the prompt, every time. Then the model cannot miss anything, and we do not need search.”

Amira does not argue. She tests it. She sent 20 of Grace’s questions to both chat models twice. The first time, each question came with the 5 passages from search. The second time, it came with all 185 passages of the collection (recorded on Azure). Here is the honest result, by the course’s automatic check:

Model	Five passages from search	The whole collection
chat-small (gpt-6-luna)	16 of 20 right	19 of 20 right
chat-strong (gpt-6.1-sol)	15 of 20 right	18 of 20 right

On this small collection, the whole collection answered as well or better. It found the answer to Q13, “Does Larkfield sell extended warranties?”, which is only in a French document that search missed. Tomás is right today. So what is wrong with his shortcut? Look at what each request contains.

The median request with five passages from search had 905 input tokens. With all 185 passages, it had about 17,850: 20 times more, for every question. Recorded on Azure.

Text description of the chart

A column chart with two columns. Five passages from search: a median of 905 input tokens per question (rerank, all 67 questions). All 185 passages of the collection: a median of about 17,850 input tokens per question (from 17,840 to 17,870, by the length of the question; 20 questions). The second column is about 20 times the first. Both counts are the provider’s own token counts, and they include the instructions and the question.

Twenty times the tokens means about twenty times the price of the input. At Larkfield’s 45 questions a day, this is the cost for each model (US dollars):

Model	Search, per day	Whole collection, per day
chat-small	0.008	0.09
chat-strong	0.14	1.70

These are small sums for 36 documents. But four things change as the assistant grows:

Scale. 36 short documents already need 17,850 tokens. Every new policy, guide and version makes every question bigger and slower.
Access. The whole collection has the 6 staff-only procedures in it. With all of them in the prompt, chat-small answered Q01, a double charge, from the staff procedure: “After confirming both charges, Finance refunds the extra charge within 3 business days.” A customer-facing assistant must never send staff text to the model for a customer.
Versions. Both returns policies, 30 and 45 days, are in every prompt. Only the model’s reading decides which one applies.
Citations. The model must copy one ID out of 185. In this run it produced the course’s only made-up citation: c23617e24d2, an ID that does not exist. You meet it again in the next lesson.

So the shortcut works today and fails tomorrow. A long context window does not fix it: a model can accept far more text than it reads well. Context and memory in the LLMs course measured this on a product handbook.

The idea: a context budget

The context is the text that goes into the prompt with the question: here, the passages. The context budget is the most tokens that the passages may use. Assembling the context means four rules:

Keep the search order. The best passage comes first. If something must go, it is the weakest.
Remove near-duplicates. Two passages from different documents that say almost the same words use the budget twice for one fact.
Stop at the budget. A passage that does not fit stays out.
Keep the ID and the source. Each passage enters the prompt with a header: its ID, title, type, version, dates and section.

The project estimates tokens as the number of characters divided by 4. That is rough, but good enough for a budget; the provider reports the real count afterwards. Larkfield’s budget is 1,500 estimated tokens of passages.

How close is “almost the same words”? The project compares the sets of words of the two passages, without the title line. The share of words they have in common is the Jaccard similarity: the shared words divided by all the different words of both. Two passages with a Jaccard similarity of 0.9 or more are near-duplicates.

A tiny example: five passages for one question

Grace’s question Q30: “How long is the warranty on the pump of a pressure washer?” Reranked search, top 5, as of 2026-10-09, gives these passages:

Rank	Passage ID	Document · section	Estimated tokens
1	`ff8235f48619`	warranty-claims-procedure · Pressure washer pumps (staff)	104
2	`d9a5bd8ea531`	warranty-policy · Warranty periods by product category	314
3	`50f5713033a7`	product-care-pw2400 · Winter storage	92
4	`08705a17f766`	product-care-pw2200 · Winter storage	92
5	`34a0b9b8055e`	product-care-pw2400 · Using it safely	107

Passages 3 and 4 are the Winter storage sections of the PW-2400 and PW-2200 guides. Without their titles, both have the same 27 words. So 27 words are shared, out of 27 different words, and the Jaccard similarity is 27 ÷ 27 = 1.0. Now follow the rules, one passage at a time:

Assembling the context for Q30 with the project's rules. One near-duplicate is dropped; the four other passages fit in the budget.

Text description of the diagram

On the left, the top 5 of reranked search for Q30, with their estimated tokens: 1 ff8235f48619, 104; 2 d9a5bd8ea531, 314; 3 50f5713033a7, 92; 4 08705a17f766, 92; 5 34a0b9b8055e, 107. They go one by one, best first, to the first question: near-duplicate of a passage already kept? For 08705a17f766 the answer is yes: the PW-2200 guide says word for word what the PW-2400 guide says, so it is dropped. The others go to the second question: do the tokens used so far plus this passage fit in 1,500? All four fit, so each one is kept, with its header (ID, title, type, version, dates, section). The result: a context of 4 passages and 617 estimated tokens, best first. No passage is dropped for the budget in this example.

Passage 1 uses 104 tokens, passage 2 brings the total to 418, passage 3 to 510. Passage 4 is a near-duplicate of passage 3, so it is dropped. Passage 5 brings the total to 617. The budget of 1,500 is not reached.

The rank decides which copy stays. The PW-2400 guide ranked 3rd, so it stays, even though both guides say the same.

Explore: a smaller budget

The budget of 1,500 did not matter for Q30. Now make it smaller. Predict first, then move the slider.

Explore

Which passages fit in the budget?

Budget: 1,500 tokens
Used: 617 tokens
Passages in the prompt: 4 of 5
Passages with the answer: 2 of 2

Dark: tokens used by the passages in the prompt. Light: the whole budget. The scale ends at 1,500 tokens.

Budget 1,500 tokens: 4 of 5 passages in the prompt, 617 tokens used.

Rank 1: ff8235f48619 warranty-claims-procedure · Pressure washer pumps · staff holds the answer

Pressure washer pumps - If the pump of a PW-2200 fails within 3 years, send the pump kit PK-2200-07 instead of a new p…

104tokens

In the prompttokens 1–104
Rank 2: d9a5bd8ea531 warranty-policy · Warranty periods by product category holds the answer

Warranty periods by product category Product category | Examples | Warranty period | Exceptions Hand garden tools | sp…

314tokens

In the prompttokens 105–418
Rank 3: 50f5713033a7 product-care-pw2400 · Winter storage

Winter storage Drain all water from the pump and hoses before temperatures fall below 0 °C. Store it indoors, above 0…

92tokens

In the prompttokens 419–510
Rank 4: 08705a17f766 product-care-pw2200 · Winter storage

Winter storage Drain all water from the pump and hoses before temperatures fall below 0 °C. Store it indoors, above 0…

92tokens

Droppednear-duplicate of 50f5713033a7
Rank 5: 34a0b9b8055e product-care-pw2400 · Using it safely

Using it safely Never run the pressure washer for more than 2 minutes with the trigger released: the pump overheats. …

107tokens

In the prompttokens 511–617

The real top 5 of Q30 (reranked search), and what the project's context.assemble() keeps at each budget. Every state was computed by the project's own code.

Text description of the explorer

The explorer lists the five passages of Q30 in search order, with their estimated tokens: 1 ff8235f48619, the warranty claims procedure (staff), 104, holds the answer; 2 d9a5bd8ea531, the warranty table, 314, holds the answer; 3 50f5713033a7, PW-2400 winter storage, 92; 4 08705a17f766, PW-2200 winter storage, 92, a near-duplicate of 3; 5 34a0b9b8055e, PW-2400 safe use, 107. One slider sets the budget, from 100 to 1,000 tokens in steps of 50, then 1,500.

From 650 to 1,500, 4 passages are in the prompt (617 tokens); only the near-duplicate is dropped. At 550 and 600, the safe-use passage is also over budget: 3 passages, 510 tokens. At 450 and 500, only passages 1 and 2 are in (418 tokens). At 350 and 400, passages 1, 3 and 5 are in (303 tokens): the warranty table no longer fits, and smaller passages take its place. From 200 to 300, passages 1 and 3 are in (196 tokens). At 150, only passage 1 is in. At 100, only passage 3 fits, so no passage with the answer is left.

At 300 tokens, the warranty table is out. It is the biggest passage: after the procedure, it does not fit. The rule “skip what does not fit, try the next” then fills the space with the winter-storage passage, which does not help at all. For an agent, the staff procedure still says “within 3 years”. For a customer-facing assistant, the access filter removes that procedure, and no passage with the answer is left.

So a budget is a limit, not a goal. It protects cost and time. Set it so that your largest useful passage, here a table, always fits.

The header: the ID and the source travel with the text

Each passage enters the prompt with a header line. This is passage 1 of Q30, exactly as the model sees it:

[ff8235f48619] Warranty claims procedure | procedure, version 3, in force from 2026-01-01, no end date | section: Pressure washer pumps
Warranty claims procedure
Pressure washer pumps

- If the pump of a PW-2200 fails within 3 years, send the pump kit PK-2200-07 instead of a new pressure washer.
- For the PW-2400, send the pump kit PK-2400-03.
- Offer a technician visit only if the customer cannot fit the kit.

Each part has a job:

The ID in square brackets is what the model cites. The next lesson checks every citation against it.
The type (policy, procedure, help article, supplier content) matters when two documents disagree. Larkfield’s rule 2 says that a policy wins over a help article.
The version and the dates let the model choose the version in force, and let a person check that choice.
The section tells the agent where to look in the document.

Without the header, the model sees anonymous text. It cannot cite it, and nobody can check the answer.

Two versions are never merged

The near-duplicate rule has one exception, and it is the most important line of the code. Look at the Return window sections of the returns policy, version 3 and version 4. Without their titles, they share 24 of their 26 different words. The only two words that differ are “30” and “45”. Their Jaccard similarity is 24 ÷ 26 = 0.923, above 0.9.

So a blind near-duplicate rule would delete one of them, and the deleted word would be the whole answer. The project never merges two passages of the same document. Their differences are the point, and the date filter of Module 3 already decides which version is in force.

Larkfield’s budget in numbers

Amira ran assemble() on the reranked top 5 of all 67 questions. The passages used a median of 501 estimated tokens, and at most 871. So the budget of 1,500 never cut a passage at top 5. Only Q30 lost a near-duplicate. With top 10, the median was 1,071 tokens, still under the budget, and six questions lost near-duplicates from the PW guides.

The budget is a guard. It matters when somebody raises the top k, makes the chunks longer, or adds many documents. Then cost and time stay the same per question, and the weakest passages go first.

Build it: the context in the project

The project gets one new file, policy_assistant/context.py, and its tests. Work in ~/projects/policy-assistant, with .venv active.

Check the starting point:
Terminal window
```
python -m pytest
```
The output ends with 48 passed.

Make the file policy_assistant/context.py:

"""Assemble the context: choose the passages that go into the prompt, within a token budget.

1. Keep the search order (the best passage first).
2. Remove near-duplicates: two passages whose words are almost the same (Jaccard similarity of their
   word sets >= 0.9, title line left out), for example the same paragraph in the PW-2200 and PW-2400
   guides. The first (better-ranked) one stays. Two versions of the same document are NOT merged:
   their differences (30 days or 45 days) are the point, and the date filter decides between them.
3. Stop before the budget is used up. Tokens are estimated as characters / 4 (good enough for a
   budget; the provider reports the real count).
Every passage keeps its chunk ID and its source, so that the answer can cite it.
"""

import math
import re
from dataclasses import dataclass, field

BUDGET = 1500       # estimated tokens of passages per question
SIMILAR = 0.9


def estimate_tokens(text: str) -> int:
    return math.ceil(len(text) / 4)


def word_set(text: str) -> set[str]:
    body = text.split("\n", 1)[1] if "\n" in text else text
    return set(re.findall(r"\w+", body.lower()))


def jaccard(a: str, b: str) -> float:
    wa, wb = word_set(a), word_set(b)
    return len(wa & wb) / len(wa | wb) if wa | wb else 1.0


def format_passage(chunk) -> str:
    dates = f"in force from {chunk.effective_from}" + (f" to {chunk.effective_to}" if chunk.effective_to else ", no end date")
    return (f"[{chunk.chunk_id}] {chunk.title} | {chunk.doc_type}, version {chunk.version}, {dates} | "
            f"section: {chunk.section or '-'}\n{chunk.text}")


@dataclass
class Context:
    chunks: list = field(default_factory=list)
    dropped: list = field(default_factory=list)     # (chunk_id, reason)
    tokens: int = 0

    @property
    def ids(self) -> list[str]:
        return [c.chunk_id for c in self.chunks]

    def text(self) -> str:
        return "\n\n".join(format_passage(c) for c in self.chunks)


def assemble(chunks: list, budget: int = BUDGET) -> Context:
    ctx = Context()
    for chunk in chunks:
        twin = next((c for c in ctx.chunks if c.doc_id != chunk.doc_id and jaccard(c.text, chunk.text) >= SIMILAR), None)
        if twin is not None:
            ctx.dropped.append((chunk.chunk_id, f"near-duplicate of {twin.chunk_id}"))
            continue
        cost = estimate_tokens(format_passage(chunk))
        if ctx.tokens + cost > budget:
            ctx.dropped.append((chunk.chunk_id, "over budget"))
            continue
        ctx.chunks.append(chunk)
        ctx.tokens += cost
    return ctx

Line	Meaning
`text.split("\n", 1)[1]`	Leave out the first line, the document title. Otherwise the titles “PW-2200” and “PW-2400” would make two equal paragraphs look different.
`c.doc_id != chunk.doc_id`	Compare only passages of different documents: two versions of one document are never merged.
`estimate_tokens(format_passage(chunk))`	The header counts too: the model reads it.
`continue` after “over budget”	Skip the passage and try the next, smaller one. You saw the effect at 300 tokens.
`ctx.dropped`	Every dropped passage, with the reason. The next lesson prints it.

Make the file tests/test_context.py:

from types import SimpleNamespace

from policy_assistant.context import assemble, estimate_tokens, format_passage, jaccard


def chunk(cid, doc_id, text, version="1"):
    return SimpleNamespace(chunk_id=cid, doc_id=doc_id, version=version, title="T", doc_type="guide",
                           effective_from="2026-01-01", effective_to="", section="S", text=text)


SAFETY = "Never run the pressure washer for more than 2 minutes with the trigger released: the pump overheats."


def test_estimate_tokens():
    assert estimate_tokens("abcd" * 10) == 10


def test_near_duplicates_from_two_documents_are_removed():
    a = chunk("a", "product-care-pw2200", "PW-2200 guide\n" + SAFETY)
    b = chunk("b", "product-care-pw2400", "PW-2400 guide\n" + SAFETY)
    assert jaccard(a.text, b.text) == 1.0                       # the title line is left out
    ctx = assemble([a, b])
    assert ctx.ids == ["a"] and ctx.dropped == [("b", "near-duplicate of a")]


def test_two_versions_of_one_document_are_both_kept():
    v3 = chunk("v3", "returns-policy", "Returns\nYou can return an unused item within 30 days.", "3")
    v4 = chunk("v4", "returns-policy", "Returns\nYou can return an unused item within 30 days.", "4")
    assert assemble([v3, v4]).ids == ["v3", "v4"]


def test_budget():
    chunks = [chunk(str(i), f"d{i}", " ".join(f"w{i}x{j}" for j in range(200))) for i in range(5)]
    cost = estimate_tokens(format_passage(chunks[0]))      # 342 estimated tokens per passage
    ctx = assemble(chunks, budget=2 * cost + 10)
    assert len(ctx.ids) == 2 and ctx.tokens == 2 * cost
    assert [reason for _, reason in ctx.dropped] == ["over budget"] * 3


def test_a_passage_shows_its_id_dates_and_type():
    text = format_passage(chunk("3fa9c2d1b0e4", "d", "Title\nbody"))
    assert text.startswith("[3fa9c2d1b0e4] T | guide, version 1, in force from 2026-01-01, no end date | section: S\n")

The tests use small made-up passages (SimpleNamespace), so they need no index. The versions test uses the same text twice on purpose: even identical versions stay.

Run the new tests, then all of them:
Terminal window
```
python -m pytest tests/test_context.py
python -m pytest
```
The first output ends with 5 passed, and the second with 53 passed.

Start python in the project folder, and assemble the context of Q30. The first lines load the two local models, as search does:

from policy_assistant.store import Store
from policy_assistant.embed import LocalEmbedder
from policy_assistant.rerank import Reranker
from policy_assistant.search import Retriever
from policy_assistant.filters import Filters
from policy_assistant.context import assemble, estimate_tokens, format_passage

retriever = Retriever(Store("index.sqlite"), LocalEmbedder(), Reranker())
question = "How long is the warranty on the pump of a pressure washer?"
results = retriever.search(question, method="rerank", k=5, filters=Filters(as_of="2026-10-09", audience="staff"))
for r in results:
    print(r.chunk.chunk_id, r.chunk.doc_id, estimate_tokens(format_passage(r.chunk)))

context = assemble([r.chunk for r in results])
print(context.ids, context.tokens)
print(context.dropped)

Press Enter on the empty line after the for loop. Python may also print lines about loading the models. The output:

ff8235f48619 warranty-claims-procedure 104
d9a5bd8ea531 warranty-policy 314
50f5713033a7 product-care-pw2400 92
08705a17f766 product-care-pw2200 92
34a0b9b8055e product-care-pw2400 107
['ff8235f48619', 'd9a5bd8ea531', '50f5713033a7', '34a0b9b8055e'] 617
[('08705a17f766', 'near-duplicate of 50f5713033a7')]

These are the numbers of the diagram. Keep Python open for the next section.

Your turn: the budget of 300

Predict first: with budget=300, which IDs does assemble() keep for Q30, and what are the reasons for the others? Then type, in the same Python session:

small = assemble([r.chunk for r in results], budget=300)
print(small.ids, small.tokens)
print(small.dropped)

Show the answer

['ff8235f48619', '50f5713033a7'] 196
[('d9a5bd8ea531', 'over budget'), ('08705a17f766', 'near-duplicate of 50f5713033a7'), ('34a0b9b8055e', 'over budget')]

The warranty table (314 tokens) does not fit after the procedure (104), so it is skipped. The winter-storage passage (92) still fits, so it takes the space. Its PW-2200 copy is dropped as a near-duplicate, and the safe-use passage (107) would bring the total to 303. The context now holds the staff procedure, which says “within 3 years”, and a passage about frost. A customer-facing assistant would not see the procedure, so it would have no answer. Leave Python with exit().

When to use a budget, and what it costs

Use a budget for every model call that carries retrieved text. It keeps the cost and the time of a question predictable, and it makes the weakest passages, not random ones, go first.

The whole collection in one prompt is a fair choice in one case. The collection is small and does not grow, it has no access labels and no versions, and the cost per question does not matter. Larkfield’s collection fails three of these conditions.

What the budget costs you:

It can drop the answer. A budget smaller than your largest useful passage, here the warranty table, removes it. Check the size of your tables.
The estimate is rough. Characters ÷ 4 fits English text. French text, product codes and IDs use more tokens per character. Leave room.
The near-duplicate rule can be wrong. It compares words, not meaning. Two short passages from two different documents that differ in one important number can look the same. The PW guides’ Parts sections differ only in part numbers, and their similarity is 0.818: below 0.9, so both stay. Check the threshold on your own documents.

Common mistake

The near-duplicate rule that deleted the new policy

Symptom: Tomás simplifies assemble(): “Why only across documents? A copy is a copy.” He removes c.doc_id != chunk.doc_id. One test fails: test_two_versions_of_one_document_are_both_kept. He thinks the test is too strict and deletes it. Now, whenever the date filter is off, a returns question can lose one of the two Return window passages. No message warns anyone.

Cause: the Return window sections of versions 3 and 4 have a Jaccard similarity of 0.923. The changed rule drops the second one as a near-duplicate. Whichever version ranks first is the only one that the model sees, so the date no longer decides.

Fix: never merge two passages of the same document. Keep the test, and add a real case: the version 3 and version 4 Return window passages must both stay. Let the date filter, not a similarity score, choose the version.

What it means for a real AI service: a rule that removes “copies” must know what a copy is. In versioned documents, the small differences are the content.

Knowledge check

Check your understanding

On these 20 questions, the whole collection in the prompt answered as well as search, or better. Give two reasons why Larkfield still uses search with a budget.

Show the answer

Any two of these. The cost: about 17,850 input tokens per question instead of about 900. For chat-strong at 45 questions a day, that is US$1.70 instead of US$0.14, and it grows with every document. Access: staff-only procedures go to the model for every question; chat-small answered Q01 from the staff procedure. Versions: both returns policies are always in the prompt, and only the model’s reading chooses. Citations: with 185 IDs, the model made up a citation (c23617e24d2). The result on 36 documents does not hold for a collection that grows.

Check your understanding

Predict: a question's top 5 holds the version 3 and the version 4 Return window passages (Jaccard 0.923), and the date filter is off. What does assemble() keep, and why?

Show the answer

It keeps both, if they fit in the budget. The near-duplicate rule only compares passages of different documents, and both passages belong to the returns policy. Their only different words, 30 and 45, are the answer. Each passage keeps its header with its version and dates, so the model and a person can see which version applies.

Check your understanding

Why does each passage enter the prompt with a header that holds its ID, type, version, dates and section?

Show the answer

The ID lets the model cite the passage, and lets code check the citation. The type tells which document wins when two disagree (a policy wins over a help article). The version and dates show which version is in force on the question’s date. The section tells the agent where to look. Without the header, the text is anonymous: the answer cannot cite it and nobody can check it.

Recap

The context is the text that goes into the prompt with the question. The context budget limits it: Larkfield’s is 1,500 estimated tokens (characters ÷ 4).
assemble() keeps the search order and drops near-duplicates from different documents (Jaccard similarity of 0.9 or more). It skips passages over the budget, and keeps each passage’s header: ID, title, type, version, dates and section.
Two versions of one document are never merged: version 3 and version 4 differ only in “30” and “45”.
The whole collection in one prompt answered as well on 36 documents. It cost 20 times the tokens, put staff text in every prompt and produced one made-up citation. It does not scale.
A budget that is too small can drop the best passage. Make your largest useful passage fit.

Next, you ask the model for an answer made of claims, each with the IDs of its passages, and you check every citation in code.

Next: Answer and cite