M10: Advanced RAG Patterns

# chunking_strategies.py
# pip install nltk sentence-transformers numpy
# python -c "import nltk; nltk.download('punkt')"

from __future__ import annotations
import re
import numpy as np
import nltk
from sentence_transformers import SentenceTransformer
from typing import List

# WHAT: Shared sentence encoder used by the semantic chunker
# WHY: Load once at module level — model initialization is expensive (~3s)
# GOTCHA: all-MiniLM-L6-v2 runs on CPU; for GPU, set device='cuda'
_encoder = None
def _get_encoder():
    global _encoder
    if _encoder is None:
        _encoder = SentenceTransformer("all-MiniLM-L6-v2")
    return _encoder


# ── Strategy 1: Fixed-size with overlap ───────────────────────────────────────
# WHAT: Naive baseline — split every chunk_size chars with overlap chars repeated
# WHY: Overlap ensures a sentence cut at a boundary is still findable from either chunk
# GOTCHA: overlap must be < chunk_size; typical ratio is 10-20%
def fixed_size_chunks(text: str, chunk_size: int = 512, overlap: int = 64) -> List[str]:
    chunks = []
    start = 0
    while start < len(text):
        end = min(start + chunk_size, len(text))
        chunks.append(text[start:end])
        start += chunk_size - overlap
    return chunks


# ── Strategy 2: Sentence-aware chunking ───────────────────────────────────────
# WHAT: Group sentences until a target_chars limit is reached; start a new chunk
# WHY: Sentences are the atomic unit of meaning in most prose documents
# GOTCHA: Very long sentences (technical specs, legal clauses) can produce
#         over-sized chunks; add a fallback fixed split for sentences > 1000 chars
def sentence_aware_chunks(text: str, target_chars: int = 500) -> List[str]:
    try:
        sentences = nltk.sent_tokenize(text)
    except LookupError:
        nltk.download("punkt", quiet=True)
        sentences = nltk.sent_tokenize(text)

    chunks: List[str] = []
    current: List[str] = []
    current_len = 0

    for sent in sentences:
        if current_len + len(sent) > target_chars and current:
            chunks.append(" ".join(current))
            current = []
            current_len = 0
        current.append(sent)
        current_len += len(sent) + 1

    if current:
        chunks.append(" ".join(current))
    return chunks


# ── Strategy 3: Semantic chunking ─────────────────────────────────────────────
# WHAT: Split where cosine distance between consecutive sentence embeddings exceeds threshold
# WHY: Detects actual topic shifts, not just length limits
# GOTCHA: threshold=0.3 works for technical docs; use 0.2 for narrative prose
#         where topics transition more gradually
def semantic_chunks(text: str, threshold: float = 0.3) -> List[str]:
    try:
        sentences = nltk.sent_tokenize(text)
    except LookupError:
        nltk.download("punkt", quiet=True)
        sentences = nltk.sent_tokenize(text)

    if len(sentences) <= 1:
        return sentences

    encoder = _get_encoder()
    embeddings = encoder.encode(sentences, show_progress_bar=False)

    # Compute cosine distance between consecutive sentences
    def cosine_dist(a: np.ndarray, b: np.ndarray) -> float:
        return 1.0 - float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-8))

    chunks: List[str] = []
    current: List[str] = [sentences[0]]

    for i in range(1, len(sentences)):
        dist = cosine_dist(embeddings[i - 1], embeddings[i])
        if dist > threshold:
            # Topic shift detected — flush current group
            chunks.append(" ".join(current))
            current = [sentences[i]]
        else:
            current.append(sentences[i])

    if current:
        chunks.append(" ".join(current))
    return chunks


# ── Strategy 4: Recursive character splitting ─────────────────────────────────
# WHAT: Try each separator in order; recurse on pieces still larger than max_size
# WHY: Respects document structure (paragraphs > sentences > words) without
#      requiring any NLP — fast and reliable on structured docs
# GOTCHA: separators list must be ordered from coarsest to finest
def recursive_chunks(
    text: str,
    max_size: int = 500,
    separators: List[str] | None = None,
) -> List[str]:
    if separators is None:
        separators = ["\n\n", "\n", ". ", "? ", "! ", " "]

    def _split(text: str, seps: List[str]) -> List[str]:
        if len(text) <= max_size:
            return [text] if text.strip() else []
        sep, *rest = seps
        if not sep:  # last resort: character split
            return [text[i:i + max_size] for i in range(0, len(text), max_size)]
        parts = text.split(sep)
        result: List[str] = []
        for part in parts:
            part = part.strip()
            if not part:
                continue
            if len(part) <= max_size:
                result.append(part)
            else:
                result.extend(_split(part, rest))
        return result

    return _split(text, separators)


