Retrieval-Augmented Generation for Agents

from pathlib import Path


def load_text_file(path: str) -> str:
    """Load a plain text or markdown file."""
    return Path(path).read_text(encoding="utf-8")


def load_documents(directory: str, extensions: tuple = (".txt", ".md")) -> list[dict]:
    """Load all documents from a directory.

    This simple loader handles text and markdown files. In production,
    you would use specialized loaders for different formats:
    - PDF: PyPDF2, pdfplumber, or unstructured
    - DOCX: python-docx
    - HTML: BeautifulSoup
    - Code: tree-sitter for syntax-aware parsing
    """
    documents = []
    for path in Path(directory).rglob("*"):
        if path.suffix in extensions:
            documents.append({
                "content": path.read_text(encoding="utf-8"),
                "source": str(path),
                "filename": path.name,
            })
    return documents

def fixed_size_chunks(
    text: str, chunk_size: int = 500, overlap: int = 50
) -> list[str]:
    """Split text into fixed-size chunks with overlap.

    The overlap parameter is important: without it, a sentence that
    happens to fall right at a chunk boundary would be split in half,
    with the first part in one chunk and the second in the next.
    Overlap ensures that boundary sentences appear in both chunks,
    so at least one chunk contains the complete sentence.

    Args:
        text: The input text to chunk.
        chunk_size: Maximum number of characters per chunk.
        overlap: Number of characters to overlap between chunks.

    Returns:
        List of text chunks.
    """
    chunks = []
    start = 0
    while start < len(text):
        end = start + chunk_size
        chunk = text[start:end]
        chunks.append(chunk.strip())
        start = end - overlap
    return [c for c in chunks if c]  # Remove empty chunks

import re


def sentence_based_chunks(
    text: str, max_chunk_size: int = 500
) -> list[str]:
    """Split text into chunks at sentence boundaries.

    This preserves sentence integrity -- no sentence is ever split
    across two chunks. Sentences are grouped together until the
    chunk reaches the target size, then a new chunk begins.

    The trade-off: chunk sizes are variable. Some chunks may be
    much shorter than max_chunk_size (if a single sentence is very
    long), while others may approach it closely.
    """
    # Simple sentence splitting (production systems use spaCy or nltk)
    sentences = re.split(r'(?<=[.!?])\s+', text)

    chunks = []
    current_chunk = []
    current_size = 0

    for sentence in sentences:
        sentence_size = len(sentence)
        if current_size + sentence_size > max_chunk_size and current_chunk:
            chunks.append(" ".join(current_chunk))
            current_chunk = []
            current_size = 0
        current_chunk.append(sentence)
        current_size += sentence_size

    if current_chunk:
        chunks.append(" ".join(current_chunk))

    return chunks

from sentence_transformers import SentenceTransformer


def generate_embeddings(
    chunks: list[str], model_name: str = "all-MiniLM-L6-v2"
) -> list[list[float]]:
    """Generate embeddings for a list of text chunks.

    Each chunk is converted into a fixed-size vector (e.g., 384
    dimensions for MiniLM) that captures its semantic meaning.
    Similar texts get similar vectors, enabling similarity search.

    Args:
        chunks: List of text strings to embed.
        model_name: Name of the sentence-transformers model to use.

    Returns:
        List of embedding vectors.
    """
    model = SentenceTransformer(model_name)
    embeddings = model.encode(chunks, show_progress_bar=True)
    return embeddings.tolist()

import chromadb


