Retrieval: Finding the Right Context

Chunk a document so facts survive the boundaries, rank the chunks, take the top k, and build a prompt from what you found.

Syntaxtop = heapq.nlargest(k, chunks, key=lambda c: cosine(query_vec, c.vector))

Retrieval-augmented generation is three ordinary steps around one model call: split your documents into chunks, find the chunks most related to the question, and paste those into the prompt. The model supplies the language; your store supplies the facts.

Chunking, and why boundaries matter

A chunk is the unit you retrieve, so it has to be self-contained. Cut a document into fixed-size windows and a sentence like "the limit is 60 per minute" can end up with its subject in the previous chunk, which makes it unretrievable and, worse, quotable out of context. Two cheap fixes: overlap the windows by 10 to 20 percent so every sentence appears whole somewhere, and prefer natural boundaries (headings, paragraphs) over character counts where the format gives you any.

Storing and ranking

Each chunk is stored with its vector, its text and where it came from. For a few thousand chunks, a list and heapq.nlargest are genuinely enough; a dedicated vector database earns its place when the store is large or has to be shared between processes. Keep the source with the chunk so the answer can cite it.

Assembling the prompt

Order the chunks best-first, fence them, and instruct the model to answer only from the fenced context and to say so when the answer is not there. Put the question after the context. Then stop adding chunks when the token budget runs out, rather than hoping the window stretches.

What trips people up

Retrieval failures look like model failures. When an answer is wrong, print the chunks you retrieved first: nine times in ten the right chunk was never in the prompt.

Example

Try it yourself
Loading editor…
Press Run to execute the code.

When to use it

  • Answering questions over internal documentation the model has never seen
  • Grounding a support assistant in the current pricing page
  • Citing the source paragraph alongside a generated answer
  • Working out why a correct document was never retrieved

More examples

Chunking on structure, not characters

Structure-aware chunking for documents that have structure. A heading is a better boundary than a character offset because the author already decided that is where one idea ends, and carrying the heading into the source field gives you a citation for free. Fall back to size-based packing only inside a section that is too long.

Example · python
from dataclasses import dataclass

@dataclass
class Chunk:
    text: str
    source: str        # file and heading, so the answer can cite it
    tokens: int

def split_markdown(path, heading="## ", max_tokens=400):
    """Split on headings first; only fall back to size inside a long section."""
    body = path.read_text(encoding="utf-8")
    sections = body.split("\n" + heading)
    out = []
    for section in sections:
        title = section.splitlines()[0].strip()
        for part in pack_paragraphs(section, max_tokens):   # never mid-sentence
            out.append(Chunk(part, f"{path.name}#{title}", count_tokens(part)))
    return out

The whole pipeline, end to end

Retrieve, fit to a budget, fence, ask. Two details earn their keep: the source label inside each chunk, which makes citations possible and lets you audit a wrong answer, and the system prompt's explicit permission to say the answer is absent, without which the model will invent one.

Example · python
def answer(question: str, k: int = 4) -> str:
    query_vec = embed_one(question)
    top = heapq.nlargest(k, store, key=lambda c: cosine(query_vec, c.vector))

    budget = 12_000                         # tokens for context, see the cost lesson
    used = []
    for chunk in top:
        if chunk.tokens > budget:
            break
        used.append(chunk)
        budget -= chunk.tokens

    context = "\n\n".join(f"[{c.source}]\n{c.text}" for c in used)
    reply = client.messages.create(
        model="claude-opus-5-5", max_tokens=1024,
        system="Answer only from the provided context. Cite the [source] you "
               "used. If the context does not contain the answer, say so.",
        messages=[{"role": "user",
                   "content": f"<context>\n{context}\n</context>\n\n{question}"}],
    )
    return "".join(b.text for b in reply.content if b.type == "text")

Debugging retrieval before blaming the model

The first thing to run when a RAG answer is wrong. Printing the retrieved chunks with their scores separates a retrieval bug from a generation bug in one step, and it is the only way to tell whether raising k, changing the chunk size, or rewriting the query is the fix.

Example · python
def explain(question, k=5):
    query_vec = embed_one(question)
    for chunk in heapq.nlargest(k, store, key=lambda c: cosine(query_vec, c.vector)):
        print(f"{cosine(query_vec, chunk.vector):.3f}  {chunk.source}")
        print("    " + chunk.text[:100].replace("\n", " "))

explain("what does a 429 carry")
# If the chunk holding the answer is not in this list, no prompt change will
# fix the answer: the problem is chunking, the query, or k being too small.

Related

py ai embeddingspy ai tokenspy ai promptspy ai callpy data pathlib

Discussion

  • Be the first to comment on this lesson.