# ── Demo: compare strategies on one passage ───────────────────────────────────
if __name__ == "__main__":
    sample = (
        "Metformin is the first-line pharmacological treatment for type 2 diabetes. "
        "It works by decreasing hepatic glucose production and improving insulin sensitivity. "
        "Common side effects include nausea, vomiting, and diarrhea, particularly when starting treatment. "
        "These GI effects usually subside after a few weeks. "
        "Serious but rare side effects include lactic acidosis, which requires immediate medical attention. "
        "Kidney function should be monitored every 3-6 months because metformin is excreted unchanged by the kidneys. "
        "Patients with eGFR below 30 mL/min should not use metformin. "
        "Drug interactions include iodinated contrast agents, which may cause acute kidney injury when combined."
    )

    for name, fn in [
        ("Fixed+Overlap",  lambda t: fixed_size_chunks(t, 200, 30)),
        ("Sentence-Aware", lambda t: sentence_aware_chunks(t, 300)),
        ("Semantic",       lambda t: semantic_chunks(t, 0.3)),
        ("Recursive",      lambda t: recursive_chunks(t, 250)),
    ]:
        chunks = fn(sample)
        print(f"\n{'='*50}")
        print(f"{name}: {len(chunks)} chunks")
        for i, c in enumerate(chunks):
            print(f"  [{i+1}] ({len(c)}c) {c[:80]}{'...' if len(c)>80 else ''}")

# hybrid_search.py
# pip install rank-bm25 chromadb sentence-transformers

from rank_bm25 import BM25Okapi
import chromadb
from sentence_transformers import SentenceTransformer
from typing import List, Dict, Any
import numpy as np

# WHAT: Build both indexes from the same document list
# WHY: Keeping them in sync ensures BM25 rank and dense rank cover identical docs
# GOTCHA: Use doc IDs consistently — ChromaDB returns them by ID, BM25 by list index
class HybridSearchIndex:
    def __init__(self, model_name: str = "all-MiniLM-L6-v2"):
        self.encoder = SentenceTransformer(model_name)
        self.client = chromadb.Client()
        # Delete collection if it exists (for clean demo runs)
        try:
            self.client.delete_collection("hybrid_docs")
        except Exception:
            pass
        self.collection = self.client.create_collection("hybrid_docs")
        self.docs: List[str] = []
        self.doc_ids: List[str] = []
        self.bm25: BM25Okapi | None = None

    def add_documents(self, docs: List[str], ids: List[str]) -> None:
        self.docs = docs
        self.doc_ids = ids

        # WHAT: Build BM25 index — tokenize by lowercasing and splitting on spaces
        # GOTCHA: For better BM25 performance, add stop-word removal and stemming
        tokenized = [d.lower().split() for d in docs]
        self.bm25 = BM25Okapi(tokenized)

        # WHAT: Embed and add to ChromaDB
        embeddings = self.encoder.encode(docs, show_progress_bar=False).tolist()
        self.collection.add(
            documents=docs,
            embeddings=embeddings,
            ids=ids,
        )
        print(f"Indexed {len(docs)} documents")

    # WHAT: Reciprocal Rank Fusion — merge ranked lists from BM25 and dense
    # WHY: RRF(k=60) is robust: k prevents a rank-1 result from completely
    #      dominating the merged score; default k=60 is well-validated
    @staticmethod
    def _rrf(rankings: List[List[str]], k: int = 60) -> Dict[str, float]:
        scores: Dict[str, float] = {}
        for ranked_list in rankings:
            for rank, doc_id in enumerate(ranked_list, start=1):
                scores[doc_id] = scores.get(doc_id, 0.0) + 1.0 / (k + rank)
        return scores

    def search(self, query: str, top_k: int = 5, fetch_k: int = 20) -> List[Dict[str, Any]]:
        """Hybrid search: returns top_k docs after BM25 + dense + RRF merge."""

        # ── BM25 lexical retrieval ────────────────────────────────────────────
        bm25_scores = self.bm25.get_scores(query.lower().split())
        bm25_ranked = [
            self.doc_ids[i]
            for i in np.argsort(bm25_scores)[::-1][:fetch_k]
        ]

        # ── Dense semantic retrieval ──────────────────────────────────────────
        query_embedding = self.encoder.encode([query], show_progress_bar=False).tolist()
        dense_results = self.collection.query(
            query_embeddings=query_embedding,
            n_results=min(fetch_k, len(self.docs)),
        )
        dense_ranked = dense_results["ids"][0]  # already sorted by similarity

        # ── RRF merge ────────────────────────────────────────────────────────
        rrf_scores = self._rrf([bm25_ranked, dense_ranked])
        sorted_ids = sorted(rrf_scores, key=rrf_scores.get, reverse=True)[:top_k]

        # ── Build result objects ──────────────────────────────────────────────
        id_to_doc = dict(zip(self.doc_ids, self.docs))
        return [
            {
                "id": doc_id,
                "text": id_to_doc[doc_id],
                "rrf_score": rrf_scores[doc_id],
            }
            for doc_id in sorted_ids
            if doc_id in id_to_doc
        ]


