Why Your RAG Pipeline Assembles Context Wrong

# Wrong way: the most common context assembly pattern in production RAGdef assemble_context_naive(query: str, vectorstore, llm) -> str:    """    What most RAG tutorials show, and what most production systems ship with.        Problems:    1. Ordering: chunks returned in retrieval score order, not LLM-attention order.       The top-scored chunk goes first. But if there are 5 chunks, positions 2-4       sit in the middle of the context window - the attention dead zone.           2. No deduplication: if two source documents contain the same paragraph,       both are retrieved and both are passed to the LLM. The model reads       the same sentence twice, paying token cost twice, and potentially       amplifying that content relative to non-duplicated information.           3. No budget policy: the retrieved context is concatenated until complete.       If the system prompt is long, or conversation history is long,       or the 5 retrieved chunks are each 500 tokens, the total context       may exceed the budget the system was designed for. No enforcement.           4. No compression: if the budget is exceeded, truncation drops everything       past the limit. The chunk at position 5 may be partially or fully cut.       The model has no idea the context was truncated.    """    docs = vectorstore.similarity_search(query, k=5)        # Concatenated in retrieval order, no processing    context = "\n\n".join([doc.page_content for doc in docs])        prompt = f"Context:\n{context}\n\nQuestion: {query}\nAnswer:"    return llm.invoke(prompt).content

from dataclasses import dataclass@dataclassclass ScoredChunk:    text: str    score: float       # Cross-encoder reranker score from Part 4    doc_id: str    source: strdef bookend_ordering(chunks: list[ScoredChunk]) -> list[ScoredChunk]:    """    Place highest-scored chunks at first and last positions.    Fill middle positions with remaining chunks in descending score order.        This directly counteracts the U-shaped attention bias documented by    Liu et al. (TACL 2024): the model attends most strongly to beginning    and end. Putting your best evidence there maximizes the probability    the model uses it.        For k=5 chunks:    Input order (reranker score desc): [A, B, C, D, E]    Bookend order:                     [A, C, D, E, B]    A (best) at position 1, B (second best) at position 5.    C, D, E fill the middle.    """    if not chunks:        return []    if len(chunks) == 1:        return chunks        sorted_chunks = sorted(chunks, key=lambda x: x.score, reverse=True)        best = sorted_chunks[0]    second_best = sorted_chunks[1]    middle = sorted_chunks[2:]        # Bookend: best first, second-best last, rest in middle    return [best] + middle + [second_best]def ordering_for_size(chunks: list[ScoredChunk]) -> list[ScoredChunk]:    """    Alternative: descending relevance for small sets (k <= 3),    bookend for larger sets (k > 3) where mid-context degradation    becomes significant.    """    if len(chunks) <= 3:        return sorted(chunks, key=lambda x: x.score, reverse=True)    return bookend_ordering(chunks)

import hashlibimport numpy as npfrom sentence_transformers import SentenceTransformerimport hashlibimport numpy as npfrom sentence_transformers import SentenceTransformerdef exact_dedup(chunks: list[ScoredChunk]) -> list[ScoredChunk]:    """    Hash-based exact deduplication.    Removes verbatim duplicates: same text from two different source documents.    Preserves the higher-scored copy when duplicates are found.        Cost: O(n) - fast, runs before semantic dedup.    """    seen_hashes: set[str] = set()    deduped: list[ScoredChunk] = []        # Sort by score desc so we keep the higher-scored copy on collision    for chunk in sorted(chunks, key=lambda x: x.score, reverse=True):        content_hash = hashlib.md5(chunk.text.strip().encode()).hexdigest()        if content_hash not in seen_hashes:            seen_hashes.add(content_hash)            deduped.append(chunk)        return dedupeddef semantic_dedup(    chunks: list[ScoredChunk],    embed_model: SentenceTransformer,    similarity_threshold: float = 0.92,) -> list[ScoredChunk]:    """    Embedding-based semantic deduplication.    Removes near-duplicate chunks: same content paraphrased or slightly reformatted.        Threshold calibration:      0.95+: only near-verbatim duplicates removed      0.90-0.94: removes close paraphrases (recommended range)      0.85-0.89: aggressive; may remove genuinely distinct chunks on the same topic        Cost: one embedding call per chunk. Run exact dedup first to reduce n.    """    if len(chunks) <= 1:        return chunks        texts = [c.text for c in chunks]    embeddings = embed_model.encode(texts, normalize_embeddings=True)        kept_indices: list[int] = []    dropped: set[int] = set()        # Greedy: keep highest-scored; drop anything too similar to what we kept    sorted_indices = sorted(        range(len(chunks)),        key=lambda i: chunks[i].score,        reverse=True,    )        for idx in sorted_indices:        if idx in dropped:            continue                kept_indices.append(idx)                # Mark similar chunks as dropped        for other_idx in sorted_indices:            if other_idx == idx or other_idx in dropped:                continue            sim = float(np.dot(embeddings[idx], embeddings[other_idx]))            if sim >= similarity_threshold:                dropped.add(other_idx)        return [chunks[i] for i in sorted(kept_indices)]def dedup_pipeline(    chunks: list[ScoredChunk],    embed_model: SentenceTransformer,    semantic_threshold: float = 0.92,) -> list[ScoredChunk]:    """    Run exact dedup first (fast), then semantic dedup on the remainder.    """    after_exact = exact_dedup(chunks)    return semantic_dedup(after_exact, embed_model, semantic_threshold)

