Naive RAG Is Dead. Here's What Replaced It.

import anthropic
import json
 
client = anthropic.Anthropic()
 
async def route_query(query: str) -> str:
    """Classify a query as simple or complex."""
    response = client.messages.create(
        model="claude-haiku-4-5-20251001",
        max_tokens=100,
        messages=[{
            "role": "user",
            "content": (
                f"Classify this query as SIMPLE or COMPLEX.\n"
                f"SIMPLE: single fact, direct lookup, unambiguous.\n"
                f"COMPLEX: multi-part, requires reasoning across documents, "
                f"ambiguous, or comparative.\n\n"
                f"Query: {query}\n\n"
                f"Respond with only SIMPLE or COMPLEX."
            )
        }]
    )
    classification = response.content[0].text.strip().upper()
    return "simple" if "SIMPLE" in classification else "complex"

async def decompose_query(query: str) -> list[str]:
    """Break a complex query into retrieval-optimized sub-queries."""
    response = client.messages.create(
        model="claude-sonnet-4-6-20250514",
        max_tokens=1024,
        messages=[{
            "role": "user",
            "content": (
                f"Break this query into 2-4 independent sub-queries, each "
                f"targeting a specific piece of information needed to answer "
                f"the original query. Each sub-query should be self-contained "
                f"and optimized for retrieval from a document store.\n\n"
                f"Query: {query}\n\n"
                f"Return a JSON array of strings. No explanation."
            )
        }]
    )
    return json.loads(response.content[0].text)

from dataclasses import dataclass
 
@dataclass
class RetrievedChunk:
    content: str
    source: str
    score: float
    method: str  # "semantic" or "keyword"
 
 
async def hybrid_retrieve(
    query: str,
    vector_store,
    bm25_index,
    top_k: int = 10
) -> list[RetrievedChunk]:
    """Retrieve using both semantic and keyword search, then fuse."""
    # Semantic search via embeddings
    semantic_results = await vector_store.query(
        query=query,
        top_k=top_k,
        include_metadata=True
    )
 
    # Keyword search via BM25
    keyword_results = bm25_index.search(query, top_k=top_k)
 
    # Reciprocal Rank Fusion (RRF)
    scores: dict[str, float] = {}
    chunks: dict[str, RetrievedChunk] = {}
 
    for rank, result in enumerate(semantic_results):
        doc_id = result["id"]
        scores[doc_id] = scores.get(doc_id, 0) + 1 / (60 + rank)
        chunks[doc_id] = RetrievedChunk(
            content=result["text"],
            source=result["metadata"]["source"],
            score=0,  # will be updated
            method="semantic"
        )
 
    for rank, result in enumerate(keyword_results):
        doc_id = result.doc_id
        scores[doc_id] = scores.get(doc_id, 0) + 1 / (60 + rank)
        if doc_id not in chunks:
            chunks[doc_id] = RetrievedChunk(
                content=result.text,
                source=result.source,
                score=0,
                method="keyword"
            )
        else:
            chunks[doc_id].method = "both"
 
    # Sort by fused score
    ranked = sorted(scores.items(), key=lambda x: x[1], reverse=True)
    result = []
    for doc_id, score in ranked[:top_k]:
        chunk = chunks[doc_id]
        chunk.score = score
        result.append(chunk)
 
    return result

@dataclass
class JudgmentResult:
    is_sufficient: bool
    relevant_chunks: list[int]  # indices of relevant chunks
    missing_info: str | None    # what's missing, if insufficient
    suggested_requery: str | None  # rewritten query to try
 
 
async def judge_retrieval(
    query: str,
    chunks: list[RetrievedChunk],
    max_retries: int = 2
) -> tuple[list[RetrievedChunk], JudgmentResult]:
    """Judge whether retrieved chunks can answer the query."""
    chunks_text = "\n\n".join(
        f"[Chunk {i}] (source: {c.source})\n{c.content}"
        for i, c in enumerate(chunks)
    )
 
    response = client.messages.create(
        model="claude-sonnet-4-6-20250514",
        max_tokens=1024,
        messages=[{
            "role": "user",
            "content": (
                f"You are a retrieval quality judge. Given a query and "
                f"retrieved chunks, determine:\n"
                f"1. Which chunks are relevant (by index)\n"
                f"2. Whether the relevant chunks contain sufficient "
                f"information to answer the query\n"
                f"3. If insufficient, what information is missing\n"
                f"4. If insufficient, suggest a rewritten query\n\n"
                f"Query: {query}\n\n"
                f"Retrieved chunks:\n{chunks_text}\n\n"
                f"Respond in JSON:\n"
                f'{{"is_sufficient": bool, "relevant_chunks": [int], '
                f'"missing_info": str|null, "suggested_requery": str|null}}'
            )
        }]
    )
 
    judgment = JudgmentResult(**json.loads(response.content[0].text))
 
    # Filter to only relevant chunks
    relevant = [chunks[i] for i in judgment.relevant_chunks if i < len(chunks)]
 
    return relevant, judgment

@dataclass
class GroundingResult:
    is_grounded: bool
    unsupported_claims: list[str]
    confidence: float  # 0-1
 
 
