Advanced RAG Techniques: Beyond Naive Chunking in Production

from dataclasses import dataclass
from typing import Optional

@dataclass
class QueryTransformConfig:
    num_queries: int = 3
    model: str = "gpt-4o-mini"
    temperature: float = 0.0

async def generate_search_queries(
    original_query: str,
    config: QueryTransformConfig,
) -> list[str]:
    """Generate diverse search queries from user input."""
    prompt = f"""Generate {config.num_queries} different search queries 
    that would help find documents answering this question.
    Return one query per line, no numbering.
    
    User question: {original_query}"""
    
    response = await llm_client.complete(
        prompt,
        model=config.model,
        temperature=config.temperature,
    )
    queries = [q.strip() for q in response.strip().split("\n") if q.strip()]
    return [original_query] + queries  # Always include original


async def multi_query_retrieve(
    query: str,
    vector_store,
    embed_pipeline,
    top_k: int = 10,
    filters: Optional[dict] = None,
) -> list[dict]:
    """Retrieve using multiple query variants, fuse with RRF."""
    queries = await generate_search_queries(query, QueryTransformConfig())
    
    all_results = []
    for q in queries:
        embedding = await embed_pipeline.embed_query(q)
        results = await vector_store.similarity_search(
            embedding=embedding,
            top_k=top_k,
            filter=filters,
        )
        all_results.append(results)
    
    return reciprocal_rank_fusion(all_results)


def reciprocal_rank_fusion(
    result_lists: list[list[dict]],
    k: int = 60,
) -> list[dict]:
    """Combine ranked lists using Reciprocal Rank Fusion."""
    scores: dict[str, float] = {}
    doc_map: dict[str, dict] = {}
    
    for results in result_lists:
        for rank, item in enumerate(results):
            doc_id = item["id"]
            scores[doc_id] = scores.get(doc_id, 0) + 1 / (k + rank + 1)
            doc_map[doc_id] = item
    
    ranked_ids = sorted(scores.keys(), key=lambda x: scores[x], reverse=True)
    return [doc_map[doc_id] for doc_id in ranked_ids]

async def hyde_retrieve(
    query: str,
    vector_store,
    embed_pipeline,
    top_k: int = 10,
    filters: Optional[dict] = None,
) -> list[dict]:
    """Hypothetical Document Embeddings retrieval."""
    hyde_prompt = f"""Write a detailed, factual paragraph that would 
    answer this question as if it appeared in a knowledge base article.
    Do not say "I don't know" — write the most likely answer.
    
    Question: {query}"""
    
    hypothetical_doc = await llm_client.complete(hyde_prompt, temperature=0.0)
    hyde_embedding = await embed_pipeline.embed_query(hypothetical_doc)
    
    return await vector_store.similarity_search(
        embedding=hyde_embedding,
        top_k=top_k,
        filter=filters,
    )

from dataclasses import dataclass, field
from typing import Optional

@dataclass
class HierarchicalChunk:
    child_id: str
    child_text: str
    child_embedding: list[float]
    parent_id: str
    parent_text: str
    metadata: dict = field(default_factory=dict)