# ── Demo ──────────────────────────────────────────────────────────────────────
if __name__ == "__main__":
    docs = [
        "Metformin lowers blood sugar by reducing hepatic glucose production.",
        "FDA approval number 123-456 was granted for metformin HCl 500mg tablets.",
        "Side effects of metformin include nausea, diarrhea, and abdominal discomfort.",
        "Drug interactions: iodinated contrast agents may cause lactic acidosis.",
        "Metformin is contraindicated in patients with eGFR below 30 mL/min.",
        "FDA approval number 789-012 covers the extended-release formulation.",
        "Kidney function monitoring every 6 months is required for all metformin users.",
    ]
    ids = [f"doc-{i}" for i in range(len(docs))]

    index = HybridSearchIndex()
    index.add_documents(docs, ids)

    for query in ["FDA approval number 123-456", "what are the side effects"]:
        print(f"\nQuery: {query!r}")
        results = index.search(query, top_k=3)
        for r in results:
            print(f"  [{r['rrf_score']:.4f}] {r['text'][:80]}")

# reranker.py
# pip install sentence-transformers openai

from sentence_transformers import CrossEncoder
from openai import OpenAI
import json
from typing import List, Dict, Any

# WHAT: Cross-encoder re-ranker using local model
# WHY: cross-encoder/ms-marco-MiniLM-L-6-v2 is specifically trained
#      on MS MARCO passage ranking — best single-model choice for RAG reranking
# GOTCHA: load the model once at startup; each predict() call is a forward pass
#         on all (query, doc) pairs simultaneously — batch efficiently
class CrossEncoderReranker:
    def __init__(self, model: str = "cross-encoder/ms-marco-MiniLM-L-6-v2"):
        self.model = CrossEncoder(model)

    def rerank(
        self,
        query: str,
        candidates: List[Dict[str, Any]],
        top_k: int = 5,
    ) -> List[Dict[str, Any]]:
        if not candidates:
            return []

        # WHAT: Build (query, document) pairs for the cross-encoder
        # GOTCHA: cross-encoder scores all pairs in one batch call —
        #         much faster than looping with individual predict() calls
        pairs = [(query, c["text"]) for c in candidates]
        scores = self.model.predict(pairs)   # returns numpy array of logit scores

        # Sort by score descending, return top_k
        ranked = sorted(
            zip(candidates, scores),
            key=lambda x: x[1],
            reverse=True,
        )
        return [
            {**cand, "rerank_score": float(score)}
            for cand, score in ranked[:top_k]
        ]


# WHAT: LLM-based re-ranker using Mistral — useful when relevance criteria
#       are domain-specific and can't be captured by a general-purpose cross-encoder
# WHY: For specialized domains (clinical trials, legal precedent, UCC filings),
#      Mistral can apply domain logic the MS MARCO model never saw
# GOTCHA: 20 Ollama calls per query adds 2-8s latency — use only when accuracy
#         matters more than speed and for offline batch pipelines
class LLMReranker:
    def __init__(self, base_url: str = "http://localhost:11434/v1"):
        self.client = OpenAI(base_url=base_url, api_key="ollama")

    def _score_candidate(self, query: str, text: str) -> float:
        try:
            resp = self.client.chat.completions.create(
                model="mistral",
                messages=[
                    {
                        "role": "system",
                        "content": (
                            "You are a relevance judge. Score how relevant the document "
                            "is to the query on a scale of 0.0 to 1.0. "
                            "Respond with JSON only: {\"score\": 0.0}"
                        ),
                    },
                    {
                        "role": "user",
                        "content": f"Query: {query}\n\nDocument: {text[:500]}",
                    },
                ],
                max_tokens=32,
                temperature=0,
            )
            raw = resp.choices[0].message.content or "{}"
            # Strip markdown fences if present
            raw = raw.strip().lstrip("```json").lstrip("```").rstrip("```")
            return float(json.loads(raw).get("score", 0.0))
        except Exception:
            return 0.0

    def rerank(
        self,
        query: str,
        candidates: List[Dict[str, Any]],
        top_k: int = 5,
    ) -> List[Dict[str, Any]]:
        scored = [
            {**c, "rerank_score": self._score_candidate(query, c["text"])}
            for c in candidates
        ]
        return sorted(scored, key=lambda x: x["rerank_score"], reverse=True)[:top_k]


