Ch 3 — Advanced RAG¶

"""
hyde.py — Hypothetical Document Embeddings for improved recall.
"""

from __future__ import annotations

import anthropic
import numpy as np
from sentence_transformers import SentenceTransformer


def generate_hypothetical_document(
    query: str,
    model: str = "claude-haiku-4-5",
) -> str:
    """Generate a plausible (not necessarily accurate) answer to the query."""
    client = anthropic.Anthropic()
    message = client.messages.create(
        model=model,
        max_tokens=256,
        system=(
            "Write a concise, factual-sounding passage that directly answers "
            "the question. The passage should look like an excerpt from a "
            "reference document. Do not hedge or say you are generating a hypothetical."
        ),
        messages=[{"role": "user", "content": query}],
    )
    return message.content[0].text  # type: ignore[union-attr]


def hyde_embed(
    query: str,
    embed_model: SentenceTransformer,
    n_hypothetical: int = 3,
) -> np.ndarray:
    """
    Return a HyDE embedding by averaging embeddings of N hypothetical documents.

    Averaging multiple hypothetical documents reduces variance from any single
    hallucinated answer.
    """
    hypotheticals = [generate_hypothetical_document(query) for _ in range(n_hypothetical)]
    embeddings: np.ndarray = embed_model.encode(
        hypotheticals, normalize_embeddings=True, convert_to_numpy=True
    )
    mean_emb = embeddings.mean(axis=0)
    # Re-normalise the averaged embedding
    mean_emb = mean_emb / np.linalg.norm(mean_emb)
    return mean_emb

def expand_query(query: str) -> list[str]:
    """Generate N expanded versions of the query for multi-query retrieval."""
    client = anthropic.Anthropic()
    message = client.messages.create(
        model="claude-haiku-4-5",
        max_tokens=256,
        system=(
            "You are a search query optimiser. Given a query, generate 3 alternative "
            "phrasings that cover different aspects or use different vocabulary. "
            "Output one query per line, no numbering."
        ),
        messages=[{"role": "user", "content": f"Original query: {query}"}],
    )
    lines = message.content[0].text.strip().split("\n")  # type: ignore[union-attr]
    return [query] + [line.strip() for line in lines if line.strip()]

"""
reranker.py — Cross-encoder re-ranking of retrieved candidates.

Install: pip install sentence-transformers
"""

from __future__ import annotations

from sentence_transformers import CrossEncoder


def rerank(
    query: str,
    candidates: list[str],
    model_name: str = "cross-encoder/ms-marco-MiniLM-L-6-v2",
    top_k: int = 5,
) -> list[tuple[float, str]]:
    """
    Re-rank candidate passages using a cross-encoder model.

    Args:
        query: The user's original query.
        candidates: List of retrieved passages to re-rank.
        model_name: Cross-encoder model identifier.
        top_k: Number of top passages to return.

    Returns:
        List of (score, passage) tuples sorted by score descending.
    """
    model = CrossEncoder(model_name)
    pairs = [(query, passage) for passage in candidates]
    scores: list[float] = model.predict(pairs).tolist()

    ranked = sorted(zip(scores, candidates), key=lambda x: x[0], reverse=True)
    return ranked[:top_k]

"""
multi_query.py — Multi-query retrieval with union and deduplication.
"""

from __future__ import annotations

import anthropic
from sentence_transformers import SentenceTransformer