def build_index(
    chunks: list[str],
    sources: list[str],
    collection_name: str = "documents",
) -> chromadb.Collection:
    """Build a vector index from text chunks.

    ChromaDB handles embedding generation internally using its
    default model, so we only need to provide the text chunks.
    The 'cosine' distance metric is used for similarity search,
    which measures the angle between vectors (direction similarity)
    rather than their absolute distance.

    Args:
        chunks: List of text chunks.
        sources: List of source identifiers (one per chunk).
        collection_name: Name for the vector collection.

    Returns:
        A ChromaDB collection with the indexed documents.
    """
    client = chromadb.Client()
    collection = client.create_collection(
        name=collection_name,
        metadata={"hnsw:space": "cosine"},  # Use cosine similarity
    )

    # ChromaDB generates embeddings automatically using its default model
    collection.add(
        documents=chunks,
        metadatas=[{"source": src} for src in sources],
        ids=[f"chunk_{i}" for i in range(len(chunks))],
    )

    return collection

import numpy as np


def cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
    """Compute cosine similarity between two vectors.

    The formula: cos(theta) = (A . B) / (|A| * |B|)

    Where A . B is the dot product (sum of element-wise products)
    and |A| is the magnitude (Euclidean norm) of vector A.
    """
    return float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-8))

# Example: FAISS with IVF + PQ for large-scale search
import faiss
import numpy as np


def build_faiss_index(
    embeddings: np.ndarray, nlist: int = 100, m: int = 8
) -> faiss.Index:
    """Build a FAISS index with IVF and Product Quantization.

    This combines two techniques:
    - IVF (Inverted File): Clusters vectors into nlist groups.
      At query time, only the nearest clusters are searched.
    - PQ (Product Quantization): Compresses each vector into m
      sub-vectors, reducing memory by ~16x.

    Together, they enable searching millions of vectors in
    milliseconds using modest hardware.

    Args:
        embeddings: Matrix of embeddings (n_docs x dim).
        nlist: Number of clusters for IVF.
        m: Number of sub-quantizers for PQ.

    Returns:
        Trained FAISS index.
    """
    dim = embeddings.shape[1]

    # Create the index: IVF with PQ compression
    quantizer = faiss.IndexFlatL2(dim)
    index = faiss.IndexIVFPQ(quantizer, dim, nlist, m, 8)

    # Train the index on the data -- IVF needs to learn cluster
    # centers, and PQ needs to learn quantization codebooks
    index.train(embeddings)
    index.add(embeddings)

    return index


def search_faiss(
    index: faiss.Index, query_embedding: np.ndarray, top_k: int = 5
) -> tuple[np.ndarray, np.ndarray]:
    """Search the FAISS index for nearest neighbors.

    Returns:
        Tuple of (distances, indices) arrays.
    """
    query = query_embedding.reshape(1, -1)
    distances, indices = index.search(query, top_k)
    return distances[0], indices[0]

import math
from collections import Counter


