Contextual retrieval and query transformation

The big idea

Retrieval compares two pieces of text: a chunk and a query. Both are often worse than they look. A chunk cut out of its document loses the context around it: "Free plans: 20 per day. Paid plans: no limit." does not say which product or which limit. Two such chunks from different products look identical to any search. And the query is whatever the user typed: "is there a cap?", in the middle of a conversation, says nothing searchable at all.

So improve both ends. Contextual retrievalPrepending a short, model-written description of where a chunk sits in its document (title, section, what it is about) before embedding and keyword-indexing it. gives every chunk back a sentence or two of context before it is indexed. Query transformationRewriting, expanding or splitting the user's question with a language model before retrieval, so the text sent to search matches how the answer is written. turns the user's question into one or more queries that search well.

Think of a box of photocopied pages from many manuals. A page that just says "Limit: 20 per day" is useless until someone writes "Notes app, export settings" at the top. And when a colleague asks "is there a cap?", you first work out "they mean the Notes export limit", then go looking.

How it works

Contextual chunks

For each chunk, send the whole document and the chunk to a language model with an instruction like "write a short context that situates this chunk within the document, to improve search retrieval". The answer, usually a sentence or two, is prepended to the chunk, and the combined text is what gets embedded and keyword-indexed. A cheaper variant prepends just the document title and heading path from your metadata ( Document ingestion and chunking ); the model-written version also captures facts from elsewhere in the document, such as which company or quarter a table refers to.

This costs one model call per chunk at ingestion, not per query. Because every call for the same document repeats that document, prompt caching ( Prompt caching and cost control ) makes it much cheaper.

Query rewriting

Users write short, vague, misspelled or conversational questions. A rewrite step asks a model: given the conversation so far and this message, write a standalone search query. "Is there a cap?" after a discussion of exporting notes becomes "Notes export limit per day". This is nearly mandatory for chat interfaces, where follow-up questions rarely make sense on their own.

Multi-query

Generate several phrasings of the question, search with each, and fuse the results with reciprocal rank fusion ( Hybrid search and reranking ). Different phrasings hit different vocabulary, so recall goes up, at the cost of several searches per question.

HyDE: search with a hypothetical answer

Questions and answers are written differently, so their embeddings are not always close. HyDEHypothetical Document Embeddings: asking a model to write a plausible answer to the query and searching with the embedding of that hypothetical answer instead of the question. asks a model to write a plausible answer first, then searches with that. The fake answer may contain wrong details, but it is written like the real documents, so it lands in the right neighbourhood of the embedding space; the real documents then supply the facts.

Decomposition

Questions that combine several facts, such as "which allows more per day on a free plan, Notes exports or Mail sending?", need chunks that no single query will rank together. Split the question into sub-questions, retrieve for each, and give the model all the results. Agents take this further, deciding what to search next based on what they found ( The agent loop (ReAct) ).

Try it

Retrieval here is a toy word-overlap search, so you can see exactly why each chunk matches.

Try this:

1. Ask the daily export limit question with contextual chunks off. The top 2 are the general export page and the Mail limit chunk: the Notes limit is missing, and the model would likely answer 500, the wrong product's number. Turn contextual chunks on: the Notes limit chunk ranks first.
2. Pick the vague follow-up "is there a cap?". As asked, nothing matches. Choose "rewritten", which uses the conversation to produce "What is the daily export limit for Notes plans?": with contextual chunks on, the right chunk is found. Try HyDE and multi-query too.
3. Pick the comparison question with contextual chunks off. Both limit chunks are retrieved, but neither says which product it belongs to, so the model cannot tell which number is which. Context is not only for ranking; it is what makes a retrieved chunk usable.

Going deeper

Contextual BM25 in code

# Contextual retrieval in miniature: prepend each chunk's document context, then index with BM25.
# `context_for` stands in for the LLM call that writes 50-100 tokens of context per chunk.
import math
import re

DOCS = {"Notes help": {"Export limits": "Free plans: 20 per day. Paid plans: no limit.",
                       "Export": "You can export notes as Markdown or PDF."},
        "Mail help": {"Sending limits": "Free plans: 500 per day. Paid plans: no limit.",
                      "Attachments": "Attachments up to 25 MB are allowed."}}
CHUNKS = [(doc, sec, text) for doc, secs in DOCS.items() for sec, text in secs.items()]

def context_for(doc, sec):
    return f"This chunk is from the {doc}, section '{sec}'."

tok = lambda s: re.findall(r"[a-z0-9]+", s.lower())

def bm25_rank(texts, query, k1=1.2, b=0.75):
    docs = [tok(t) for t in texts]
    avg = sum(map(len, docs)) / len(docs)
    df = {w: sum(w in d for d in docs) for w in set(tok(query))}
    def score(d):
        s = 0.0
        for w in set(tok(query)):
            if df[w] == 0:
                continue
            idf = math.log(1 + (len(docs) - df[w] + 0.5) / (df[w] + 0.5))
            tf = d.count(w)
            s += idf * tf * (k1 + 1) / (tf + k1 * (1 - b + b * len(d) / avg))
        return s
    return sorted(range(len(docs)), key=lambda i: -score(docs[i])), [round(score(d), 2) for d in docs]

query = "notes export limit per day"
for label, texts in [("plain chunks", [t for _, _, t in CHUNKS]),
                     ("contextual chunks", [f"{context_for(d, s)} {t}" for d, s, t in CHUNKS])]:
    order, scores = bm25_rank(texts, query)
    print(f"{label}:")
    for i in order[:2]:
        d, s, t = CHUNKS[i]
        print(f"  {scores[i]:5.2f}  [{d} / {s}] {t}")

plain chunks:
   2.44  [Notes help / Export] You can export notes as Markdown or PDF.
   2.00  [Notes help / Export limits] Free plans: 20 per day. Paid plans: no limit.
contextual chunks:
   3.37  [Notes help / Export limits] Free plans: 20 per day. Paid plans: no limit.
   2.02  [Mail help / Sending limits] Free plans: 500 per day. Paid plans: no limit.

On plain chunks, the general export page outranks the limit chunk, and the Notes and Mail limit chunks score exactly the same (2.00): they are the same words. With one sentence of context prepended, the Notes limit chunk jumps to first and the two limit chunks become distinguishable.

Costs and risks of query transformation

Every transformation adds a model call, typically tens to hundreds of milliseconds, before retrieval can start. Use a small, fast model for it, and run variants in parallel. Rewrites can also go wrong, dropping a crucial term or answering the question themselves, so log the transformed queries and include them in error analysis ( Error analysis ). A common safeguard is to search with both the original and the rewritten query.

When not to bother

If your knowledge base is small enough to fit in the prompt, you may not need retrieval at all ( Prompting vs RAG vs fine-tuning: choosing ). And if evaluation shows retrieval already finds the right chunks for nearly every question, extra stages only add latency. Add each technique when your evaluation set shows the failure it fixes ( Evaluating RAG ).

Trade-offs

Common mistakes

• Embedding chunks that do not say which product, version or document they come from.
• Searching with raw follow-up messages such as "and for the paid plan?" in a chat interface.
• Adding every transformation by default and doubling latency without measuring the gain.
• Not logging transformed queries, so a bad rewrite looks like a retrieval failure.
• Generating context for chunks without caching the document across calls, multiplying cost.

In the wild

Check yourself

Recap