Reranking and Context Assembly
Reranking and context assembly turn a broad retrieval result into the compact evidence packet that a generator can actually use. The outcome is not just a higher score; it is an ordered, cited set of passages that answers the user’s question within the model’s token budget while keeping enough source information to verify the answer.
In this generative AI course, this lesson sits after retrieval because vector search alone usually returns plausible neighbors, not the best final context. A RAG system may retrieve twenty to one hundred chunks, many of them redundant, stale, or only loosely related. Reranking estimates which candidates are most useful for the specific query. Context assembly then chooses, deduplicates, orders, and labels those passages so the model receives evidence instead of a noisy document pile.
What Reranking Changes
Initial retrieval is normally optimized for speed and recall. A dense embedding index, lexical BM25 index, or hybrid retriever finds candidates that might contain the answer. Reranking is a second stage optimized for precision. It compares the original query with each candidate, often using a cross-encoder or another scoring model that reads the query and passage together. Because that comparison is more expensive than vector distance, it is applied only to a candidate set such as the top 50 retrieval hits.
The key internal difference is interaction. A bi-encoder embedding system encodes the query and document separately, then compares vectors. A cross-encoder reads text pairs jointly, so it can notice exact conditions, negations, dates, entities, and whether a passage actually answers the question. For example, a query about “refund window for annual plans” may retrieve chunks about refunds, billing, cancellations, and plans. The reranker should lift the chunk containing “annual plans within 30 days” above a generic cancellation policy.
Pipeline Anatomy
A typical context pipeline has five records: query, candidate, scored candidate, selected passage, and prompt evidence. A candidate should carry at least doc_id, chunk_id, raw text, retrieval score, token count, source metadata, and permission metadata. A scored candidate adds one or more reranking signals. A selected passage adds assembly decisions such as final position, citation label, and any trimming that occurred.
Scores are not interchangeable. Cosine similarity, BM25 score, recency boost, authority weight, and cross-encoder relevance live on different scales. Production systems either normalize them, learn a weighted combination, or use rules such as “filter by permission and freshness, rerank by cross-encoder, then break ties by source priority.” The important design point is to store component scores separately so a bad answer can be diagnosed later.
from dataclasses import dataclass
from math import log
@dataclass(frozen=True)
class Candidate:
doc_id: str
chunk_id: str
text: str
retriever_score: float
cross_encoder_score: float
tokens: int
section: str
def normalized_score(candidate: Candidate) -> float:
retrieval_component = log(1 + max(candidate.retriever_score, 0.0))
return 0.35 * retrieval_component + 0.65 * candidate.cross_encoder_score
candidate = Candidate("policy", "p3", "Refunds are available for annual plans within 30 days.", 0.82, 0.94, 11, "Billing")
print(round(normalized_score(candidate), 3))
This first example defines a candidate and combines a fast retriever score with a more precise cross-encoder score. The printed value is 0.821. The exact formula is intentionally simple, but the structure matters: retrieval contributes enough to preserve broadly relevant candidates, while the cross-encoder dominates the final order.
From Ranked List to Context
After reranking, the highest-scoring chunks are not automatically the right prompt. Adjacent chunks can repeat the same paragraph, one source may crowd out all others, and a long chunk can consume the entire budget. Context assembly solves a constrained selection problem: maximize answer coverage and relevance while fitting a fixed token budget and preserving provenance.
Useful assembly terms include deduplication, which removes near-identical passages; diversity, which avoids selecting five chunks that all say the same thing; window expansion, which adds neighboring chunks when an answer depends on surrounding context; compression, which extracts only the relevant sentences; and citation mapping, which keeps generated claims tied to original sources. These operations should be deterministic when possible, because stable context makes evaluation and troubleshooting much easier.
from dataclasses import dataclass
@dataclass(frozen=True)
class Passage:
doc_id: str
chunk_id: str
text: str
score: float
tokens: int
def dedupe_passages(passages: list[Passage], overlap_threshold: float = 0.75) -> list[Passage]:
selected: list[Passage] = []
for passage in sorted(passages, key=lambda p: p.score, reverse=True):
words = set(passage.text.lower().split())
is_duplicate = False
for kept in selected:
kept_words = set(kept.text.lower().split())
overlap = len(words & kept_words) / max(1, min(len(words), len(kept_words)))
if overlap >= overlap_threshold:
is_duplicate = True
break
if not is_duplicate:
selected.append(passage)
return selected
items = [
Passage("handbook", "a", "Employees may expense train travel with manager approval.", 0.91, 9),
Passage("handbook", "b", "Employees may expense rail travel with manager approval.", 0.86, 9),
Passage("security", "c", "Laptop encryption is required before international travel.", 0.72, 8),
]
print([p.chunk_id for p in dedupe_passages(items)])
The second example keeps the highest-scoring travel expense passage, drops a near-duplicate phrased with “rail” instead of “train,” and keeps the encryption passage. Its deterministic output is ['a', 'c']. In a real system, duplicate detection might use shingles or embeddings, but the behavior is the same: remove repeated evidence before spending prompt tokens.
Budgeting, Ordering, and Citations
Token budget is the assembly constraint most teams feel first. The prompt must include system instructions, user question, evidence, tool results, and space for the answer. If the model has room for 8,000 input tokens, the context assembler might allocate 1,000 to instructions and conversation, 5,500 to evidence, and leave the remainder as buffer. The evidence budget should be explicit; otherwise a single large retrieved document can silently push out instructions or answer space.
Ordering affects model behavior. Score order puts the strongest evidence first. Source order preserves document flow, which helps when chunks form a procedure. Chronological order helps incident reports and policy histories. Many RAG applications use a hybrid: select by score and diversity, then present selected chunks in source order with citation labels. This keeps the prompt readable without letting early, weak chunks win selection.
from dataclasses import dataclass
@dataclass(frozen=True)
class Passage:
doc_id: str
chunk_id: str
text: str
score: float
tokens: int
section: str
def assemble_context(passages: list[Passage], token_budget: int) -> list[Passage]:
chosen: list[Passage] = []
used = 0
seen_sections: set[str] = set()
ordered = sorted(passages, key=lambda p: (p.section in seen_sections, -p.score))
while ordered:
passage = ordered.pop(0)
if used + passage.tokens <= token_budget:
chosen.append(passage)
used += passage.tokens
seen_sections.add(passage.section)
ordered = sorted(ordered, key=lambda p: (p.section in seen_sections, -p.score))
return sorted(chosen, key=lambda p: (p.doc_id, p.chunk_id))
candidates = [
Passage("runbook", "01", "Service owners rotate API keys every 90 days.", 0.93, 8, "security"),
Passage("runbook", "02", "Expired keys return 401 until the client refreshes credentials.", 0.89, 10, "failure"),
Passage("faq", "07", "Key rotation does not change customer-visible request limits.", 0.78, 8, "limits"),
]
print([(p.doc_id, p.chunk_id) for p in assemble_context(candidates, 17)])
The third example assembles a small context under a 17-token budget. It selects chunks from different sections when possible, then returns them in source order for presentation. The expected output is [('runbook', '01'), ('faq', '07')]. The failure chunk is relevant, but it no longer fits after security and limits coverage are chosen.
from dataclasses import dataclass
@dataclass(frozen=True)
class Passage:
doc_id: str
chunk_id: str
text: str
score: float
tokens: int
def build_prompt(question: str, passages: list[Passage]) -> str:
evidence = []
for index, passage in enumerate(passages, start=1):
evidence.append(f"[{index}] {passage.doc_id}#{passage.chunk_id}: {passage.text}")
return "Answer only from the evidence. Cite source numbers.\nQuestion: " + question + "\n" + "\n".join(evidence)
prompt = build_prompt("When are refunds available?", [Passage("policy", "p3", "Refunds are available for annual plans within 30 days.", 0.94, 11)])
print(prompt)
The fourth example shows the final prompt shape. Its output starts with an instruction to answer only from evidence, repeats the user question, and labels the source as [1] policy#p3. Those labels are not cosmetic. They are the link that lets the answer cite the selected passage and lets an evaluator check whether each claim is supported.
Design Choices and Trade-offs
The first choice is reranker type. A cross-encoder is usually more accurate but slower and more expensive than vector scoring. A lightweight LLM judge can handle nuanced relevance instructions, but it may be less stable and harder to calibrate. A learned ranker can combine many features, but it needs training data. For many applications, a practical starting point is hybrid retrieval for recall, cross-encoder reranking for the top candidates, and deterministic assembly rules.
The second choice is chunk strategy. Small chunks improve precise matching and citation quality, but they can lose necessary context. Large chunks preserve context, but they waste tokens and can bury the answer. Window expansion is a compromise: retrieve and score small chunks, then include one neighboring chunk only when the selected passage refers to something outside itself.
The third choice is diversity. Maximum marginal relevance and source caps reduce redundancy, but they can exclude several supporting passages from the same document. Use diversity when the user asks a broad question, and reduce it when the user asks for a specific clause, number, or procedure where one document should dominate.
Failure Modes and Troubleshooting
A common symptom is a confident answer citing irrelevant passages. The cause is often scale mixing: a large BM25 score overwhelms a reranker score, or a recency boost lifts a newer but weaker document. Diagnose it by logging each component score for the final selected passages and nearby rejected passages. Correct it by normalizing scores, changing the tie-breaker, or adding a relevance threshold after reranking.
Another symptom is an answer that misses an obvious exception, such as “refunds are unavailable after usage begins.” The cause is usually over-aggressive deduplication or context compression that removed the exception sentence. Reproduce the query, inspect the pre-deduped candidates, and compare raw chunks with assembled context. Correct it by lowering the duplicate threshold, preserving negation-bearing sentences, or expanding the selected chunk window.
A third symptom is unstable citations: repeated runs answer correctly but cite different sources. The cause may be nondeterministic ordering, equal scores without a stable tie-breaker, or asynchronous retrieval results merged in arrival order. Diagnose by replaying with fixed inputs and checking whether selected passage IDs change. Correct it with deterministic sorting by score, source priority, document ID, and chunk ID.
A fourth symptom is latency spikes. Cross-encoding too many candidates, calling the reranker serially, or assembling oversized prompts are common causes. Measure candidate count, reranker batch size, token count, and model latency separately. Correct the issue by reducing the first-stage candidate count, batching reranker calls, caching stable query-document scores where appropriate, or enforcing a hard evidence budget.
Security, Performance, and Reliability
Reranking must respect document permissions before scoring and assembly. Do not retrieve a broad corpus and rely on the generator to ignore unauthorized chunks. The candidate set should already be filtered by tenant, role, and data policy. The assembled context should also strip secrets that are not needed for the answer and carry provenance that can be audited.
Performance depends on narrowing expensive work. Retrieve broadly enough for recall, but rerank only the plausible set. Track recall-oriented metrics before reranking and answer-quality metrics after assembly, because a reranker cannot rescue a document that retrieval never returned. Reliability improves when the assembler has explicit degraded behavior: if reranking times out, use a conservative lexical-plus-vector order, reduce the answer scope, or ask the user for clarification instead of fabricating missing evidence.
Hands-on Lab
Prerequisites: Python 3, a terminal, and no external model service. This lab simulates a RAG context assembler so you can verify the mechanics without network calls.
- Create an in-memory list of passages with document IDs, chunk IDs, text, scores, token estimates, and sections.
- Run the score-combination function from the first example for each candidate and sort descending.
- Run deduplication from the second example and confirm that paraphrased duplicates do not both survive.
- Run the assembler with a small token budget, then increase the budget and observe which additional passage appears.
- Build the final prompt with citation labels and inspect whether every selected passage still has its original source ID.
Verification: the deduplication example should print ['a', 'c'], the budget example should print [('runbook', '01'), ('faq', '07')], and the prompt should contain [1] policy#p3. If those values differ, check sorting order, token counts, and whether your duplicate threshold changed. Cleanup is simply deleting the scratch script; no services, indexes, or credentials are created.
Assessment Exercises
- Given ten retrieved chunks where the top five all repeat the same policy paragraph, design an assembly rule that preserves one copy while still allowing exceptions from lower-ranked chunks to appear.
- A reranker improves offline relevance scores but production answers get worse. What measurements would you inspect to decide whether retrieval recall, score calibration, or context ordering caused the regression?
- For a legal assistant, would you present selected passages in score order or document order? Explain which user behavior and citation requirement drives your choice.
- Write a relevance threshold policy for the case where all reranked candidates are weak. What should the application return instead of a normal answer?
- Modify the token-budget example so each source can contribute at most one passage unless the user asks for a step-by-step procedure. What behavior should change?
Summary
Reranking and context assembly are the precision layer of RAG. Retrieval gathers possible evidence; reranking estimates query-specific usefulness; assembly turns selected passages into a bounded, ordered, cited prompt. Strong systems keep component scores visible, deduplicate carefully, budget tokens explicitly, preserve source IDs, and make failure states diagnosable. The result is a generator that answers from the best available evidence instead of from whatever chunks happened to arrive first.