class BM25:
    """A simple BM25 implementation for sparse retrieval.

    BM25 (Best Matching 25) ranks documents based on:
    1. Term Frequency (TF): How often does the query term appear
       in this document? More occurrences = more relevant.
    2. Inverse Document Frequency (IDF): How rare is this term
       across all documents? Rare terms are more informative.
    3. Document Length Normalization: Longer documents naturally
       contain more terms, so we normalize for length.

    The name "BM25" comes from it being the 25th in a series of
    ranking functions explored by Robertson et al. in the 1990s.

    Parameters:
    - k1 controls term frequency saturation. Higher k1 means
      additional term occurrences have more impact. Default: 1.5.
    - b controls document length normalization. b=1 means full
      normalization, b=0 means no normalization. Default: 0.75.
    """

    def __init__(self, documents: list[str], k1: float = 1.5, b: float = 0.75):
        self.k1 = k1
        self.b = b
        self.documents = documents
        self.doc_count = len(documents)

        # Tokenize and compute statistics
        self.doc_tokens = [doc.lower().split() for doc in documents]
        self.doc_lengths = [len(tokens) for tokens in self.doc_tokens]
        self.avg_doc_length = sum(self.doc_lengths) / self.doc_count

        # Compute document frequency for each term
        self.df = {}
        for tokens in self.doc_tokens:
            for token in set(tokens):  # set() to count each term once per doc
                self.df[token] = self.df.get(token, 0) + 1

    def _idf(self, term: str) -> float:
        """Compute inverse document frequency for a term.

        IDF is high for rare terms and low for common terms.
        The formula includes smoothing (+0.5) to avoid division
        by zero and the +1 to ensure IDF is always positive.
        """
        df = self.df.get(term, 0)
        return math.log((self.doc_count - df + 0.5) / (df + 0.5) + 1)

    def score(self, query: str, doc_index: int) -> float:
        """Compute BM25 score for a query-document pair.

        The score is the sum of each query term's contribution:
        IDF * (TF * (k1 + 1)) / (TF + k1 * (1 - b + b * dl/avgdl))

        where TF is term frequency, dl is document length, and
        avgdl is average document length.
        """
        query_tokens = query.lower().split()
        doc_tokens = self.doc_tokens[doc_index]
        doc_length = self.doc_lengths[doc_index]
        tf_counts = Counter(doc_tokens)

        score = 0.0
        for term in query_tokens:
            if term not in tf_counts:
                continue
            tf = tf_counts[term]
            idf = self._idf(term)
            numerator = tf * (self.k1 + 1)
            denominator = tf + self.k1 * (
                1 - self.b + self.b * doc_length / self.avg_doc_length
            )
            score += idf * numerator / denominator

        return score

    def search(self, query: str, top_k: int = 5) -> list[tuple[int, float]]:
        """Search for the most relevant documents.

        Returns:
            List of (document_index, score) tuples sorted by score descending.
        """
        scores = [(i, self.score(query, i)) for i in range(self.doc_count)]
        scores.sort(key=lambda x: x[1], reverse=True)
        return scores[:top_k]

def reciprocal_rank_fusion(
    rankings: list[list[tuple[str, float]]],
    k: int = 60,
) -> list[tuple[str, float]]:
    """Combine multiple rankings using Reciprocal Rank Fusion.

    RRF is elegant in its simplicity: it assigns a score of
    1/(k + rank) to each document in each ranking, then sums
    the scores across rankings.

    The k parameter (default 60) is a smoothing constant that
    prevents top-ranked documents from dominating. Without it,
    the #1 result from one system would always beat the #2 result
    from both systems combined.

    Example: If document X is ranked #1 by dense retrieval and
    #3 by sparse retrieval:
    - Dense score: 1/(60+1) = 0.0164
    - Sparse score: 1/(60+3) = 0.0159
    - Total: 0.0323

    If document Y is ranked #2 by both:
    - Dense score: 1/(60+2) = 0.0161
    - Sparse score: 1/(60+2) = 0.0161
    - Total: 0.0323

    Both documents score equally, which makes intuitive sense:
    being good in both systems is as valuable as being great
    in one and decent in the other.

    Args:
        rankings: List of ranked lists, each containing (doc_id, score) tuples.
        k: Constant to prevent high scores for top-ranked documents (default: 60).

    Returns:
        Fused ranking as a list of (doc_id, fused_score) tuples.
    """
    fused_scores: dict[str, float] = {}

    for ranking in rankings:
        for rank, (doc_id, _) in enumerate(ranking):
            if doc_id not in fused_scores:
                fused_scores[doc_id] = 0.0
            fused_scores[doc_id] += 1.0 / (k + rank + 1)

    # Sort by fused score descending
    fused = sorted(fused_scores.items(), key=lambda x: x[1], reverse=True)
    return fused