async def check_grounding(
    answer: str,
    chunks: list[RetrievedChunk],
    query: str
) -> GroundingResult:
    """Verify that the answer is grounded in the retrieved chunks."""
    context = "\n\n".join(c.content for c in chunks)
 
    response = client.messages.create(
        model="claude-sonnet-4-6-20250514",
        max_tokens=1024,
        messages=[{
            "role": "user",
            "content": (
                f"You are a grounding checker. Verify that every factual "
                f"claim in the answer is directly supported by the provided "
                f"context. Flag any claims that are not supported.\n\n"
                f"Question: {query}\n\n"
                f"Context:\n{context}\n\n"
                f"Answer:\n{answer}\n\n"
                f"Respond in JSON:\n"
                f'{{"is_grounded": bool, "unsupported_claims": [str], '
                f'"confidence": float}}'
            )
        }]
    )
 
    return GroundingResult(**json.loads(response.content[0].text))

async def agentic_rag(
    query: str,
    vector_store,
    bm25_index,
    max_retrieval_attempts: int = 3
) -> dict:
    """Full agentic RAG pipeline with routing, judging, and grounding."""
 
    # Step 1: Route
    complexity = await route_query(query)
 
    if complexity == "simple":
        # Fast path: retrieve and generate directly
        chunks = await hybrid_retrieve(query, vector_store, bm25_index, top_k=5)
        answer = await generate_answer(query, chunks)
        return {"answer": answer, "sources": [c.source for c in chunks], "path": "simple"}
 
    # Step 2: Decompose complex queries
    sub_queries = await decompose_query(query)
 
    # Step 3: Retrieve for each sub-query
    all_chunks: list[RetrievedChunk] = []
    for sub_q in sub_queries:
        chunks = await hybrid_retrieve(sub_q, vector_store, bm25_index, top_k=5)
        all_chunks.extend(chunks)
 
    # Deduplicate by content
    seen = set()
    unique_chunks = []
    for chunk in all_chunks:
        if chunk.content not in seen:
            seen.add(chunk.content)
            unique_chunks.append(chunk)
 
    # Step 4: Judge retrieval quality (with retry loop)
    current_query = query
    for attempt in range(max_retrieval_attempts):
        relevant_chunks, judgment = await judge_retrieval(current_query, unique_chunks)
 
        if judgment.is_sufficient:
            break
 
        if judgment.suggested_requery and attempt < max_retrieval_attempts - 1:
            # Re-retrieve with the judge's suggested query
            new_chunks = await hybrid_retrieve(
                judgment.suggested_requery, vector_store, bm25_index, top_k=5
            )
            unique_chunks.extend(new_chunks)
            current_query = judgment.suggested_requery
    else:
        # Exhausted retries — generate with best available context
        relevant_chunks = unique_chunks[:10]
 
    # Step 5: Generate answer
    answer = await generate_answer(query, relevant_chunks)
 
    # Step 6: Grounding check
    grounding = await check_grounding(answer, relevant_chunks, query)
 
    if not grounding.is_grounded and grounding.unsupported_claims:
        # Regenerate with explicit grounding instruction
        answer = await generate_answer(
            query,
            relevant_chunks,
            system_suffix=(
                "\n\nIMPORTANT: Only state facts directly supported by the "
                "provided context. Do not infer or extrapolate. If the context "
                "does not contain enough information, say so explicitly."
            )
        )
 
    return {
        "answer": answer,
        "sources": list(set(c.source for c in relevant_chunks)),
        "grounding_confidence": grounding.confidence,
        "retrieval_attempts": attempt + 1,
        "path": "complex"
    }
 
 
async def generate_answer(
    query: str,
    chunks: list[RetrievedChunk],
    system_suffix: str = ""
) -> str:
    """Generate an answer from retrieved chunks."""
    context = "\n\n".join(
        f"[Source: {c.source}]\n{c.content}" for c in chunks
    )
    response = client.messages.create(
        model="claude-sonnet-4-6-20250514",
        max_tokens=4096,
        system=(
            "Answer the user's question based only on the provided context. "
            "Cite sources when possible. If the context doesn't contain "
            "enough information, say so." + system_suffix
        ),
        messages=[{
            "role": "user",
            "content": f"Context:\n{context}\n\nQuestion: {query}"
        }]
    )
    return response.content[0].text

@dataclass
class EvalCase:
    query: str
    expected_sources: list[str]  # doc IDs that should be retrieved
    expected_answer_contains: list[str]  # key facts the answer must include
 
 
async def evaluate_pipeline(
    eval_set: list[EvalCase],
    pipeline_fn,
    vector_store,
    bm25_index
) -> dict:
    """Evaluate a RAG pipeline against a ground-truth eval set."""
    metrics = {
        "retrieval_recall": [],      # did we find the right docs?
        "answer_completeness": [],   # did the answer include key facts?
        "hallucination_rate": [],    # did the answer include wrong facts?
        "latency_ms": [],
        "cost_usd": [],
    }
 
    for case in eval_set:
        start = time.time()
        result = await pipeline_fn(case.query, vector_store, bm25_index)
        elapsed = (time.time() - start) * 1000
 
        # Retrieval recall
        retrieved_sources = set(result["sources"])
        expected = set(case.expected_sources)
        recall = len(retrieved_sources & expected) / len(expected) if expected else 1.0
        metrics["retrieval_recall"].append(recall)
 
        # Answer completeness
        answer_lower = result["answer"].lower()
        found = sum(1 for fact in case.expected_answer_contains if fact.lower() in answer_lower)
        completeness = found / len(case.expected_answer_contains) if case.expected_answer_contains else 1.0
        metrics["answer_completeness"].append(completeness)
 
        metrics["latency_ms"].append(elapsed)
 
    return {k: sum(v) / len(v) for k, v in metrics.items() if v}