# ── Demo ──────────────────────────────────────────────────────────────────────
if __name__ == "__main__":
    candidates = [
        {"id": "doc-0", "text": "Metformin lowers blood sugar through hepatic glucose reduction."},
        {"id": "doc-1", "text": "Side effects of metformin: nausea, diarrhea, lactic acidosis."},
        {"id": "doc-2", "text": "FDA approval number 123-456 granted for metformin 500mg."},
        {"id": "doc-3", "text": "Drug interactions: contrast agents plus metformin risk lactic acidosis."},
        {"id": "doc-4", "text": "Kidney function monitoring schedule for metformin patients."},
    ]
    query = "What are the serious adverse effects of metformin?"

    print("Cross-Encoder Re-ranking:")
    reranker = CrossEncoderReranker()
    for r in reranker.rerank(query, candidates, top_k=3):
        print(f"  [{r['rerank_score']:.3f}] {r['text'][:70]}")

# multi_query.py
# pip install openai sentence-transformers numpy

from openai import OpenAI
from sentence_transformers import SentenceTransformer
import json
import numpy as np
from typing import List, Dict, Any, Callable

client = OpenAI(base_url="http://localhost:11434/v1", api_key="ollama")
encoder = SentenceTransformer("all-MiniLM-L6-v2")

# WHAT: Ask Mistral to generate N semantically diverse query variants
# WHY: Diverse phrasings surface documents that use different terminology
#      for the same concept — crucial for technical and medical domains
# GOTCHA: Ask for JSON array output; Mistral sometimes wraps it in markdown
#         fences — strip them before parsing
def expand_query(query: str, n: int = 3) -> List[str]:
    resp = client.chat.completions.create(
        model="mistral",
        messages=[
            {
                "role": "system",
                "content": (
                    f"Generate {n} alternative phrasings of the user's search query. "
                    "Each variant should use different vocabulary while preserving the "
                    "original intent. Include technical synonyms, related terms, and "
                    "alternative framings. Return a JSON array of strings only."
                ),
            },
            {"role": "user", "content": query},
        ],
        response_format={"type": "json_object"},
        temperature=0.5,
        max_tokens=256,
    )
    raw = resp.choices[0].message.content or "[]"
    # GOTCHA: response_format may still return {"queries": [...]} — handle both shapes
    parsed = json.loads(raw)
    if isinstance(parsed, list):
        variants = parsed
    else:
        variants = next(
            (v for v in parsed.values() if isinstance(v, list)), []
        )
    return [query] + [str(v) for v in variants[:n]]   # include original


# WHAT: Deduplicate a list of retrieved docs by embedding similarity
# WHY: Multiple queries often retrieve the same doc under different rankings
#      Text dedup misses paraphrases; embedding dedup is more robust
# GOTCHA: threshold=0.95 is conservative — lower it (0.85) for cleaner dedup
#         at the cost of occasionally removing genuinely different docs
def deduplicate_by_embedding(
    docs: List[Dict[str, Any]],
    threshold: float = 0.95,
) -> List[Dict[str, Any]]:
    if not docs:
        return []
    texts = [d["text"] for d in docs]
    embeddings = encoder.encode(texts, show_progress_bar=False)

    keep = []
    for i, doc in enumerate(docs):
        is_dup = False
        for j in keep:
            sim = float(np.dot(embeddings[i], embeddings[j]) /
                        (np.linalg.norm(embeddings[i]) * np.linalg.norm(embeddings[j]) + 1e-8))
            if sim > threshold:
                is_dup = True
                break
        if not is_dup:
            keep.append(i)

    return [docs[i] for i in keep]


# WHAT: Full multi-query retrieval pipeline
# WHY: Combining all variants before dedup gives the re-ranker the richest pool
def multi_query_retrieve(
    query: str,
    search_fn: Callable[[str], List[Dict[str, Any]]],
    n_variants: int = 3,
    dedup_threshold: float = 0.95,
) -> List[Dict[str, Any]]:
    variants = expand_query(query, n_variants)
    print(f"Query variants: {variants}")

    # Retrieve for each variant; label source variant for debugging
    all_results: List[Dict[str, Any]] = []
    seen_ids = set()
    for v in variants:
        results = search_fn(v)
        for r in results:
            if r.get("id") not in seen_ids:
                all_results.append({**r, "_source_variant": v})
                seen_ids.add(r.get("id"))

    # Final embedding-level dedup before re-ranking
    return deduplicate_by_embedding(all_results, threshold=dedup_threshold)