import tiktokenfrom dataclasses import dataclass, field@dataclassclass ContextBudget:    """    Explicit token budget for each context zone.    Total must be less than model max_tokens to leave room for generation.        Recommended starting allocation for a 128K model (Claude Sonnet 4.6):      system_prompt:    4,000 tokens  (stable instructions)      retrieved_docs:  32,000 tokens  (retrieved context)      conversation:    8,000 tokens   (rolling history)      query:           1,000 tokens   (current query)      generation:      4,000 tokens   (output budget)      buffer:          1,000 tokens   (safety margin)        Adjust per use case:      - High-context analytical queries: increase retrieved_docs, reduce conversation      - Customer support chat: increase conversation, reduce retrieved_docs      - Single-turn Q&A: set conversation to 0    """    system_prompt_tokens: int = 4_000    retrieved_docs_tokens: int = 32_000    conversation_tokens: int = 8_000    query_tokens: int = 1_000    generation_tokens: int = 4_000    buffer_tokens: int = 1_000        @property    def total(self) -> int:        return (            self.system_prompt_tokens            + self.retrieved_docs_tokens            + self.conversation_tokens            + self.query_tokens            + self.generation_tokens            + self.buffer_tokens        )def count_tokens(text: str, model: str = "gpt-4o") -> int:    """Count tokens using the model's own tokenizer."""    enc = tiktoken.encoding_for_model(model)    return len(enc.encode(text))def fit_chunks_to_budget(    chunks: list[ScoredChunk],    token_budget: int,    model: str = "gpt-4o",) -> list[ScoredChunk]:    """    Select chunks that fit within the retrieved_docs token budget.        Strategy: include chunks greedily by score until budget is exhausted.    A chunk that does not fully fit is dropped, not truncated mid-sentence.    Truncating mid-chunk is worse than excluding the chunk entirely -    the model receives an incomplete semantic unit.        Log when chunks are excluded: this is a signal that either the budget    is too tight or the retrieval is returning too many large chunks.    """    fitted: list[ScoredChunk] = []    tokens_used = 0    excluded = []        for chunk in chunks:  # Already ordered by bookend_ordering        chunk_tokens = count_tokens(chunk.text, model)                if tokens_used + chunk_tokens <= token_budget:            fitted.append(chunk)            tokens_used += chunk_tokens        else:            excluded.append((chunk.doc_id, chunk_tokens))        if excluded:        print(            f"Budget: excluded {len(excluded)} chunks "            f"({sum(t for _, t in excluded)} tokens) from context. "            f"Used {tokens_used}/{token_budget} tokens."        )        return fitted

