Advanced RAG: Contextual Retrieval, Late Chunking, and Hybrid Search Beyond Naive Splitting

# Naive chunking (what most tutorials show)
def naive_chunk(text: str, chunk_size: int = 512,
                overlap: int = 50) -> list[str]:
    tokens = text.split()  # Simplified tokenization
    chunks = []
    for i in range(0, len(tokens), chunk_size - overlap):
        chunk = " ".join(tokens[i:i + chunk_size])
        chunks.append(chunk)
    return chunks
 
# Problems with naive_chunk:
# 1. Splits mid-sentence
# 2. Splits mid-code-block
# 3. No awareness of document structure
# 4. Embedded chunk has no document context

from anthropic import AsyncAnthropic
from typing import Optional
import asyncio
 
client = AsyncAnthropic()
 
CONTEXT_GENERATION_PROMPT = """<document>
{document_content}
</document>
 
Here is a chunk from the above document:
<chunk>
{chunk_content}
</chunk>
 
Please give a short succinct context (50-100 tokens) to situate this chunk
within the overall document for the purpose of improving search retrieval.
Answer only with the succinct context and nothing else."""
 
async def generate_chunk_context(
    document: str,
    chunk: str,
    model: str = "claude-haiku-4-5-20251001",  # Use fast, cheap model
) -> str:
    """
    Generate context for a chunk using a language model.
    Uses a small, fast model since context generation runs for every chunk.
    """
    prompt = CONTEXT_GENERATION_PROMPT.format(
        document_content=document[:50000],  # Truncate very long docs
        chunk_content=chunk,
    )
 
    response = await client.messages.create(
        model=model,
        max_tokens=150,
        messages=[{"role": "user", "content": prompt}],
    )
 
    return response.content[0].text.strip()
 
async def create_contextual_chunks(
    document: str,
    chunks: list[str],
    model: str = "claude-haiku-4-5-20251001",
    batch_size: int = 20,
) -> list[str]:
    """
    Generate contextual versions of all chunks.
    Returns chunks with prepended context.
    """
    contextual_chunks = []
 
    # Process in batches to control concurrency
    for i in range(0, len(chunks), batch_size):
        batch = chunks[i:i + batch_size]
 
        contexts = await asyncio.gather(
            *[generate_chunk_context(document, chunk, model) for chunk in batch]
        )
 
        for chunk, context in zip(batch, contexts):
            contextual_chunk = f"{context}\n\n{chunk}"
            contextual_chunks.append(contextual_chunk)
 
    return contextual_chunks
 
# Cost estimation
def estimate_contextual_retrieval_cost(
    n_chunks: int,
    avg_document_tokens: int = 10000,
    avg_chunk_tokens: int = 500,
    haiku_price_per_million_input: float = 0.80,  # USD
    haiku_price_per_million_output: float = 4.00,
) -> dict:
    """
    Estimate cost of contextual retrieval preprocessing.
    With prompt caching, the document tokens cost 10% of normal.
    """
    # Without caching
    input_tokens_no_cache = n_chunks * (avg_document_tokens + avg_chunk_tokens)
    output_tokens = n_chunks * 100  # ~100 tokens per context
    cost_no_cache = (
        input_tokens_no_cache / 1e6 * haiku_price_per_million_input +
        output_tokens / 1e6 * haiku_price_per_million_output
    )
 
    # With prompt caching (document portion cached after first chunk)
    # First chunk: full document cost; subsequent chunks: 10% of document cost
    input_first_chunk = avg_document_tokens + avg_chunk_tokens
    input_cached_chunks = (n_chunks - 1) * (avg_document_tokens * 0.10 + avg_chunk_tokens)
    cost_with_cache = (
        (input_first_chunk + input_cached_chunks) / 1e6 * haiku_price_per_million_input +
        output_tokens / 1e6 * haiku_price_per_million_output
    )
 
    return {
        "n_chunks": n_chunks,
        "cost_no_cache_usd": round(cost_no_cache, 4),
        "cost_with_cache_usd": round(cost_with_cache, 4),
        "cache_savings_percent": round((1 - cost_with_cache / cost_no_cache) * 100, 1),
        "cost_per_chunk_usd": round(cost_with_cache / n_chunks, 5),
    }
 
# For 1000 chunks from a 10K-token document:
# Without cache: ~$0.82
# With cache: ~$0.08 (prompt caching reduces cost by ~90%)
print(estimate_contextual_retrieval_cost(1000))