from sentence_transformers import CrossEncoder


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

    The cross-encoder processes each (query, document) pair through
    a transformer model, producing a relevance score. This is
    significantly more accurate than bi-encoder similarity because
    the model can attend to interactions between query terms and
    document terms.

    Typical pipeline:
    1. Initial retrieval: Get 20-50 candidates (fast, approximate)
    2. Re-ranking: Score all candidates with cross-encoder (slow, precise)
    3. Return top-k after re-ranking

    Args:
        query: The search query.
        documents: List of candidate documents from initial retrieval.
        model_name: Cross-encoder model to use.
        top_k: Number of top documents to return.

    Returns:
        List of (document, score) tuples sorted by relevance.
    """
    model = CrossEncoder(model_name)
    pairs = [(query, doc) for doc in documents]
    scores = model.predict(pairs)

    scored_docs = list(zip(documents, scores))
    scored_docs.sort(key=lambda x: x[1], reverse=True)
    return scored_docs[:top_k]

def hyde_retrieval(query: str, llm_call, retriever) -> list[str]:
    """Hypothetical Document Embedding retrieval.

    1. Ask the LLM to generate a hypothetical answer.
    2. Use the hypothetical answer's embedding for retrieval.
    3. Return the retrieved real documents.

    The hypothetical answer does not need to be correct -- it just
    needs to be in the right semantic neighborhood. Even an
    inaccurate hypothesis will use the right terminology and
    concepts, making it a better retrieval query than the original
    question.

    Reference: Gao et al., 2023 - "Precise Zero-Shot Dense Retrieval
    without Relevance Labels" (ACL 2023).
    """
    # Step 1: Generate hypothetical answer
    hypothesis = llm_call(
        prompt=f"Write a short paragraph that would answer this question: {query}"
    )

    # Step 2: Use the hypothetical answer for retrieval
    results = retriever.search(hypothesis, top_k=5)

    return results

DECOMPOSE_PROMPT = """Break the following complex question into 2-4 simpler
sub-questions that, when answered together, would answer the original question.

Original question: {query}

Sub-questions:
1."""


def decompose_and_retrieve(query: str, llm_call, retriever) -> list[str]:
    """Decompose a complex query and retrieve for each sub-query.

    Example: "Compare the performance and cost of GPT-4 and Claude 3.5"
    might decompose into:
    1. "What is the performance of GPT-4 on standard benchmarks?"
    2. "What is the performance of Claude 3.5 on standard benchmarks?"
    3. "What is the pricing of GPT-4?"
    4. "What is the pricing of Claude 3.5?"

    Each sub-query retrieves different documents, and together they
    provide the information needed to answer the original question.
    """
    # Generate sub-questions
    response = llm_call(prompt=DECOMPOSE_PROMPT.format(query=query))
    sub_queries = [q.strip() for q in response.split("\n") if q.strip()]

    # Retrieve for each sub-query, deduplicating results
    all_results = []
    seen = set()
    for sub_query in sub_queries:
        results = retriever.search(sub_query, top_k=3)
        for doc in results:
            if doc not in seen:
                all_results.append(doc)
                seen.add(doc)

    return all_results

def iterative_retrieval(
    query: str, llm_call, retriever, max_steps: int = 3
) -> list[str]:
    """Iteratively retrieve and refine.

    At each step, the LLM examines current results and generates
    a follow-up query to fill in gaps. This mimics how a human
    researcher works: you read initial sources, identify gaps in
    your understanding, then search for more information to fill
    those gaps.

    The process terminates when the LLM determines it has enough
    information or when the maximum number of steps is reached.
    """
    all_retrieved = []

    current_query = query
    for step in range(max_steps):
        # Retrieve
        results = retriever.search(current_query, top_k=3)
        all_retrieved.extend(results)

        # Ask LLM if we have enough information
        context = "\n".join(all_retrieved)
        follow_up = llm_call(
            prompt=(
                f"Original question: {query}\n\n"
                f"Information retrieved so far:\n{context}\n\n"
                "Is this enough to answer the question? If not, what specific "
                "information is still missing? Generate a follow-up search query "
                "to find the missing information. If enough, respond with 'SUFFICIENT'."
            )
        )

        if "SUFFICIENT" in follow_up.upper():
            break
        current_query = follow_up

    return all_retrieved

AGENT_SYSTEM_PROMPT = """You are a helpful assistant with access to a knowledge base.
You have the following tools available:

