M10: Advanced RAG Patterns

from anthropic import Anthropic
import chromadb

client = Anthropic()
chroma = chromadb.PersistentClient(path="./vector_db")

# Two corpora — the agent picks which to query per call.
INDICES = {
    "product_docs":    chroma.get_collection("product_docs"),
    "support_tickets": chroma.get_collection("support_tickets"),
}

# === The retrieval tool. Claude decides corpus + query + k. =================
tools = [{
    "name": "search",
    "description": (
        "Search a corpus by semantic similarity. "
        "Use 'product_docs' for behavior/spec questions, "
        "'support_tickets' for known-issues / customer cases. "
        "Call multiple times if the first result is insufficient or off-topic."
    ),
    "input_schema": {
        "type": "object",
        "properties": {
            "corpus": {"type": "string", "enum": list(INDICES)},
            "query":  {"type": "string", "description": "Natural-language query"},
            "k":      {"type": "integer", "minimum": 1, "maximum": 8, "default": 4},
        },
        "required": ["corpus", "query"],
    },
}]

def run_tool(name: str, args: dict) -> str:
    if name != "search":
        raise ValueError(f"Unknown tool: {name}")
    idx = INDICES.get(args["corpus"])
    if idx is None:
        return f"ERROR: unknown corpus '{args['corpus']}'"
    res = idx.query(query_texts=[args["query"]], n_results=args.get("k", 4))
    docs = res["documents"][0]
    return "\n---\n".join(f"[chunk {i+1}] {d}" for i, d in enumerate(docs)) or "(no hits)"

# === The agentic loop — standard M12 ReAct, just with `search` as the tool. ===
def agentic_rag(question: str, *, max_iters: int = 6) -> str:
    messages = [{"role": "user", "content": question}]
    for step in range(max_iters):
        resp = client.messages.create(
            model="claude-sonnet-4-6",
            max_tokens=2048,
            tools=tools,
            messages=messages,
        )
        if resp.stop_reason == "end_turn":
            return "".join(b.text for b in resp.content if b.type == "text")

        messages.append({"role": "assistant", "content": resp.content})
        tool_results = []
        for block in resp.content:
            if block.type == "tool_use":
                try:
                    result = run_tool(block.name, block.input)
                except Exception as e:
                    result = f"TOOL_ERROR: {e}"
                tool_results.append({
                    "type": "tool_result", "tool_use_id": block.id, "content": result,
                })
        messages.append({"role": "user", "content": tool_results})

    # Hit the step cap — escalate per M17 rather than guess.
    raise RuntimeError(f"Agentic RAG exceeded {max_iters} retrieval steps")

import math
from collections import Counter

class BM25:
    """Simple BM25 implementation for keyword search."""
    def __init__(self, documents: list[str], k1: float = 1.5, b: float = 0.75):
        self.k1 = k1
        self.b = b
        self.docs = documents
        self.doc_len = [len(d.split()) for d in documents]
        self.avg_dl = sum(self.doc_len) / len(self.doc_len)
        self.doc_count = len(documents)
        # Build inverted index: word -> set of doc indices
        self.doc_freqs: dict[str, int] = {}
        self.term_freqs: list[Counter] = []
        for doc in documents:
            tf = Counter(doc.lower().split())
            self.term_freqs.append(tf)
            for term in tf:
                self.doc_freqs[term] = self.doc_freqs.get(term, 0) + 1

    def score(self, query: str, top_k: int = 10) -> list[tuple[int, float]]:
        """Score all documents against query, return top-k (index, score) pairs."""
        query_terms = query.lower().split()
        scores = []
        for i in range(self.doc_count):
            s = 0.0
            for term in query_terms:
                if term not in self.doc_freqs:
                    continue
                df = self.doc_freqs[term]
                idf = math.log((self.doc_count - df + 0.5) / (df + 0.5) + 1)
                tf = self.term_freqs[i].get(term, 0)
                numerator = tf * (self.k1 + 1)
                denominator = tf + self.k1 * (1 - self.b + self.b * self.doc_len[i] / self.avg_dl)
                s += idf * numerator / denominator
            scores.append((i, s))
        scores.sort(key=lambda x: x[1], reverse=True)
        return scores[:top_k]