def multi_query_retrieve(
    query: str,
    index,  # VectorIndex from Ch 1
    embed_model: SentenceTransformer,
    n_queries: int = 4,
    k_per_query: int = 5,
) -> list[str]:
    """
    Generate N paraphrased queries, retrieve k results per query,
    return deduplicated union.

    Returns:
        List of unique retrieved passages, deduplicated by exact text match.
    """
    client = anthropic.Anthropic()
    message = client.messages.create(
        model="claude-haiku-4-5",
        max_tokens=512,
        system=(
            "Generate exactly {n} alternative phrasings of the user's question. "
            "Each phrasing should approach the topic from a different angle. "
            "Output one phrasing per line, no numbering, no preamble."
        ).format(n=n_queries - 1),
        messages=[{"role": "user", "content": query}],
    )
    extra_queries = [
        line.strip()
        for line in message.content[0].text.strip().split("\n")  # type: ignore[union-attr]
        if line.strip()
    ]
    all_queries = [query] + extra_queries[:n_queries - 1]

    seen: set[str] = set()
    unique_passages: list[str] = []

    for q in all_queries:
        results = index.search(q, k=k_per_query)
        for _score, passage in results:
            if passage not in seen:
                seen.add(passage)
                unique_passages.append(passage)

    return unique_passages

"""
parent_child.py — Parent-child chunk indexing and retrieval.
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np
from sentence_transformers import SentenceTransformer


@dataclass
class ParentChunk:
    parent_id: str
    parent_text: str
    children: list[str]


def build_parent_child_index(
    paragraphs: list[str],
    embed_model: SentenceTransformer,
    child_sentences: int = 2,
) -> tuple[np.ndarray, list[str], list[str]]:
    """
    Split each paragraph into child sentences, embed children for retrieval,
    map each child back to its parent paragraph.

    Returns:
        (child_embeddings, child_texts, parent_texts_per_child)
    """
    import nltk
    nltk.download("punkt_tab", quiet=True)

    child_texts: list[str] = []
    parent_texts: list[str] = []

    for para in paragraphs:
        sentences = nltk.sent_tokenize(para)
        for i in range(0, len(sentences), child_sentences):
            child = " ".join(sentences[i : i + child_sentences]).strip()
            if child:
                child_texts.append(child)
                parent_texts.append(para)  # map child → its parent paragraph

    embeddings: np.ndarray = embed_model.encode(
        child_texts, normalize_embeddings=True, convert_to_numpy=True
    )
    return embeddings, child_texts, parent_texts

def compress_context(
    query: str,
    passage: str,
    model: str = "claude-haiku-4-5",
) -> str:
    """Extract only the sentences from `passage` that are relevant to `query`."""
    client = anthropic.Anthropic()
    message = client.messages.create(
        model=model,
        max_tokens=512,
        system=(
            "You are a context compressor. Given a passage and a query, extract "
            "ONLY the sentences from the passage that are directly relevant to "
            "the query. Preserve the exact wording. If no sentences are relevant, "
            "return the empty string."
        ),
        messages=[{
            "role": "user",
            "content": f"Query: {query}\n\nPassage:\n{passage}",
        }],
    )
    compressed = message.content[0].text.strip()  # type: ignore[union-attr]
    return compressed if compressed else passage  # fall back to full passage

"""
ragas_eval.py — Evaluate a RAG pipeline with RAGAS.