from anthropic import Anthropicclient = Anthropic()def compress_chunk(chunk: ScoredChunk, target_tokens: int) -> ScoredChunk:    """    Compress a chunk to approximately target_tokens using a fast LLM.        Use for chunks that are relevant but exceed the per-chunk budget.    Preserve the factual content; strip narrative structure and examples.        When to compress vs drop:    - Compress: chunk score above relevance floor, but chunk too large    - Drop: chunk score below relevance floor (established in Part 4)        Note: compression introduces an LLM call per chunk.    Use only when budget is genuinely tight, not as a default operation.    """    prompt = f"""Compress the following text to approximately {target_tokens} tokens.Preserve all specific facts, numbers, dates, names, and technical details.Remove examples, analogies, and repetitive explanations.Output only the compressed text, nothing else.Text to compress:{chunk.text}"""        response = client.messages.create(        model="claude-haiku-4-5-20251001",  # Fast model for compression        max_tokens=target_tokens + 100,     # Slight headroom        messages=[{"role": "user", "content": prompt}]    )        compressed_text = response.content[0].text    return ScoredChunk(        text=compressed_text,        score=chunk.score,        doc_id=chunk.doc_id,        source=chunk.source,    )def selective_compress(    chunks: list[ScoredChunk],    token_budget: int,    relevance_floor: float = 0.0,    max_chunk_tokens: int = 2_000,    model: str = "gpt-4o",) -> list[ScoredChunk]:    """    Selectively compress chunks that exceed max_chunk_tokens    while leaving high-relevance, appropriately-sized chunks untouched.        Order of operations:    1. Drop chunks below relevance_floor entirely    2. Compress chunks above max_chunk_tokens to max_chunk_tokens    3. Return compressed + untouched chunks within budget    """    filtered = [c for c in chunks if c.score >= relevance_floor]        result: list[ScoredChunk] = []    tokens_used = 0        for chunk in filtered:        chunk_tokens = count_tokens(chunk.text, model)                if chunk_tokens > max_chunk_tokens:            # Compress: too large to include verbatim            compressed = compress_chunk(chunk, target_tokens=max_chunk_tokens)            compressed_tokens = count_tokens(compressed.text, model)                        if tokens_used + compressed_tokens <= token_budget:                result.append(compressed)                tokens_used += compressed_tokens        else:            if tokens_used + chunk_tokens <= token_budget:                result.append(chunk)                tokens_used += chunk_tokens        return result

def assemble_context(    query: str,    reranked_chunks: list[ScoredChunk],    budget: ContextBudget,    embed_model: SentenceTransformer,    citation_format: bool = True,    model: str = "gpt-4o",) -> tuple[str, list[str]]:    """    Four-stage context assembly pipeline.        Returns (assembled_context_string, citation_list)        Stage 1: Deduplication - remove redundant information    Stage 2: Budget fitting - select chunks within token budget    Stage 3: Compression  - compress oversized chunks if needed    Stage 4: Ordering     - bookend for LLM attention optimization        call after reranking (Part 4), before prompt construction.    """    if not reranked_chunks:        return "", []        # Stage 1: Dedup    deduped = dedup_pipeline(reranked_chunks, embed_model)        # Stage 2 + 3: Fit to budget with selective compression    compressed = selective_compress(        deduped,        token_budget=budget.retrieved_docs_tokens,        model=model,    )        # Stage 4: Reorder for attention optimization    ordered = ordering_for_size(compressed)        # Format with citations if requested    if citation_format:        formatted_chunks = []        citations = []                for i, chunk in enumerate(ordered, 1):            formatted_chunks.append(                f"[{i}] Source: {chunk.source}\n{chunk.text}"            )            citations.append(chunk.source)                context = "\n\n---\n\n".join(formatted_chunks)        return context, citations    else:        context = "\n\n".join([c.text for c in ordered])        return context, [c.source for c in ordered]def rag_with_context_assembly(    query: str,    retriever,    reranker,    embed_model: SentenceTransformer,    llm,    budget: ContextBudget,) -> dict:    """    Full RAG pipeline with context assembly as an explicit stage.        Parts 1-4 (retrieval, chunking, embedding, reranking) produce the candidates.    Context assembly (Part 8) determines what the LLM actually reads.    """    # Parts 1-4: retrieve and rerank    raw_docs = retriever.invoke(query)        pairs = [(query, doc.page_content) for doc in raw_docs]    rerank_scores = reranker.predict(pairs)        scored_chunks = [        ScoredChunk(            text=doc.page_content,            score=float(score),            doc_id=doc.metadata.get("doc_id", f"doc_{i}"),            source=doc.metadata.get("source", "unknown"),        )        for i, (doc, score) in enumerate(zip(raw_docs, rerank_scores))    ]        # Part 8: context assembly    context, citations = assemble_context(        query=query,        reranked_chunks=scored_chunks,        budget=budget,        embed_model=embed_model,    )        if not context:        return {"answer": "No relevant content found.", "citations": []}        prompt = f"""Use the following sources to answer the question.Cite sources inline using [1], [2] notation.{context}Question: {query}Answer:"""        answer = llm.invoke(prompt).content        return {        "answer": answer,        "citations": citations,        "chunks_retrieved": len(raw_docs),        "chunks_after_dedup_and_budget": len(reranked_chunks),    }