def reciprocal_rank_fusion(
    ranked_lists: list[list[tuple[int, float]]],
    k: int = 60
) -> list[tuple[int, float]]:
    """Merge multiple ranked lists using RRF.

    Each ranked_list is [(doc_index, score), ...] in descending order.
    Returns fused list sorted by RRF score.
    """
    fused_scores: dict[int, float] = {}
    for ranked_list in ranked_lists:
        for rank, (doc_idx, _) in enumerate(ranked_list):
            fused_scores[doc_idx] = fused_scores.get(doc_idx, 0) + 1.0 / (k + rank + 1)
    result = sorted(fused_scores.items(), key=lambda x: x[1], reverse=True)
    return result


def hybrid_search(
    query: str,
    bm25: BM25,
    vector_store,        # Your ChromaDB collection from M09
    embedding_fn,        # Function to embed text -> vector
    top_k: int = 10,
    bm25_weight: float = 0.5
) -> list[tuple[int, float]]:
    """Run hybrid search: BM25 + vector, fused with RRF."""
    try:
        # BM25 keyword search
        bm25_results = bm25.score(query, top_k=top_k * 2)

        # Vector similarity search (using ChromaDB from M09)
        query_embedding = embedding_fn(query)
        vector_results_raw = vector_store.query(
            query_embeddings=[query_embedding],
            n_results=top_k * 2
        )
        # Convert to (index, score) format
        vector_results = [
            (int(doc_id), 1.0 - dist)  # Convert distance to similarity
            for doc_id, dist in zip(
                vector_results_raw["ids"][0],
                vector_results_raw["distances"][0]
            )
        ]

        # Fuse with RRF
        fused = reciprocal_rank_fusion([bm25_results, vector_results])
        return fused[:top_k]

    except Exception as e:
        print(f"Hybrid search error: {e}")
        # Fallback to vector-only search
        return vector_results[:top_k] if vector_results else []

import anthropic
import json

client = anthropic.Anthropic()  # reads ANTHROPIC_API_KEY env var

async def rerank_with_claude(
    query: str,
    documents: list[dict],   # [{"id": "...", "text": "..."}, ...]
    top_k: int = 5,
    timeout: float = 30.0
) -> list[dict]:
    """Re-rank documents using Claude as a cross-encoder.

    Returns top_k documents sorted by Claude's relevance score.
    """
    if not documents:
        return []

    # Build the scoring prompt
    doc_list = "\n\n".join(
        f"[Document {i+1}] (ID: {doc['id']})\n{doc['text'][:500]}"
        for i, doc in enumerate(documents)
    )

    try:
        response = client.messages.create(
            model="claude-sonnet-4-6",
            max_tokens=1024,
            messages=[{
                "role": "user",
                "content": f"""Score each document's relevance to the query on a scale of 1-10.
Use the FULL scale: 1=completely irrelevant, 5=somewhat related, 10=directly answers the query.

Query: {query}

{doc_list}

Respond with ONLY a JSON array of objects, each with "doc_id" and "score" fields.
Example: [{{"doc_id": "abc", "score": 8}}, {{"doc_id": "def", "score": 3}}]"""
            }]
        )

        # Parse Claude's scoring response
        scores = json.loads(response.content[0].text)

        # Build lookup and sort
        score_map = {s["doc_id"]: s["score"] for s in scores}
        ranked = sorted(
            documents,
            key=lambda d: score_map.get(d["id"], 0),
            reverse=True
        )
        return ranked[:top_k]

    except json.JSONDecodeError as e:
        print(f"Re-rank parse error: {e}. Returning original order.")
        return documents[:top_k]
    except anthropic.APITimeoutError:
        print("Re-rank timeout. Returning original order.")
        return documents[:top_k]
    except Exception as e:
        print(f"Re-rank error: {e}. Returning original order.")
        return documents[:top_k]

async def hyde_transform(
    query: str,
    domain_context: str = "UCC filings and commercial lending"
) -> str:
    """Generate a hypothetical document that answers the query.

    The hypothesis is embedded instead of the raw query,
    producing better vector matches against real documents.
    """
    try:
        response = client.messages.create(
            model="claude-sonnet-4-6",
            max_tokens=300,
            messages=[{
                "role": "user",
                "content": f"""Write a short paragraph (3-5 sentences) that would appear in a
document about {domain_context} and directly answers this question:

Question: {query}

Write as if this paragraph is from an actual document, not as a response
to a user. Use technical language appropriate for the domain.
Do NOT start with "Here is" or similar preamble."""
            }]
        )
        return response.content[0].text
    except Exception as e:
        print(f"HyDE generation failed: {e}. Using original query.")
        return query