class HierarchicalRetriever:
    def __init__(self, vector_store, embed_pipeline):
        self.vector_store = vector_store
        self.embed_pipeline = embed_pipeline
        self.parent_cache: dict[str, str] = {}

    async def index_document(
        self,
        doc_id: str,
        text: str,
        child_size: int = 300,
        parent_size: int = 1500,
        overlap: int = 50,
    ) -> int:
        """Split into parent chunks, then child chunks within each parent."""
        parents = self._split_text(text, parent_size, overlap)
        indexed = 0
        
        for p_idx, parent_text in enumerate(parents):
            parent_id = f"{doc_id}_p{p_idx}"
            self.parent_cache[parent_id] = parent_text
            
            children = self._split_text(parent_text, child_size, overlap // 3)
            for c_idx, child_text in enumerate(children):
                child_id = f"{parent_id}_c{c_idx}"
                embedding = (await self.embed_pipeline.embed_documents([child_text]))[0]
                
                await self.vector_store.upsert({
                    "id": child_id,
                    "embedding": embedding,
                    "text": child_text,
                    "metadata": {
                        "parent_id": parent_id,
                        "doc_id": doc_id,
                        "parent_text": parent_text,
                    },
                })
                indexed += 1
        
        return indexed

    async def retrieve(
        self,
        query: str,
        top_k: int = 5,
        filters: Optional[dict] = None,
    ) -> list[dict]:
        """Retrieve child chunks, deduplicate parents, return parent context."""
        query_embedding = await self.embed_pipeline.embed_query(query)
        child_results = await self.vector_store.similarity_search(
            embedding=query_embedding,
            top_k=top_k * 3,
            filter=filters,
        )
        
        seen_parents: set[str] = set()
        parent_results = []
        
        for child in child_results:
            parent_id = child["metadata"]["parent_id"]
            if parent_id not in seen_parents:
                seen_parents.add(parent_id)
                parent_results.append({
                    "id": parent_id,
                    "text": child["metadata"]["parent_text"],
                    "matched_child": child["text"],
                    "score": child.get("score", 0),
                })
            if len(parent_results) >= top_k:
                break
        
        return parent_results

    @staticmethod
    def _split_text(text: str, size: int, overlap: int) -> list[str]:
        chunks = []
        start = 0
        while start < len(text):
            end = start + size
            chunks.append(text[start:end])
            start = end - overlap
        return chunks

CONTEXTUALIZE_PROMPT = """<document>
{document}
</document>

Here is a chunk from the document:
<chunk>
{chunk}
</chunk>

Write a short context (2-3 sentences) explaining what this chunk 
is about within the document. Include document title, section, 
and relevant entities. Output only the context, no preamble."""

async def contextualize_chunks(
    document: str,
    chunks: list[str],
    doc_title: str = "",
) -> list[dict]:
    """Enrich chunks with LLM-generated context before embedding."""
    enriched = []
    
    for i, chunk in enumerate(chunks):
        context = await llm_client.complete(
            CONTEXTUALIZE_PROMPT.format(document=document[:8000], chunk=chunk),
            model="gpt-4o-mini",
            temperature=0.0,
        )
        
        contextualized_text = f"{context.strip()}\n\n{chunk}"
        enriched.append({
            "original_chunk": chunk,
            "context_prefix": context.strip(),
            "contextualized_text": contextualized_text,
            "chunk_index": i,
            "doc_title": doc_title,
        })
    
    return enriched


async def index_with_context(
    doc_id: str,
    document: str,
    title: str,
    vector_store,
    embed_pipeline,
    chunk_size: int = 800,
) -> int:
    """Full pipeline: chunk → contextualize → embed → store."""
    raw_chunks = split_on_headers(document, chunk_size)
    enriched = await contextualize_chunks(document, raw_chunks, title)
    
    texts = [e["contextualized_text"] for e in enriched]
    embeddings = await embed_pipeline.embed_documents(texts)
    
    for e, embedding in zip(enriched, embeddings):
        await vector_store.upsert({
            "id": f"{doc_id}_chunk_{e['chunk_index']}",
            "embedding": embedding,
            "text": e["original_chunk"],  # Store original for LLM context
            "metadata": {
                "doc_id": doc_id,
                "title": title,
                "context_prefix": e["context_prefix"],
            },
        })
    
    return len(enriched)

from dataclasses import dataclass, field
from typing import Optional

@dataclass
class Entity:
    id: str
    name: str
    entity_type: str
    properties: dict = field(default_factory=dict)

@dataclass
class Relationship:
    source_id: str
    target_id: str
    relation_type: str
    properties: dict = field(default_factory=dict)


class GraphRAGRetriever:
    def __init__(self, vector_store, graph_store, embed_pipeline):
        self.vector_store = vector_store
        self.graph = graph_store
        self.embed_pipeline = embed_pipeline

    async def extract_entities(self, text: str) -> tuple[list[Entity], list[Relationship]]:
        """Extract entities and relationships from text using LLM."""
        prompt = f"""Extract entities and relationships from this text.
        Return JSON with "entities" and "relationships" arrays.
        
        Text: {text[:4000]}"""
        
        result = await llm_client.complete_json(prompt, temperature=0.0)
        entities = [Entity(**e) for e in result["entities"]]
        relationships = [Relationship(**r) for r in result["relationships"]]
        return entities, relationships

    async def retrieve(
        self,
        query: str,
        top_k: int = 5,
        hop_depth: int = 2,
    ) -> list[dict]:
        """Graph-augmented retrieval: vector search + graph traversal."""
        query_embedding = await self.embed_pipeline.embed_query(query)
        seed_chunks = await self.vector_store.similarity_search(
            embedding=query_embedding,
            top_k=top_k,
        )
        
        # Step 2: Extract entity mentions from seed chunks
        seed_entities = set()
        for chunk in seed_chunks:
            entities = chunk.get("metadata", {}).get("entities", [])
            seed_entities.update(entities)
        
        # Step 3: Graph traversal from seed entities
        related_entities = await self.graph.traverse(
            start_nodes=list(seed_entities),
            max_hops=hop_depth,
        )
        
        # Step 4: Retrieve chunks linked to traversed entities
        graph_chunks = await self.vector_store.search_by_entities(
            entity_ids=[e.id for e in related_entities],
            top_k=top_k,
        )
        
        # Step 5: Merge and deduplicate
        all_chunks = {c["id"]: c for c in seed_chunks + graph_chunks}
        return list(all_chunks.values())[:top_k * 2]

from enum import Enum
from dataclasses import dataclass

class RetrievalAction(str, Enum):
    SEARCH = "search"
    REFINE_QUERY = "refine_query"
    SUFFICIENT = "sufficient"

@dataclass
class AgenticRAGState:
    original_query: str
    current_query: str
    retrieved_chunks: list[dict]
    iteration: int = 0
    max_iterations: int = 3


async def agentic_retrieve(
    query: str,
    retriever,
    max_iterations: int = 3,
) -> list[dict]:
    """Agent-driven retrieval with self-evaluation loop."""
    state = AgenticRAGState(
        original_query=query,
        current_query=query,
        retrieved_chunks=[],
        max_iterations=max_iterations,
    )
    
    while state.iteration < state.max_iterations:
        # Retrieve with current query
        new_chunks = await retriever.retrieve(
            state.current_query,
            top_k=10,
        )
        state.retrieved_chunks.extend(new_chunks)
        state.retrieved_chunks = deduplicate_by_id(state.retrieved_chunks)
        
        # Agent evaluates: sufficient or need more?
        evaluation = await evaluate_retrieval(
            query=state.original_query,
            chunks=state.retrieved_chunks,
        )
        
        if evaluation.action == RetrievalAction.SUFFICIENT:
            break
        
        if evaluation.action == RetrievalAction.REFINE_QUERY:
            state.current_query = evaluation.refined_query
        
        state.iteration += 1
    
    return state.retrieved_chunks[:7]


async def evaluate_retrieval(query: str, chunks: list[dict]) -> dict:
    """LLM evaluates whether retrieved chunks can answer the query."""
    chunk_summaries = "\n---\n".join(
        c["text"][:300] for c in chunks[:5]
    )
    
    prompt = f"""Given this question and retrieved document chunks, 
    decide if the chunks contain enough information to answer.
    
    Question: {query}
    
    Retrieved chunks:
    {chunk_summaries}
    
    Respond with JSON:
    {{"action": "sufficient" | "refine_query" | "search",
      "refined_query": "new query if refine_query",
      "reasoning": "brief explanation"}}"""
    
    return await llm_client.complete_json(prompt, temperature=0.0)

async def advanced_rerank(
    query: str,
    candidates: list[dict],
    top_k: int = 5,
) -> list[dict]:
    """Multi-signal reranking pipeline."""
    if not candidates:
        return []
    
    # Cross-encoder scores
    pairs = [(query, c["text"]) for c in candidates]
    ce_scores = cross_encoder.predict(pairs)
    
    # LLM relevance scores (batch for efficiency)
    llm_scores = await batch_llm_relevance(query, candidates)
    
    for i, candidate in enumerate(candidates):
        candidate["final_score"] = (
            0.40 * normalize(ce_scores[i])
            + 0.20 * normalize(candidate.get("vector_score", 0))
            + 0.15 * normalize(candidate.get("bm25_score", 0))
            + 0.15 * metadata_relevance(query, candidate.get("metadata", {}))
            + 0.10 * llm_scores[i]
        )
    
    ranked = sorted(candidates, key=lambda x: x["final_score"], reverse=True)
    deduped = deduplicate_chunks(ranked, threshold=0.85)
    return deduped[:top_k]


async def batch_llm_relevance(
    query: str,
    candidates: list[dict],
) -> list[float]:
    """Score chunk relevance using LLM judge."""
    prompt = f"""Rate each chunk's relevance to the query on a 0-10 scale.
    Query: {query}
    
    {format_chunks_for_scoring(candidates)}
    
    Return JSON array of scores, one per chunk."""
    
    scores = await llm_client.complete_json(prompt, temperature=0.0)
    return [s / 10.0 for s in scores]

Stage 1 (Baseline):     Chunk → Embed → Vector Search → LLM
Stage 2 (Retrieval+):  + Hybrid Search + Reranking
Stage 3 (Transform):    + Query Transformation + Parent-Child
Stage 4 (Context):      + Contextual Retrieval + Metadata Filters
Stage 5 (Advanced):     + Graph RAG OR Agentic Retrieval (pick one)

async def evaluate_advanced_rag(
    test_cases: list[dict],
    pipeline,
) -> dict:
    """Evaluate full RAG pipeline with retrieval + generation metrics."""
    metrics = {
        "precision_at_5": [],
        "mrr": [],
        "answer_faithfulness": [],
        "latency_ms": [],
    }
    
    for case in test_cases:
        start = time.monotonic()
        
        chunks = await pipeline.retrieve(case["query"])
        answer = await pipeline.generate(case["query"], chunks)
        
        latency = (time.monotonic() - start) * 1000
        metrics["latency_ms"].append(latency)
        
        # Retrieval metrics
        retrieved_ids = {c["metadata"]["doc_id"] for c in chunks}
        expected_ids = set(case["expected_doc_ids"])
        hits = len(retrieved_ids & expected_ids)
        metrics["precision_at_5"].append(hits / min(len(chunks), 5))
        
        # Answer faithfulness (does answer match retrieved context?)
        faithfulness = await score_faithfulness(answer, chunks, case["query"])
        metrics["answer_faithfulness"].append(faithfulness)
    
    return {
        "avg_precision_at_5": mean(metrics["precision_at_5"]),
        "avg_mrr": mean(metrics["mrr"]),
        "avg_faithfulness": mean(metrics["answer_faithfulness"]),
        "p95_latency_ms": percentile(metrics["latency_ms"], 95),
    }