import torch
import numpy as np
from transformers import AutoTokenizer, AutoModel
 
def late_chunking_embed(
    text: str,
    chunk_boundaries: list[tuple[int, int]],  # (start_char, end_char) pairs
    model_name: str = "jinaai/jina-embeddings-v2-base-en",
    max_length: int = 8192,
) -> list[np.ndarray]:
    """
    Embed chunks using late chunking: process full document, then chunk embeddings.
 
    chunk_boundaries: character-level boundaries for each chunk.
    Returns: list of embeddings, one per chunk.
    """
    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
    model = AutoModel.from_pretrained(model_name, trust_remote_code=True)
    model.eval()
 
    # Tokenize the full document
    encoding = tokenizer(
        text,
        return_tensors="pt",
        return_offsets_mapping=True,  # Get char-to-token mapping
        max_length=max_length,
        truncation=True,
    )
 
    offset_mapping = encoding.pop("offset_mapping")[0]  # [n_tokens, 2]
 
    # Forward pass on full document
    with torch.no_grad():
        outputs = model(**encoding)
 
    # Token-level embeddings: [1, n_tokens, hidden_size]
    token_embeddings = outputs.last_hidden_state[0]  # [n_tokens, hidden_size]
 
    # Convert char boundaries to token boundaries using offset mapping
    chunk_embeddings = []
 
    for start_char, end_char in chunk_boundaries:
        # Find tokens that fall within this chunk's character range
        token_mask = (
            (offset_mapping[:, 0] >= start_char) &
            (offset_mapping[:, 1] <= end_char) &
            (offset_mapping[:, 0] < offset_mapping[:, 1])  # Exclude special tokens
        )
 
        chunk_tokens = token_embeddings[token_mask]
 
        if len(chunk_tokens) == 0:
            # Fallback: use nearest token
            chunk_tokens = token_embeddings[:1]
 
        # Mean pool the token embeddings for this chunk
        chunk_emb = chunk_tokens.mean(dim=0)
 
        # L2 normalize
        chunk_emb = chunk_emb / (chunk_emb.norm() + 1e-8)
        chunk_embeddings.append(chunk_emb.numpy())
 
    return chunk_embeddings
 
def get_sentence_boundaries(text: str) -> list[tuple[int, int]]:
    """
    Get character boundaries for sentence-level chunks.
    Better than fixed-size for late chunking since sentences are natural units.
    """
    import re
    boundaries = []
    # Split on sentence boundaries
    sentences = re.split(r'(?<=[.!?])\s+', text)
 
    pos = 0
    for sentence in sentences:
        start = text.find(sentence, pos)
        if start == -1:
            continue
        end = start + len(sentence)
        boundaries.append((start, end))
        pos = end
 
    return boundaries
 
# Benchmark: late chunking vs standard chunking on BeIR datasets
# Dataset          | Standard | Late Chunking | Improvement
# SciFact          | 64.20%   | 66.10%        | +1.9% nDCG@10
# TREC-COVID       | 63.36%   | 64.70%        | +1.3% nDCG@10
# NFCorpus         | 23.46%   | 29.98%        | +6.5% nDCG@10
# Effect increases with document length (as context from earlier text matters more)

from rank_bm25 import BM25Okapi
import numpy as np
from typing import Optional
 