async def advanced_rag_query(
    query: str,
    bm25: BM25,
    vector_store,
    embedding_fn,
    documents: list[dict],
    top_k: int = 5
) -> list[dict]:
    """Full advanced RAG pipeline: HyDE + Hybrid Search + Re-rank."""
    # Step 1: Generate hypothetical answer for better embedding
    hypothesis = await hyde_transform(query)

    # Step 2: Run hybrid search with BOTH original query and hypothesis
    original_results = hybrid_search(query, bm25, vector_store, embedding_fn, top_k=top_k * 2)
    hyde_embedding = embedding_fn(hypothesis)
    hyde_results_raw = vector_store.query(
        query_embeddings=[hyde_embedding], n_results=top_k * 2
    )
    hyde_results = [
        (int(doc_id), 1.0 - dist)
        for doc_id, dist in zip(
            hyde_results_raw["ids"][0], hyde_results_raw["distances"][0]
        )
    ]

    # Step 3: Fuse all three result sets (BM25, vector, HyDE-vector)
    bm25_only = bm25.score(query, top_k=top_k * 2)
    all_fused = reciprocal_rank_fusion([
        bm25_only,
        [(idx, sc) for idx, sc in original_results],
        hyde_results
    ])

    # Step 4: Gather top candidates as documents
    candidate_docs = []
    seen = set()
    for doc_idx, _ in all_fused[:top_k * 2]:
        if doc_idx not in seen and doc_idx < len(documents):
            candidate_docs.append(documents[doc_idx])
            seen.add(doc_idx)

    # Step 5: Re-rank with Claude
    reranked = await rerank_with_claude(query, candidate_docs, top_k=top_k)
    return reranked

from dataclasses import dataclass

@dataclass
class TestCase:
    query: str
    relevant_doc_ids: list[str]   # Ground truth: which docs answer this
    expected_answer_fragments: list[str]  # Key phrases the answer should contain

@dataclass
class EvalResult:
    precision: float    # Retrieved relevant / total retrieved
    recall: float       # Retrieved relevant / total relevant
    faithfulness: float # Claims grounded in context / total claims

def evaluate_retrieval(
    retrieved_ids: list[str],
    relevant_ids: list[str],
    k: int = 5
) -> tuple[float, float]:
    """Compute precision@k and recall@k."""
    top_k = set(retrieved_ids[:k])
    relevant = set(relevant_ids)

    hits = len(top_k & relevant)
    precision = hits / k if k > 0 else 0.0
    recall = hits / len(relevant) if relevant else 0.0
    return precision, recall

async def evaluate_faithfulness(
    answer: str,
    retrieved_chunks: list[str]
) -> float:
    """Use Claude to check if the answer is grounded in retrieved context."""
    context = "\n---\n".join(retrieved_chunks)
    try:
        response = client.messages.create(
            model="claude-sonnet-4-6",
            max_tokens=512,
            messages=[{
                "role": "user",
                "content": f"""Evaluate whether the answer below is fully supported by the context.
For each factual claim in the answer, check if it appears in the context.

Context:
{context}

Answer:
{answer}

Respond with JSON: {{"total_claims": N, "grounded_claims": M, "ungrounded": ["list of claims not in context"]}}"""
            }]
        )
        result = json.loads(response.content[0].text)
        total = result.get("total_claims", 1)
        grounded = result.get("grounded_claims", 0)
        return grounded / total if total > 0 else 0.0
    except Exception as e:
        print(f"Faithfulness eval error: {e}")
        return 0.0

async def run_evaluation(
    test_cases: list[TestCase],
    search_fn,  # Function: query -> (retrieved_doc_ids, answer, chunks)
    k: int = 5
) -> dict:
    """Run full evaluation across test cases."""
    precisions, recalls, faiths = [], [], []

    for tc in test_cases:
        retrieved_ids, answer, chunks = await search_fn(tc.query)
        p, r = evaluate_retrieval(retrieved_ids, tc.relevant_doc_ids, k)
        f = await evaluate_faithfulness(answer, chunks)
        precisions.append(p)
        recalls.append(r)
        faiths.append(f)
        print(f"  Query: {tc.query[:50]}... P={p:.2f} R={r:.2f} F={f:.2f}")

    return {
        "avg_precision": sum(precisions) / len(precisions),
        "avg_recall": sum(recalls) / len(recalls),
        "avg_faithfulness": sum(faiths) / len(faiths),
        "n_queries": len(test_cases)
    }