# ── Demo ──────────────────────────────────────────────────────────────────────
if __name__ == "__main__":
    variants = expand_query("Python memory leak", n=3)
    print("Expanded queries:")
    for i, v in enumerate(variants):
        print(f"  {i+1}. {v}")

# contextual_compression.py
# pip install openai

from openai import OpenAI
from typing import List, Dict, Any

client = OpenAI(base_url="http://localhost:11434/v1", api_key="ollama")

# WHAT: Ask Mistral to extract the relevant excerpt from a chunk
# WHY: The model sees the query + chunk and returns only the directly useful text
# GOTCHA: If the entire chunk is relevant, Mistral will (correctly) return
#         most of it — don't expect dramatic compression in all cases
def compress_chunk(query: str, chunk: str, max_tokens: int = 256) -> str:
    if len(chunk) < 150:
        return chunk   # too short to compress; return as-is

    resp = client.chat.completions.create(
        model="mistral",
        messages=[
            {
                "role": "system",
                "content": (
                    "Extract the portion of the provided document that is directly "
                    "relevant to the query. Return only the relevant excerpt verbatim. "
                    "If no portion is relevant, respond with: [NOT RELEVANT]"
                ),
            },
            {
                "role": "user",
                "content": f"Query: {query}\n\nDocument:\n{chunk}",
            },
        ],
        max_tokens=max_tokens,
        temperature=0,
    )
    excerpt = (resp.choices[0].message.content or "").strip()

    # GOTCHA: Fall back to full chunk if model returns nothing useful
    if not excerpt or excerpt == "[NOT RELEVANT]" or len(excerpt) < 10:
        return chunk
    return excerpt


def compress_results(
    query: str,
    results: List[Dict[str, Any]],
    max_tokens_per_chunk: int = 256,
) -> List[Dict[str, Any]]:
    """Compress all retrieved chunks and filter out non-relevant ones."""
    compressed = []
    for r in results:
        excerpt = compress_chunk(query, r["text"], max_tokens_per_chunk)
        if excerpt != "[NOT RELEVANT]":
            compressed.append({**r, "text": excerpt, "original_text": r["text"]})
    return compressed


# ── Demo ──────────────────────────────────────────────────────────────────────
if __name__ == "__main__":
    long_chunk = (
        "Metformin is a biguanide class antidiabetic agent. It was first described in "
        "the scientific literature in 1922, synthesized in 1929, and approved in France "
        "in 1957. It works primarily by reducing hepatic glucose production. "
        "The mechanism involves activation of AMP-activated protein kinase (AMPK). "
        "SERIOUS SIDE EFFECTS: Lactic acidosis — rare but potentially fatal. "
        "Risk factors include renal impairment, hepatic impairment, and iodinated contrast "
        "agents. Stop metformin 48 hours before contrast procedures. "
        "Monitor serum creatinine annually in all patients and every 3-6 months in "
        "patients with eGFR 30-60. "
        "The drug is available as generic metformin HCl and various extended-release formulations."
    )
    query = "What are the serious side effects of metformin?"
    result = compress_chunk(query, long_chunk)
    tokens_saved = len(long_chunk) // 4 - len(result) // 4
    print(f"Original: {len(long_chunk)} chars")
    print(f"Compressed: {len(result)} chars (~{tokens_saved} tokens saved)")
    print(f"\nExcerpt: {result}")

# parent_child_chunking.py
# pip install chromadb sentence-transformers nltk

import chromadb
from sentence_transformers import SentenceTransformer
import nltk
from typing import List, Dict, Any

encoder = SentenceTransformer("all-MiniLM-L6-v2")