class HybridRetriever:
    """
    Hybrid BM25 + dense vector retrieval with Reciprocal Rank Fusion.
    """
 
    def __init__(self, embed_fn: callable, k1: float = 1.5, b: float = 0.75):
        """
        embed_fn: function that takes a text and returns a dense embedding vector
        k1, b: BM25 parameters (k1=1.5, b=0.75 are BM25+ defaults)
        """
        self.embed = embed_fn
        self.k1 = k1
        self.b = b
        self.chunks: list[str] = []
        self.embeddings: list[np.ndarray] = []
        self.bm25: Optional[BM25Okapi] = None
 
    def index(self, chunks: list[str]):
        """Index chunks for both BM25 and vector retrieval."""
        self.chunks = chunks
 
        # BM25 index: tokenized chunks
        tokenized = [chunk.lower().split() for chunk in chunks]
        self.bm25 = BM25Okapi(tokenized, k1=self.k1, b=self.b)
 
        # Dense index: embeddings
        self.embeddings = np.array([self.embed(chunk) for chunk in chunks])
 
    def retrieve(self, query: str, top_k: int = 20,
                  bm25_weight: float = 0.5,
                  vector_weight: float = 0.5) -> list[tuple[int, float, str]]:
        """
        Hybrid retrieval using Reciprocal Rank Fusion.
        Returns: [(chunk_idx, score, chunk_text), ...]
        """
        # BM25 retrieval
        query_tokens = query.lower().split()
        bm25_scores = self.bm25.get_scores(query_tokens)
        bm25_ranked = np.argsort(bm25_scores)[::-1][:top_k * 2]
 
        # Dense vector retrieval
        query_emb = np.array(self.embed(query))
        # Cosine similarity
        similarities = np.dot(self.embeddings, query_emb) / (
            np.linalg.norm(self.embeddings, axis=1) * np.linalg.norm(query_emb) + 1e-8
        )
        dense_ranked = np.argsort(similarities)[::-1][:top_k * 2]
 
        # Reciprocal Rank Fusion
        rrf_k = 60  # RRF constant (60 is standard)
        scores = {}
 
        for rank, idx in enumerate(bm25_ranked):
            if idx not in scores:
                scores[idx] = 0.0
            scores[idx] += bm25_weight / (rrf_k + rank + 1)
 
        for rank, idx in enumerate(dense_ranked):
            if idx not in scores:
                scores[idx] = 0.0
            scores[idx] += vector_weight / (rrf_k + rank + 1)
 
        # Sort by combined score
        sorted_results = sorted(scores.items(), key=lambda x: x[1], reverse=True)
 
        return [
            (idx, score, self.chunks[idx])
            for idx, score in sorted_results[:top_k]
        ]
 
    def retrieve_with_explanation(self, query: str, top_k: int = 5) -> list[dict]:
        """Retrieve with BM25 and vector scores separately for debugging."""
        query_tokens = query.lower().split()
        bm25_scores = self.bm25.get_scores(query_tokens)
        query_emb = np.array(self.embed(query))
        vector_scores = np.dot(self.embeddings, query_emb) / (
            np.linalg.norm(self.embeddings, axis=1) * np.linalg.norm(query_emb) + 1e-8
        )
 
        combined = self.retrieve(query, top_k)
 
        return [
            {
                "chunk_idx": idx,
                "text": text[:200],
                "rrf_score": score,
                "bm25_score": float(bm25_scores[idx]),
                "vector_score": float(vector_scores[idx]),
                "retrieval_source": (
                    "both" if bm25_scores[idx] > np.median(bm25_scores) and
                    vector_scores[idx] > np.median(vector_scores) else
                    "bm25" if bm25_scores[idx] > np.median(bm25_scores) else "vector"
                ),
            }
            for idx, score, text in combined
        ]

MaxSim(Q, D) = Σ_i max_j (Q_i · D_j^T)

import torch
import numpy as np
from transformers import AutoTokenizer, AutoModel
 
class ColBERTRetriever:
    """
    ColBERT-style late interaction retrieval.
    Uses single BERT model for both query and document encoding.
    """
 
    def __init__(self, model_name: str = "colbert-ir/colbertv2.0"):
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModel.from_pretrained(model_name)
        self.model.eval()
        self.doc_embeddings: list[torch.Tensor] = []
        self.docs: list[str] = []
 
    def encode_query(self, query: str) -> torch.Tensor:
        """Encode query into per-token embeddings: [n_query_tokens, dim]"""
        inputs = self.tokenizer(
            f"[Q] {query}",  # Query prefix
            return_tensors="pt",
            max_length=32,    # ColBERT uses short query sequences
            truncation=True,
        )
        with torch.no_grad():
            outputs = self.model(**inputs)
        embeddings = outputs.last_hidden_state[0]  # [n_tokens, dim]
        # L2 normalize each token embedding
        return torch.nn.functional.normalize(embeddings, dim=-1)
 
    def encode_document(self, doc: str) -> torch.Tensor:
        """Encode document into per-token embeddings: [n_doc_tokens, dim]"""
        inputs = self.tokenizer(
            f"[D] {doc}",  # Document prefix
            return_tensors="pt",
            max_length=180,   # ColBERT uses longer document sequences
            truncation=True,
        )
        with torch.no_grad():
            outputs = self.model(**inputs)
        embeddings = outputs.last_hidden_state[0]
        return torch.nn.functional.normalize(embeddings, dim=-1)
 
    def maxsim(self, query_embeddings: torch.Tensor,
                doc_embeddings: torch.Tensor) -> float:
        """
        Compute ColBERT MaxSim score.
        For each query token, find max similarity to any document token.
        Sum across query tokens.
        """
        # [n_query_tokens, n_doc_tokens] similarity matrix
        similarities = query_embeddings @ doc_embeddings.T
 
        # For each query token: max similarity across document tokens
        max_similarities = similarities.max(dim=1).values
 
        # Sum across query tokens (excluding [Q] special token)
        return max_similarities[1:].sum().item()  # Skip [Q] token
 
    def index(self, documents: list[str]):
        """Pre-encode all documents."""
        self.docs = documents
        self.doc_embeddings = [self.encode_document(doc) for doc in documents]
 
    def retrieve(self, query: str, top_k: int = 10) -> list[tuple[int, float, str]]:
        """Retrieve top_k documents by MaxSim score."""
        query_emb = self.encode_query(query)
 
        scores = [
            (i, self.maxsim(query_emb, doc_emb))
            for i, doc_emb in enumerate(self.doc_embeddings)
        ]
 
        scores.sort(key=lambda x: x[1], reverse=True)
        return [(idx, score, self.docs[idx]) for idx, score in scores[:top_k]]