1. search_knowledge_base(query: str) -> list[str]
   Search the knowledge base for relevant information.

2. answer(response: str) -> None
   Provide your final answer to the user.

Guidelines:
- Only search the knowledge base when you need information you don't
  already know or when accuracy is critical.
- You may search multiple times with different queries.
- Always cite the source when using retrieved information.
- If the knowledge base doesn't have relevant information, say so
  and answer based on your general knowledge (with a caveat).

Think step by step about whether you need to search before answering.
"""


class AgenticRAG:
    """An agent that decides when and what to retrieve.

    Unlike standard RAG (which always retrieves), this agent
    uses retrieval as a tool -- it decides whether to search,
    what to search for, and when it has enough information to
    answer. This is more efficient and often more accurate.
    """

    def __init__(self, llm_call, retriever):
        self.llm_call = llm_call
        self.retriever = retriever
        self.retrieved_context = []

    def search_knowledge_base(self, query: str) -> list[str]:
        """Tool: search the knowledge base."""
        results = self.retriever.search(query, top_k=3)
        self.retrieved_context.extend(results)
        return results

    def process_query(self, user_query: str) -> str:
        """Process a user query with optional retrieval.

        The agent reasons about whether it needs to search,
        formulates search queries if needed, and synthesizes
        the results into a final answer.
        """
        messages = [
            {"role": "system", "content": AGENT_SYSTEM_PROMPT},
            {"role": "user", "content": user_query},
        ]

        # Agent reasoning loop (max 5 iterations to prevent runaway)
        for _ in range(5):
            response = self.llm_call(messages=messages)

            if "search_knowledge_base" in response:
                # Agent decided to search -- extract query and execute
                search_query = self._extract_search_query(response)
                results = self.search_knowledge_base(search_query)
                messages.append({"role": "assistant", "content": response})
                messages.append({
                    "role": "tool",
                    "content": f"Search results:\n" + "\n---\n".join(results),
                })
            elif "answer(" in response:
                # Agent is ready to give a final answer
                return self._extract_answer(response)
            else:
                return response

        return "I was unable to complete the query within the step limit."

    def _extract_search_query(self, response: str) -> str:
        """Extract the search query from an agent's tool call."""
        start = response.index("search_knowledge_base(") + len("search_knowledge_base(")
        end = response.index(")", start)
        return response[start:end].strip("\"'")

    def _extract_answer(self, response: str) -> str:
        """Extract the answer from an agent's tool call."""
        start = response.index("answer(") + len("answer(")
        end = response.rindex(")")
        return response[start:end].strip("\"'")

ROUTER_PROMPT = """Given the user's question, decide which knowledge source to query.
Available sources:
- technical_docs: API documentation, code references, technical specifications
- company_wiki: Company policies, procedures, organizational information
- research_papers: Academic papers and research findings
- none: No retrieval needed (use general knowledge)

Question: {query}

Respond with ONLY the source name."""


def route_query(query: str, llm_call) -> str:
    """Route a query to the appropriate knowledge source.

    This is like a librarian who directs you to the right section
    of the library based on your question, rather than making you
    search the entire library every time.
    """
    response = llm_call(prompt=ROUTER_PROMPT.format(query=query))
    return response.strip().lower()

Input: "What causes aurora borealis?"

Step 1: [Retrieve: Yes]  -- Model decides retrieval is needed
Step 2: Retrieved: "Aurora borealis occurs when charged particles from
        the Sun interact with Earth's magnetic field..."
Step 3: [Relevant: Yes]  -- Model confirms passage is relevant
Step 4: "Aurora borealis is caused by charged particles from the Sun
        colliding with gases in Earth's atmosphere..."
Step 5: [Supported: Fully]  -- Model confirms response is grounded
Step 6: [Useful: 5]  -- Model rates utility highly