class ParentChildIndex:
    def __init__(self):
        self.client = chromadb.Client()
        # WHAT: Two collections — children are indexed, parents are stored for retrieval
        # GOTCHA: Delete collections before re-indexing to avoid duplicate IDs
        for name in ["parent_chunks", "child_chunks"]:
            try: self.client.delete_collection(name)
            except Exception: pass
        self.parents = self.client.create_collection("parent_chunks")
        self.children = self.client.create_collection("child_chunks")

    def _sentence_chunks(self, text: str, target_chars: int) -> List[str]:
        try:
            sentences = nltk.sent_tokenize(text)
        except LookupError:
            nltk.download("punkt", quiet=True)
            sentences = nltk.sent_tokenize(text)
        chunks, current, current_len = [], [], 0
        for sent in sentences:
            if current_len + len(sent) > target_chars and current:
                chunks.append(" ".join(current))
                current, current_len = [], 0
            current.append(sent)
            current_len += len(sent) + 1
        if current:
            chunks.append(" ".join(current))
        return chunks

    def index_document(
        self,
        doc_id: str,
        text: str,
        parent_size: int = 500,
        child_size: int = 120,
    ) -> None:
        # WHAT: Create parent chunks (larger, for context-rich retrieval)
        parent_chunks = self._sentence_chunks(text, parent_size)
        parent_ids = [f"{doc_id}__parent_{i}" for i in range(len(parent_chunks))]

        parent_embeddings = encoder.encode(parent_chunks, show_progress_bar=False).tolist()
        self.parents.add(
            documents=parent_chunks,
            embeddings=parent_embeddings,
            ids=parent_ids,
        )

        # WHAT: Create child chunks from each parent and store parent_id in metadata
        # WHY: The metadata link lets us look up the parent after a child matches
        all_child_docs, all_child_ids, all_child_metas, all_child_embs = [], [], [], []
        for p_idx, (parent_text, parent_id) in enumerate(zip(parent_chunks, parent_ids)):
            children = self._sentence_chunks(parent_text, child_size)
            for c_idx, child_text in enumerate(children):
                child_id = f"{doc_id}__child_{p_idx}_{c_idx}"
                all_child_docs.append(child_text)
                all_child_ids.append(child_id)
                all_child_metas.append({"parent_id": parent_id})

        if all_child_docs:
            all_child_embs = encoder.encode(all_child_docs, show_progress_bar=False).tolist()
            self.children.add(
                documents=all_child_docs,
                embeddings=all_child_embs,
                ids=all_child_ids,
                metadatas=all_child_metas,
            )
        print(f"Indexed {len(parent_chunks)} parents, {len(all_child_docs)} children for {doc_id!r}")

    def search(self, query: str, top_k: int = 3) -> List[Dict[str, Any]]:
        """Search children, then return their corresponding parents."""
        query_emb = encoder.encode([query], show_progress_bar=False).tolist()

        # WHAT: Retrieve child matches — small chunks for precision
        child_results = self.children.query(
            query_embeddings=query_emb,
            n_results=min(top_k * 3, 20),  # fetch 3x to ensure enough unique parents
        )

        # WHAT: Look up parent IDs from child metadata; deduplicate
        seen_parent_ids = set()
        parent_ids_ordered = []
        for meta in child_results["metadatas"][0]:
            pid = meta["parent_id"]
            if pid not in seen_parent_ids:
                seen_parent_ids.add(pid)
                parent_ids_ordered.append(pid)
                if len(parent_ids_ordered) >= top_k:
                    break

        # WHAT: Fetch parent documents by ID for the final context
        if not parent_ids_ordered:
            return []
        parent_data = self.parents.get(ids=parent_ids_ordered)
        return [
            {"id": pid, "text": text}
            for pid, text in zip(parent_data["ids"], parent_data["documents"])
        ]


if __name__ == "__main__":
    long_doc = (
        "Metformin is the first-line treatment for type 2 diabetes. It acts by inhibiting "
        "complex I of the mitochondrial respiratory chain. The primary effect is reduction "
        "of hepatic glucose output. Insulin sensitivity is also improved. "
        "Gastrointestinal side effects are common at initiation: nausea, vomiting, diarrhea. "
        "These typically resolve within 2-4 weeks. Taking metformin with food reduces GI effects. "
        "Lactic acidosis is the most serious adverse effect. It is rare (incidence ~3 per 100,000 "
        "patient-years) but carries ~50% mortality. Risk factors: renal impairment (eGFR < 30), "
        "hepatic impairment, congestive heart failure, excessive alcohol use. "
        "Contraindication: hold metformin 48h before and after IV contrast procedures."
    )
    index = ParentChildIndex()
    index.index_document("metformin-overview", long_doc)
    results = index.search("serious adverse effects of metformin", top_k=2)
    print("\nParent-Child Results:")
    for r in results:
        print(f"\n[{r['id']}]\n{r['text']}")

# production_rag_pipeline.py
# pip install rank-bm25 sentence-transformers chromadb nltk openai

from __future__ import annotations
import time
from dataclasses import dataclass, field
from typing import List, Optional
from openai import OpenAI

# Import our custom modules (defined in this module set)
from chunking_strategies import semantic_chunks
from hybrid_search import HybridSearchIndex
from reranker import CrossEncoderReranker
from contextual_compression import compress_results

# WHAT: Pipeline configuration — toggle features with boolean flags
# WHY: Allows quick ablation testing without code changes
@dataclass
class PipelineConfig:
    chunk_strategy: str = "semantic"     # "fixed", "sentence", "semantic", "recursive"
    use_hybrid_search: bool = True       # BM25 + dense vs dense-only
    use_reranking: bool = True           # cross-encoder re-ranker
    use_compression: bool = False        # contextual compression (adds latency)
    retrieve_k: int = 20                 # candidates before re-ranking
    final_k: int = 5                     # chunks passed to generation
    ollama_url: str = "http://localhost:11434/v1"
    model: str = "mistral"