from sentence_transformers import CrossEncoder
 
class Reranker:
    """
    Cross-encoder reranker for retrieved chunks.
    Use as the final step after initial retrieval.
    """
 
    def __init__(self, model_name: str = "cross-encoder/ms-marco-MiniLM-L-6-v2"):
        self.model = CrossEncoder(model_name)
 
    def rerank(self, query: str, chunks: list[tuple[int, str]],
               top_n: int = 5) -> list[tuple[int, float, str]]:
        """
        Rerank retrieved chunks using cross-encoder scores.
        chunks: [(chunk_idx, chunk_text), ...]
        Returns: top_n chunks with cross-encoder scores, highest first
        """
        pairs = [(query, chunk_text) for _, chunk_text in chunks]
        scores = self.model.predict(pairs)
 
        scored = [(idx, float(score), text)
                  for (idx, text), score in zip(chunks, scores)]
        scored.sort(key=lambda x: x[1], reverse=True)
 
        return scored[:top_n]

class AdvancedRAGPipeline:
    """
    Full advanced RAG pipeline combining:
    1. Contextual chunk generation
    2. Late chunking embeddings (for long documents)
    3. Hybrid BM25 + vector retrieval
    4. Reranking
    """
 
    def __init__(self, embed_fn, reranker, llm_client):
        self.retriever = HybridRetriever(embed_fn)
        self.reranker = Reranker()
        self.llm = llm_client
 
    async def index_document(self, document: str,
                               use_late_chunking: bool = True,
                               use_contextual_retrieval: bool = True):
        """Full indexing pipeline for a document."""
        # 1. Semantic chunking (structure-aware)
        chunks = semantic_chunk(document)  # Not shown: structure-aware chunking
 
        # 2. Contextual retrieval: generate context for each chunk
        if use_contextual_retrieval:
            chunks = await create_contextual_chunks(document, chunks)
 
        # 3. Index in hybrid retriever
        if use_late_chunking and len(document.split()) < 6000:
            # Late chunking for manageable-length docs
            boundaries = [(document.find(c), document.find(c) + len(c)) for c in chunks]
            embeddings = late_chunking_embed(document, boundaries)
            self.retriever.index_with_embeddings(chunks, embeddings)
        else:
            self.retriever.index(chunks)
 
    async def query(self, query: str, top_k_retrieve: int = 20,
                     top_n_rerank: int = 5) -> dict:
        """Full retrieval + generation pipeline."""
        # 1. Hybrid retrieval
        retrieved = self.retriever.retrieve(query, top_k=top_k_retrieve)
 
        # 2. Rerank
        chunks_for_rerank = [(idx, text) for idx, _, text in retrieved]
        reranked = self.reranker.rerank(query, chunks_for_rerank, top_n=top_n_rerank)
 
        # 3. Generate answer
        context = "\n\n".join(f"[Source {i+1}]\n{text}" for i, (_, _, text) in enumerate(reranked))
        answer = await self.llm.complete(
            f"Based on these sources:\n{context}\n\nAnswer: {query}"
        )
 
        return {
            "answer": answer,
            "sources": [{"rank": i+1, "score": score, "text": text[:200]}
                        for i, (_, score, text) in enumerate(reranked)],
            "n_retrieved": len(retrieved),
            "n_reranked": len(reranked),
        }