def recall_at_k(retrieved_ids: list[str], relevant_ids: set[str], k: int) -> float:
    """Compute recall at k: fraction of relevant docs found in top k results.

    Example: If there are 4 relevant documents and 3 of them appear
    in the top 10 results, recall@10 = 3/4 = 0.75.
    """
    top_k = set(retrieved_ids[:k])
    return len(top_k & relevant_ids) / len(relevant_ids) if relevant_ids else 0.0


def precision_at_k(retrieved_ids: list[str], relevant_ids: set[str], k: int) -> float:
    """Compute precision at k: fraction of top k results that are relevant.

    Example: If 3 out of 10 top results are relevant,
    precision@10 = 3/10 = 0.3.
    """
    top_k = set(retrieved_ids[:k])
    return len(top_k & relevant_ids) / k if k > 0 else 0.0


def mrr(retrieved_ids: list[str], relevant_ids: set[str]) -> float:
    """Compute Mean Reciprocal Rank.

    If the first relevant result is at position 3,
    MRR = 1/3 = 0.333. A higher MRR means relevant results
    appear earlier in the ranking.
    """
    for i, doc_id in enumerate(retrieved_ids):
        if doc_id in relevant_ids:
            return 1.0 / (i + 1)
    return 0.0

def evaluate_faithfulness(
    question: str, answer: str, context: str, llm_call
) -> float:
    """Evaluate faithfulness of an answer to its context.

    Simplified version of the RAGAS faithfulness metric.

    The process:
    1. Extract all factual claims from the answer.
    2. Check each claim against the retrieved context.
    3. Faithfulness = (supported claims) / (total claims).

    A faithfulness of 1.0 means every claim in the answer is
    supported by the context. A faithfulness of 0.5 means half
    the claims are unsupported (hallucinated).
    """
    # Step 1: Extract claims from the answer
    claims_response = llm_call(
        prompt=f"List all factual claims made in this answer:\n\n{answer}\n\nClaims:"
    )
    claims = [c.strip() for c in claims_response.split("\n") if c.strip()]

    if not claims:
        return 1.0  # No claims to verify

    # Step 2: Check each claim against the context
    supported = 0
    for claim in claims:
        verdict = llm_call(
            prompt=(
                f"Is the following claim supported by the context?\n\n"
                f"Claim: {claim}\n\n"
                f"Context: {context}\n\n"
                f"Answer only 'Yes' or 'No'."
            )
        )
        if "yes" in verdict.lower():
            supported += 1

    return supported / len(claims)

"""
A complete RAG pipeline implementation.

This example demonstrates building a RAG system from document loading
through retrieval and generation. Each component is implemented
transparently so you can see exactly what happens at each step.

Requirements:
    pip install sentence-transformers chromadb openai
"""

import os
from pathlib import Path

import chromadb
from sentence_transformers import SentenceTransformer, CrossEncoder