class ProductionRAGPipeline:
    def __init__(self, config: PipelineConfig = None):
        self.config = config or PipelineConfig()
        self.index = HybridSearchIndex()
        self.reranker = CrossEncoderReranker() if self.config.use_reranking else None
        self.client = OpenAI(base_url=self.config.ollama_url, api_key="ollama")
        self._documents_indexed = False

    # WHAT: Index a document with the configured chunking strategy
    def index_documents(self, texts: List[str], ids: List[str]) -> None:
        all_chunks, all_ids = [], []
        for doc_id, text in zip(ids, texts):
            if self.config.chunk_strategy == "semantic":
                chunks = semantic_chunks(text, threshold=0.3)
            elif self.config.chunk_strategy == "sentence":
                from chunking_strategies import sentence_aware_chunks
                chunks = sentence_aware_chunks(text, 400)
            elif self.config.chunk_strategy == "recursive":
                from chunking_strategies import recursive_chunks
                chunks = recursive_chunks(text, 400)
            else:  # fixed
                from chunking_strategies import fixed_size_chunks
                chunks = fixed_size_chunks(text, 400, 60)
            for i, chunk in enumerate(chunks):
                all_chunks.append(chunk)
                all_ids.append(f"{doc_id}__chunk_{i}")
        self.index.add_documents(all_chunks, all_ids)
        self._documents_indexed = True

    def _retrieve(self, query: str) -> List[dict]:
        return self.index.search(query, top_k=self.config.retrieve_k)

    def _rerank(self, query: str, candidates: List[dict]) -> List[dict]:
        if self.reranker:
            return self.reranker.rerank(query, candidates, top_k=self.config.final_k)
        return candidates[:self.config.final_k]

    def _compress(self, query: str, chunks: List[dict]) -> List[dict]:
        if self.config.use_compression:
            return compress_results(query, chunks)
        return chunks

    def _generate(self, query: str, context_chunks: List[dict]) -> str:
        context = "\n\n".join(
            f"[{i+1}] {c['text']}" for i, c in enumerate(context_chunks)
        )
        resp = self.client.chat.completions.create(
            model=self.config.model,
            messages=[
                {
                    "role": "system",
                    "content": (
                        "Answer the question using only the provided context. "
                        "Cite context numbers like [1], [2] when referencing sources. "
                        "If the context doesn't contain the answer, say so explicitly."
                    ),
                },
                {
                    "role": "user",
                    "content": f"Context:\n{context}\n\nQuestion: {query}",
                },
            ],
            max_tokens=512,
            temperature=0,
        )
        return resp.choices[0].message.content or ""

    def query(self, question: str) -> dict:
        """Run the full pipeline; return answer and timing breakdown."""
        timings = {}

        t0 = time.time()
        candidates = self._retrieve(question)
        timings["retrieve_ms"] = int((time.time() - t0) * 1000)

        t0 = time.time()
        reranked = self._rerank(question, candidates)
        timings["rerank_ms"] = int((time.time() - t0) * 1000)

        t0 = time.time()
        final_chunks = self._compress(question, reranked)
        timings["compress_ms"] = int((time.time() - t0) * 1000)

        t0 = time.time()
        answer = self._generate(question, final_chunks)
        timings["generate_ms"] = int((time.time() - t0) * 1000)

        timings["total_ms"] = sum(timings.values())
        return {
            "answer": answer,
            "sources": [c["id"] for c in final_chunks],
            "timings": timings,
            "chunk_count": len(final_chunks),
        }

    def benchmark(self, queries: List[str]) -> None:
        """Print timing breakdown across multiple queries."""
        print(f"\nBenchmark: {len(queries)} queries")
        print(f"Config: hybrid={self.config.use_hybrid_search}, "
              f"rerank={self.config.use_reranking}, "
              f"compress={self.config.use_compression}")
        print("-" * 60)
        totals = {"retrieve_ms": 0, "rerank_ms": 0, "compress_ms": 0, "generate_ms": 0}
        for q in queries:
            result = self.query(q)
            for k in totals:
                totals[k] += result["timings"][k]
        n = len(queries)
        for k, v in totals.items():
            print(f"  {k:20s}: {v//n:4d}ms avg")
        print(f"  {'total_ms':20s}: {sum(totals.values())//n:4d}ms avg")