# Navigate to your M09 rag-lab directory
cd rag-lab

# Activate virtual environment
source venv/bin/activate          # macOS/Linux
# venv\Scripts\activate           # Windows

# Install dependencies (if not already installed)
pip install anthropic>=0.30.0 chromadb>=0.4.0

# Set your API key
export ANTHROPIC_API_KEY=your-key-here          # macOS/Linux
# set ANTHROPIC_API_KEY=your-key-here           # Windows CMD
# $env:ANTHROPIC_API_KEY="your-key-here"        # Windows PowerShell

import os
import glob
import json
import math
from collections import Counter
import anthropic
import chromadb

client = anthropic.Anthropic()

# ── Document Loading & Chunking (from M09) ──────────────────
def load_documents(docs_dir: str) -> list[dict]:
    docs = []
    for pattern in ["*.md", "*.txt"]:
        for path in glob.glob(os.path.join(docs_dir, pattern)):
            try:
                with open(path, "r", encoding="utf-8") as f:
                    docs.append({"content": f.read(), "source": os.path.basename(path)})
            except (IOError, UnicodeDecodeError):
                pass
    if not docs:
        raise FileNotFoundError(f"No documents found in {docs_dir}")
    return docs

def chunk_text(text: str, chunk_size: int = 500, overlap: int = 50) -> list[str]:
    if len(text) <= chunk_size:
        return [text]
    chunks, start = [], 0
    while start < len(text):
        end = start + chunk_size
        chunk = text[start:end]
        if end < len(text):
            for sep in ["\n\n", "\n", ". ", " "]:
                last_sep = chunk.rfind(sep)
                if last_sep > chunk_size * 0.5:
                    end = start + last_sep + len(sep)
                    chunk = text[start:end]
                    break
        chunks.append(chunk.strip())
        start = end - overlap
    return [c for c in chunks if c]

# ── BM25 Keyword Search ─────────────────────────────────────
class BM25:
    def __init__(self, documents: list[str], k1=1.5, b=0.75):
        self.k1, self.b = k1, b
        self.docs = documents
        self.doc_len = [len(d.split()) for d in documents]
        self.avg_dl = sum(self.doc_len) / len(self.doc_len)
        self.doc_count = len(documents)
        self.doc_freqs = {}
        self.term_freqs = []
        for doc in documents:
            tf = Counter(doc.lower().split())
            self.term_freqs.append(tf)
            for term in tf:
                self.doc_freqs[term] = self.doc_freqs.get(term, 0) + 1

    def score(self, query: str, top_k: int = 10) -> list[tuple[int, float]]:
        query_terms = query.lower().split()
        scores = []
        for i in range(self.doc_count):
            s = 0.0
            for term in query_terms:
                if term not in self.doc_freqs:
                    continue
                df = self.doc_freqs[term]
                idf = math.log((self.doc_count - df + 0.5) / (df + 0.5) + 1)
                tf = self.term_freqs[i].get(term, 0)
                s += idf * tf * (self.k1 + 1) / (tf + self.k1 * (1 - self.b + self.b * self.doc_len[i] / self.avg_dl))
            scores.append((i, s))
        scores.sort(key=lambda x: x[1], reverse=True)
        return scores[:top_k]

# ── Reciprocal Rank Fusion ───────────────────────────────────
def rrf_fuse(ranked_lists: list[list[tuple[int, float]]], k: int = 60) -> list[tuple[int, float]]:
    fused = {}
    for ranked_list in ranked_lists:
        for rank, (doc_idx, _) in enumerate(ranked_list):
            fused[doc_idx] = fused.get(doc_idx, 0) + 1.0 / (k + rank + 1)
    return sorted(fused.items(), key=lambda x: x[1], reverse=True)