class SimpleRAGPipeline:
    """A complete RAG pipeline with indexing, retrieval, and generation.

    This pipeline supports:
    - Document loading and chunking
    - Embedding generation and vector storage
    - Dense retrieval with optional re-ranking
    - Context-augmented generation

    The pipeline follows the standard RAG architecture:
    1. OFFLINE: Documents -> Chunks -> Embeddings -> Vector Store
    2. ONLINE:  Query -> Retrieve -> Augment Prompt -> Generate Answer
    """

    def __init__(
        self,
        embedding_model: str = "all-MiniLM-L6-v2",
        reranker_model: str = "cross-encoder/ms-marco-MiniLM-L-6-v2",
        chunk_size: int = 500,
        chunk_overlap: int = 50,
    ):
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap

        # Initialize models
        self.embedder = SentenceTransformer(embedding_model)
        self.reranker = CrossEncoder(reranker_model)

        # Initialize vector store (in-memory for this example)
        self.client = chromadb.Client()
        self.collection = self.client.create_collection(
            name="rag_docs",
            metadata={"hnsw:space": "cosine"},
        )
        self.doc_count = 0

    def chunk_text(self, text: str, source: str) -> list[dict]:
        """Split text into overlapping chunks with metadata."""
        chunks = []
        start = 0
        chunk_idx = 0
        while start < len(text):
            end = start + self.chunk_size
            chunk_text = text[start:end].strip()
            if chunk_text:
                chunks.append({
                    "text": chunk_text,
                    "source": source,
                    "chunk_index": chunk_idx,
                })
                chunk_idx += 1
            start = end - self.chunk_overlap
        return chunks

    def index_documents(self, documents: list[dict]) -> int:
        """Index a list of documents.

        This is the OFFLINE phase: documents are chunked, embedded,
        and stored in the vector database. This only needs to happen
        once (or when documents are updated).

        Args:
            documents: List of dicts with 'content' and 'source' keys.

        Returns:
            Number of chunks indexed.
        """
        all_chunks = []
        for doc in documents:
            chunks = self.chunk_text(doc["content"], doc["source"])
            all_chunks.extend(chunks)

        if not all_chunks:
            return 0

        # Add to vector store (ChromaDB handles embedding internally)
        self.collection.add(
            documents=[c["text"] for c in all_chunks],
            metadatas=[{"source": c["source"], "chunk_index": c["chunk_index"]} for c in all_chunks],
            ids=[f"doc_{self.doc_count + i}" for i in range(len(all_chunks))],
        )
        self.doc_count += len(all_chunks)
        return len(all_chunks)

    def retrieve(
        self,
        query: str,
        top_k: int = 5,
        use_reranking: bool = True,
        initial_k: int = 20,
    ) -> list[dict]:
        """Retrieve relevant chunks for a query.

        This is the ONLINE retrieval phase. For each query:
        1. Search the vector store for initial_k candidates (fast).
        2. Optionally re-rank with a cross-encoder (precise).
        3. Return the top_k results.

        Args:
            query: The search query.
            top_k: Number of results to return.
            use_reranking: Whether to apply cross-encoder re-ranking.
            initial_k: Number of candidates for re-ranking.

        Returns:
            List of relevant chunks with scores.
        """
        # Initial retrieval (fast, approximate)
        k = initial_k if use_reranking else top_k
        results = self.collection.query(
            query_texts=[query],
            n_results=min(k, self.doc_count),
        )

        if not results["documents"][0]:
            return []

        chunks = []
        for i, (doc, metadata, distance) in enumerate(zip(
            results["documents"][0],
            results["metadatas"][0],
            results["distances"][0],
        )):
            chunks.append({
                "text": doc,
                "source": metadata["source"],
                "initial_score": 1 - distance,  # Convert distance to similarity
            })

        # Re-rank if requested (slow, precise)
        if use_reranking and len(chunks) > top_k:
            pairs = [(query, chunk["text"]) for chunk in chunks]
            rerank_scores = self.reranker.predict(pairs)
            for chunk, score in zip(chunks, rerank_scores):
                chunk["rerank_score"] = float(score)
            chunks.sort(key=lambda x: x["rerank_score"], reverse=True)

        return chunks[:top_k]

    def build_prompt(
        self, query: str, retrieved_chunks: list[dict]
    ) -> str:
        """Build the augmented prompt with retrieved context.

        This is the AUGMENT phase: we construct a prompt that includes
        both the user's question and the retrieved context. The prompt
        instructs the LLM to answer based on the context and to
        acknowledge when the context is insufficient.

        Args:
            query: The user's question.
            retrieved_chunks: List of retrieved text chunks.

        Returns:
            The complete prompt string for the LLM.
        """
        context_parts = []
        for i, chunk in enumerate(retrieved_chunks, 1):
            context_parts.append(f"[Source {i}: {chunk['source']}]\n{chunk['text']}")

        context = "\n\n".join(context_parts)

        prompt = f"""Answer the following question based on the provided context.
If the context does not contain enough information to answer the question,
say so clearly and explain what information is missing.

Context:
{context}

Question: {query}

Answer:"""
        return prompt

    def query(
        self,
        question: str,
        llm_call,
        top_k: int = 5,
        use_reranking: bool = True,
    ) -> dict:
        """Execute a full RAG query: retrieve, augment, generate.

        This is the complete ONLINE pipeline:
        1. RETRIEVE: Find relevant chunks in the vector store.
        2. AUGMENT: Build a prompt with the retrieved context.
        3. GENERATE: Call the LLM to produce an answer.

        Args:
            question: The user's question.
            llm_call: A callable that takes a prompt string and returns a response.
            top_k: Number of chunks to retrieve.
            use_reranking: Whether to use cross-encoder re-ranking.

        Returns:
            Dict with 'answer', 'sources', and 'prompt' keys.
        """
        # Retrieve
        chunks = self.retrieve(question, top_k=top_k, use_reranking=use_reranking)

        # Augment
        prompt = self.build_prompt(question, chunks)

        # Generate
        answer = llm_call(prompt=prompt)

        return {
            "answer": answer,
            "sources": [{"source": c["source"], "text": c["text"][:200]} for c in chunks],
            "prompt": prompt,
            "num_chunks_retrieved": len(chunks),
        }