# ── Demo ──────────────────────────────────────────────────────────────────────
if __name__ == "__main__":
    docs = [
        "Metformin is a biguanide antidiabetic drug and first-line therapy for type 2 diabetes. "
        "It reduces hepatic glucose production via AMPK activation. Contraindicated with eGFR < 30. "
        "Lactic acidosis risk increases with renal impairment. Hold 48h before contrast procedures. "
        "GI side effects (nausea, diarrhea) resolve within 2-4 weeks. Take with food.",
    ]
    config = PipelineConfig(use_reranking=True, use_compression=False, retrieve_k=10, final_k=3)
    pipeline = ProductionRAGPipeline(config)
    pipeline.index_documents(docs, ["metformin-doc"])
    result = pipeline.query("What are the contraindications for metformin?")
    print("\nAnswer:", result["answer"])
    print("Timings:", result["timings"])

# lab_benchmark.py
# Measures retrieval quality improvement at each pipeline stage
# Assumes you have the M09 corpus indexed in ChromaDB at ./chroma_db/

from production_rag_pipeline import ProductionRAGPipeline, PipelineConfig
import json, time

# WHAT: Five eval questions covering different retrieval challenge types
# WHY: Each question type is designed to stress-test a specific enhancement
EVAL_QUESTIONS = [
    {
        "id": "Q1",
        "question": "What is FDA approval number 123-456?",
        "type": "exact_match",           # tests BM25 keyword search
        "expected_keyword": "123-456",
    },
    {
        "id": "Q2",
        "question": "What are the serious adverse effects of metformin?",
        "type": "semantic",              # tests semantic chunking + reranking
        "expected_keyword": "lactic acidosis",
    },
    {
        "id": "Q3",
        "question": "When should metformin be held before a procedure?",
        "type": "specific_fact",         # tests sentence-aware chunking
        "expected_keyword": "48",
    },
    {
        "id": "Q4",
        "question": "kidney monitoring metformin",
        "type": "ambiguous",             # tests multi-query expansion
        "expected_keyword": "eGFR",
    },
    {
        "id": "Q5",
        "question": "Compare mechanism and side effects of metformin",
        "type": "synthesis",             # tests parent-child or semantic chunks
        "expected_keyword": "AMPK",
    },
]

def score_results(results, expected_keyword):
    """Check if expected keyword appears in top-3 retrieved chunks."""
    for i, r in enumerate(results[:3]):
        if expected_keyword.lower() in r.get("text", "").lower():
            return {"found": True, "rank": i + 1}
    return {"found": False, "rank": None}


def run_benchmark(docs, doc_ids, questions, configs):
    for config_name, config in configs.items():
        pipeline = ProductionRAGPipeline(config)
        pipeline.index_documents(docs, doc_ids)

        print(f"\n{'='*60}")
        print(f"Config: {config_name}")
        hits = 0
        for q in questions:
            t0 = time.time()
            # Direct retrieval (no generation) for benchmarking speed
            candidates = pipeline.index.search(q["question"], top_k=config.final_k)
            if config.use_reranking and pipeline.reranker:
                candidates = pipeline.reranker.rerank(q["question"], candidates, config.final_k)
            elapsed_ms = int((time.time() - t0) * 1000)
            result = score_results(candidates, q["expected_keyword"])
            if result["found"]:
                hits += 1
            icon = "PASS" if result["found"] else "FAIL"
            rank_str = f"rank {result['rank']}" if result["found"] else "not found"
            print(f"  [{icon}] {q['id']} ({q['type']}) — {rank_str} in {elapsed_ms}ms")

        print(f"  Score: {hits}/{len(questions)} ({hits*100//len(questions)}%)")


if __name__ == "__main__":
    # Load your M09 corpus here
    docs = [
        "Metformin is first-line treatment for type 2 diabetes. FDA approval number 123-456 "
        "was granted for metformin HCl 500mg. It reduces hepatic glucose via AMPK activation. "
        "Lactic acidosis is rare but serious — stop metformin if eGFR drops below 30. "
        "Hold 48 hours before and after iodinated contrast procedures. "
        "GI side effects: nausea, diarrhea. Monitor kidney function every 6 months.",
    ]
    doc_ids = ["metformin-guide"]

    configs = {
        "Baseline (dense-only, no rerank)": PipelineConfig(
            use_hybrid_search=False, use_reranking=False, final_k=5
        ),
        "+Hybrid Search (BM25+dense)": PipelineConfig(
            use_hybrid_search=True, use_reranking=False, final_k=5
        ),
        "+Hybrid +Reranking": PipelineConfig(
            use_hybrid_search=True, use_reranking=True, final_k=5
        ),
    }
    run_benchmark(docs, doc_ids, EVAL_QUESTIONS, configs)