Lesson 16.18: Reranking, Query Transformation & Context Assembly

[ Total Knowledge Base: 1,000,000 Chunks ]
                │
                ▼  Stage 1: Fast Candidate Retrieval (Bi-Encoder / BM25)
[ Top-20 Candidate Pool ]  ──> Fast, low latency (< 20ms)
                │
                ▼  Stage 2: Deep Cross-Attention Reranking
[ Top-5 High-Precision Chunks ]  ──> Deep transformer relevance scoring
                │
                ▼
[ LLM Context Window ]
        

import os
import cohere
from dotenv import load_dotenv

load_dotenv()
co = cohere.ClientV2(api_key=os.environ["COHERE_API_KEY"])

def rerank_candidates(
    query: str,
    candidates: list[dict],
    top_n: int = 5
) -> list[dict]:
    if not candidates:
        return []

    # Structure text representation with metadata headers
    documents = [
        f"Source: {c.get('source', '')}\nSection: {c.get('section', '')}\nContent: {c['text']}"
        for c in candidates
    ]

    response = co.rerank(
        model="rerank-v4.0-pro",
        query=query,
        documents=documents,
        top_n=min(top_n, len(candidates)),
    )

    reranked = []
    for hit in response.results:
        original_chunk = candidates[hit.index]
        reranked.append({
            **original_chunk,
            "rerank_score": hit.relevance_score, # Model relevance confidence
        })

    return reranked

from sentence_transformers import CrossEncoder

# Load model (e.g., BAAI/bge-reranker-large or ms-marco)
local_reranker = CrossEncoder("BAAI/bge-reranker-large", max_length=512)

def local_rerank(query: str, candidates: list[dict], top_n: int = 5) -> list[dict]:
    if not candidates:
        return []

    # Construct (Query, Document) pairs
    pairs = [[query, c["text"]] for c in candidates]
    
    # Predict relevance scores
    scores = local_reranker.predict(pairs)

    for i, c in enumerate(candidates):
        c["rerank_score"] = float(scores[i])

    # Sort descending by relevance score
    ranked = sorted(candidates, key=lambda x: x["rerank_score"], reverse=True)
    return ranked[:top_n]

import tiktoken

tokenizer = tiktoken.get_encoding("cl100k_base")

def deduplicate_chunks(candidates: list[dict]) -> list[dict]:
    """Filters duplicate candidate chunks based on stable chunk identifiers."""
    seen = set()
    unique = []
    for c in candidates:
        cid = c.get("chunk_id")
        if cid and cid not in seen:
            seen.add(cid)
            unique.append(c)
    return unique

def assemble_context(
    reranked_chunks: list[dict],
    token_budget: int = 6000
) -> tuple[str, list[dict]]:
    """Packs reranked chunks within the token budget and formats citation blocks."""
    context_blocks = []
    selected_chunks = []
    accumulated_tokens = 0

    for idx, chunk in enumerate(reranked_chunks, 1):
        text = chunk.get("text", "").strip()
        source = chunk.get("source", "Unknown Document")
        page = chunk.get("page_number", "N/A")
        
        # Structured citation block
        block = (
            f"[Source {idx}]\n"
            f"Document: {source} (Page {page})\n"
            f"Content: {text}"
        )
        
        block_tokens = len(tokenizer.encode(block)) + 4 # Separator overhead
        
        if accumulated_tokens + block_tokens > token_budget:
            # Token budget exceeded; stop packing
            break
            
        context_blocks.append(block)
        selected_chunks.append(chunk)
        accumulated_tokens += block_tokens

    formatted_context = "\n\n---\n\n".join(context_blocks)
    return formatted_context, selected_chunks

from typing import Dict, Any, List

def retrieve_production_context(
    *,
    query: str,
    tenant_id: str,
    user_id: str,
    roles: set[str],
    candidate_limit: int = 25,
    final_limit: int = 5,
    token_budget: int = 6000
) -> Dict[str, Any]:
    # Step 1: Permission-aware hybrid candidate recall
    raw_candidates = hybrid_search(
        query=query,
        tenant_id=tenant_id,
        limit=candidate_limit
    )

    # Step 2: Deduplicate candidate chunks
    unique_candidates = deduplicate_chunks(raw_candidates)

    # Step 3: Deep Cross-Encoder reranking
    reranked = rerank_candidates(
        query=query,
        candidates=unique_candidates,
        top_n=min(len(unique_candidates), candidate_limit)
    )

    # Step 4: Authoritative server-side access & version validation
    authorized_chunks = [
        c for c in reranked 
        if verify_document_access(c["document_id"], user_id, roles)
    ][:final_limit]

    # Step 5: Pack context within strict token budget
    context_text, selected = assemble_context(
        reranked_chunks=authorized_chunks,
        token_budget=token_budget
    )

    return {
        "context": context_text,
        "citations": [
            {
                "chunk_id": c["chunk_id"],
                "source": c["source"],
                "page": c["page_number"]
            }
            for c in selected
        ]
    }

import pytest

def test_deduplication_removes_identical_chunks():
    """Verify duplicate chunks are purged correctly."""
    raw = [
        {"chunk_id": "c1", "text": "leave policy"},
        {"chunk_id": "c1", "text": "leave policy"},
        {"chunk_id": "c2", "text": "medical insurance"},
    ]
    unique = deduplicate_chunks(raw)
    assert len(unique) == 2
    assert [x["chunk_id"] for x in unique] == ["c1", "c2"]

def test_token_budget_never_exceeded():
    """Context assembler must never exceed the allocated token ceiling."""
    chunks = [
        {"chunk_id": f"c_{i}", "text": "Word " * 200, "source": "doc.pdf", "page_number": 1}
        for i in range(10)
    ]
    context, selected = assemble_context(chunks, token_budget=500)
    tokens_used = len(tokenizer.encode(context))
    assert tokens_used <= 500

def test_reranker_prioritizes_specific_answer():
    """Cross-encoder promotes exact answering chunks over generic policy text."""
    query = "How is leave calculated for employees joining after July?"
    candidates = [
        {"chunk_id": "c_gen", "text": "Employees receive 20 days annual leave.", "source": "p.pdf"},
        {"chunk_id": "c_spec", "text": "Employees joining after July receive prorated leave.", "source": "p.pdf"},
    ]
    reranked = rerank_candidates(query, candidates, top_n=2)
    assert reranked[0]["chunk_id"] == "c_spec"