# -- Usage Example ---------------------------------------------------

def main():
    """Demonstrate the RAG pipeline."""

    # Initialize the pipeline
    rag = SimpleRAGPipeline(chunk_size=300, chunk_overlap=50)

    # Sample documents
    documents = [
        {
            "content": (
                "Retrieval-Augmented Generation (RAG) is a technique that combines "
                "information retrieval with text generation. It was introduced by "
                "Lewis et al. in 2020. The key idea is to retrieve relevant documents "
                "from an external knowledge base and use them as additional context for "
                "the language model. This allows the model to access up-to-date "
                "information and reduce hallucinations. RAG has become a fundamental "
                "building block for modern AI applications."
            ),
            "source": "rag_overview.txt",
        },
        {
            "content": (
                "Vector databases are specialized databases designed to store and "
                "query high-dimensional vectors efficiently. They use approximate "
                "nearest neighbor (ANN) algorithms like HNSW to enable fast similarity "
                "search. Popular vector databases include Chroma, Pinecone, Weaviate, "
                "and Qdrant. They are essential infrastructure for RAG systems, as they "
                "store the embeddings of document chunks and enable semantic search."
            ),
            "source": "vector_databases.txt",
        },
        {
            "content": (
                "Fine-tuning is the process of updating a pre-trained model's weights "
                "on a smaller, task-specific dataset. Unlike RAG, which adds information "
                "at inference time, fine-tuning bakes knowledge into the model's "
                "parameters. Fine-tuning is better for adapting the model's style and "
                "behavior, while RAG is better for providing up-to-date factual "
                "information. Many production systems use both approaches together."
            ),
            "source": "fine_tuning.txt",
        },
    ]

    # Index documents
    num_chunks = rag.index_documents(documents)
    print(f"Indexed {num_chunks} chunks from {len(documents)} documents.\n")

    # Simulate an LLM call
    def mock_llm_call(prompt: str) -> str:
        return (
            "Based on the provided context, RAG (Retrieval-Augmented Generation) "
            "works by retrieving relevant documents from an external knowledge base "
            "and using them as additional context for the language model. This was "
            "introduced by Lewis et al. in 2020. The retrieved documents are stored "
            "as embeddings in vector databases, which use approximate nearest neighbor "
            "algorithms for efficient similarity search."
        )

    # Execute a query
    result = rag.query(
        "How does RAG work?",
        llm_call=mock_llm_call,
        top_k=3,
        use_reranking=False,
    )

    print("Question: How does RAG work?\n")
    print(f"Answer: {result['answer']}\n")
    print("Sources used:")
    for src in result["sources"]:
        print(f"  - {src['source']}: {src['text'][:100]}...")


if __name__ == "__main__":
    main()