Install: pip install ragas anthropic datasets
"""

from __future__ import annotations

from datasets import Dataset
from ragas import evaluate
from ragas.metrics import (
    answer_relevancy,
    context_precision,
    context_recall,
    faithfulness,
)

# ---------------------------------------------------------------------------
# Build the evaluation dataset
# ---------------------------------------------------------------------------
# Each row: question, answer, contexts (list[str]), ground_truth
eval_data = {
    "question": [
        "What is HNSW and how does its layered graph work?",
        "When should you use DPO instead of PPO?",
    ],
    "answer": [
        "HNSW is a graph-based ANN algorithm that builds multiple layers of proximity graphs.",
        "DPO is preferred when you have a fixed preference dataset and want simpler training.",
    ],
    "contexts": [
        [
            "HNSW (Hierarchical Navigable Small World) builds a multi-layer proximity graph...",
            "Higher layers contain fewer nodes with longer-range connections for fast traversal...",
        ],
        [
            "DPO eliminates the separate reward model and RL loop by deriving a supervised loss...",
            "PPO is preferred for online alignment where responses are generated during training...",
        ],
    ],
    "ground_truths": [
        ["HNSW builds a layered graph where higher layers provide long-range navigation and layer 0 provides dense local connections."],
        ["DPO is better when you have a fixed preference dataset; PPO is better for online alignment."],
    ],
}

dataset = Dataset.from_dict(eval_data)

# ---------------------------------------------------------------------------
# Run evaluation
# ---------------------------------------------------------------------------
result = evaluate(
    dataset=dataset,
    metrics=[faithfulness, answer_relevancy, context_recall, context_precision],
    # Specify your LLM judge here if not using the default OpenAI
    # llm=your_llm_wrapper,
)

print(result)
# Output: {'faithfulness': 0.92, 'answer_relevancy': 0.88,
#          'context_recall': 0.85, 'context_precision': 0.79}

"""
agentic_rag.py — RAG where the LLM decides when and what to retrieve.
"""

from __future__ import annotations

import json

import anthropic

# The retrieval tool the LLM can call
RETRIEVAL_TOOL = {
    "name": "retrieve_documents",
    "description": (
        "Search the knowledge base for passages relevant to a query. "
        "Call this when you need information to answer the user's question. "
        "You may call it multiple times with different queries."
    ),
    "input_schema": {
        "type": "object",
        "properties": {
            "query": {
                "type": "string",
                "description": "The search query to find relevant passages.",
            },
            "k": {
                "type": "integer",
                "description": "Number of passages to retrieve (default 5, max 10).",
                "default": 5,
            },
        },
        "required": ["query"],
    },
}


def agentic_rag(
    user_question: str,
    vector_index,  # VectorIndex from Ch 1
    max_iterations: int = 4,
) -> str:
    """
    Run an agentic RAG loop where the LLM decides when to retrieve.

    Args:
        user_question: The question to answer.
        vector_index: Prebuilt VectorIndex instance.
        max_iterations: Maximum retrieve-evaluate cycles before forcing generation.

    Returns:
        The LLM's final answer with inline citations.
    """
    client = anthropic.Anthropic()
    messages: list[dict] = [{"role": "user", "content": user_question}]
    all_retrieved: list[str] = []

    system = (
        "You are a research assistant with access to a knowledge base. "
        "Use the retrieve_documents tool to gather information before answering. "
        "You may retrieve multiple times if the initial results are insufficient. "
        "When you have enough information, provide a comprehensive answer with citations."
    )

    for _ in range(max_iterations):
        response = client.messages.create(
            model="claude-opus-4-5",
            max_tokens=1024,
            system=system,
            tools=[RETRIEVAL_TOOL],  # type: ignore[list-item]
            messages=messages,
        )

        if response.stop_reason == "end_turn":
            # LLM decided it has enough information
            return response.content[0].text  # type: ignore[union-attr]

        # Process tool calls
        tool_results = []
        for block in response.content:
            if block.type == "tool_use" and block.name == "retrieve_documents":
                inputs = block.input
                k = min(int(inputs.get("k", 5)), 10)
                results = vector_index.search(inputs["query"], k=k)
                passages = [chunk for _score, chunk in results]
                all_retrieved.extend(passages)

                context = "\n\n---\n\n".join(
                    f"[{i + 1}] {p}" for i, p in enumerate(passages)
                )
                tool_results.append({
                    "type": "tool_result",
                    "tool_use_id": block.id,
                    "content": context,
                })

        messages.append({"role": "assistant", "content": response.content})
        messages.append({"role": "user", "content": tool_results})

    # Force final answer after max iterations
    messages.append({
        "role": "user",
        "content": "Please provide your best answer based on what you have retrieved so far.",
    })
    final = client.messages.create(
        model="claude-opus-4-5",
        max_tokens=1024,
        system=system,
        messages=messages,
    )
    return final.content[0].text  # type: ignore[union-attr]