# ── Claude-Based Re-Ranking ─────────────────────────────────
def rerank_with_claude(query: str, chunks: list[dict], top_n: int = 5) -> list[dict]:
    if not chunks:
        return []
    docs_text = "\n".join(f"[Doc {i+1}] {c['text'][:300]}" for i, c in enumerate(chunks))
    prompt = (
        f"Query: {query}\n\nDocuments:\n{docs_text}\n\n"
        f"Score each document 1-10 for relevance to the query. "
        f"1 = completely irrelevant, 10 = directly answers the question.\n"
        f"Output ONLY the raw JSON array, no markdown fences, no commentary.\n"
        f"Format: [{{\"doc\": 1, \"score\": 8.5}}, ...]"
    )
    try:
        response = client.messages.create(
            model="claude-sonnet-4-6", max_tokens=512,
            messages=[{"role": "user", "content": prompt}],
        )
        raw = response.content[0].text.strip()
        # Strip common markdown code fences if Claude wraps the JSON
        if raw.startswith("```"):
            raw = raw.split("```")[1]
            if raw.startswith("json"):
                raw = raw[4:]
            raw = raw.strip()
        scores = json.loads(raw)
        scored = []
        for s in scores:
            idx = s["doc"] - 1
            if 0 <= idx < len(chunks):
                chunks[idx]["rerank_score"] = s["score"]
                scored.append(chunks[idx])
        scored.sort(key=lambda x: x.get("rerank_score", 0), reverse=True)
        return scored[:top_n]
    except Exception as e:
        print(f"  Re-ranking failed: {e}, using original order")
        return chunks[:top_n]

# ── Full Advanced Pipeline ───────────────────────────────────
def ingest_and_build(docs_dir="./docs"):
    """Load, chunk, store in ChromaDB, build BM25 index."""
    docs = load_documents(docs_dir)
    all_chunks = []
    for doc in docs:
        for i, chunk in enumerate(chunk_text(doc["content"])):
            all_chunks.append({"text": chunk, "source": doc["source"], "index": i})

    chroma = chromadb.Client()
    collection = chroma.get_or_create_collection("advanced_rag", metadata={"hnsw:space": "cosine"})
    collection.add(
        ids=[f"chunk_{i}" for i in range(len(all_chunks))],
        documents=[c["text"] for c in all_chunks],
        metadatas=[{"source": c["source"], "index": c["index"]} for c in all_chunks],
    )
    bm25 = BM25([c["text"] for c in all_chunks])
    print(f"  Loaded {len(docs)} docs, {len(all_chunks)} chunks")
    return collection, bm25, all_chunks

def advanced_query(query, collection, bm25, all_chunks, verbose=True):
    """Hybrid search + RRF + re-ranking + Claude generation."""
    # 1. BM25 search
    bm25_results = bm25.score(query, top_k=10)

    # 2. Vector search
    vector_raw = collection.query(query_texts=[query], n_results=10)
    vector_results = [(int(doc_id.split("_")[1]), 1.0 - dist)
                      for doc_id, dist in zip(vector_raw["ids"][0], vector_raw["distances"][0])]

    # 3. RRF fusion
    fused = rrf_fuse([bm25_results, vector_results])
    fused_chunks = [all_chunks[idx] for idx, _ in fused[:10]]

    if verbose:
        print(f"\n  Query: {query}")
        print(f"  BM25 top-3: {[all_chunks[i]['source'] for i, _ in bm25_results[:3]]}")
        print(f"  Vector top-3: {[all_chunks[i]['source'] for i, _ in vector_results[:3]]}")
        print(f"  Fused top-3: {[all_chunks[i]['source'] for i, _ in fused[:3]]}")

    # 4. Re-rank top candidates
    reranked = rerank_with_claude(query, fused_chunks, top_n=3)

    if verbose and reranked:
        print(f"  Re-ranked top-3: {[(c['source'], c.get('rerank_score', '?')) for c in reranked]}")

    # 5. Generate answer
    context = "\n\n---\n\n".join(
        f"[Source {i+1}: {c['source']}]\n{c['text']}" for i, c in enumerate(reranked)
    )
    response = client.messages.create(
        model="claude-sonnet-4-6", max_tokens=1024,
        system="Answer based ONLY on the provided context. Cite sources using [Source N].",
        messages=[{"role": "user", "content": f"Context:\n{context}\n\nQuestion: {query}"}],
    )
    return response.content[0].text

# ── Run Tests ────────────────────────────────────────────────
if __name__ == "__main__":
    print("── Advanced RAG Pipeline ──")
    collection, bm25, all_chunks = ingest_and_build()

    questions = [
        "What are the high risk indicators for a UCC filing?",
        "How long is a UCC-1 effective and what happens when it lapses?",
        "What is the difference between a UCC-1 and a UCC-3?",
    ]

    for q in questions:
        answer = advanced_query(q, collection, bm25, all_chunks)
        print(f"\n  Answer: {answer[:200]}...")
        print(f"  {'═' * 50}")