kvl/modules/retrieval.py

"""Module B: Retrieval Utility — measures how well the document supports information retrieval."""

from __future__ import annotations
import json
import numpy as np
import anthropic
from kvl.ingestor import Document

_QUERY_PROMPT = """Generate {n} realistic search queries that a user would type to find information in this document.
Queries should be specific, not generic. Mix question and keyword styles.

Return ONLY a JSON array of strings.

Example: ["What is the impact of X on Y?", "methodology for measuring Z", ...]

Document excerpt:
{excerpt}"""


def _call_claude(client: anthropic.Anthropic, prompt: str) -> str:
    msg = client.messages.create(
        model="claude-haiku-4-5-20251001",
        max_tokens=1024,
        messages=[{"role": "user", "content": prompt}],
    )
    return msg.content[0].text.strip()


def _generate_queries(client: anthropic.Anthropic, doc: Document, n: int = 8) -> list[str]:
    excerpt = " ".join(doc.raw.split()[:3000])
    raw = _call_claude(client, _QUERY_PROMPT.format(n=n, excerpt=excerpt))
    raw = raw.strip()
    if raw.startswith("```"):
        raw = "\n".join(raw.split("\n")[1:])
        raw = raw.rsplit("```", 1)[0]
    try:
        queries = json.loads(raw)
        return [q for q in queries if isinstance(q, str)][:n]
    except json.JSONDecodeError:
        return []


def _build_index(chunks: list, embedder) -> tuple:
    """Embed all chunks and build a FAISS flat index."""
    import faiss
    texts = [c.text for c in chunks]
    embeddings = embedder.encode(texts, normalize_embeddings=True, show_progress_bar=False)
    embeddings = np.array(embeddings, dtype="float32")
    dim = embeddings.shape[1]
    index = faiss.IndexFlatIP(dim)  # inner product = cosine when normalized
    index.add(embeddings)
    return index, embeddings


def _retrieve(query: str, index, embedder, k: int = 3) -> list[int]:
    """Return top-k chunk indices for a query."""
    q_emb = embedder.encode([query], normalize_embeddings=True, show_progress_bar=False)
    q_emb = np.array(q_emb, dtype="float32")
    _, indices = index.search(q_emb, k)
    return indices[0].tolist()


def _recall_at_k(retrieved: list[int], relevant: int) -> float:
    return 1.0 if relevant in retrieved else 0.0


def _reciprocal_rank(retrieved: list[int], relevant: int) -> float:
    for rank, idx in enumerate(retrieved, start=1):
        if idx == relevant:
            return 1.0 / rank
    return 0.0


def evaluate(client: anthropic.Anthropic, doc: Document, embedder, progress_cb=None) -> dict:
    """Return retrieval utility score (0-100) and detailed results."""
    if not doc.chunks:
        return {"score": 0, "details": [], "summary": "No chunks to index."}

    if progress_cb:
        progress_cb("Building retrieval index...")

    index, _ = _build_index(doc.chunks, embedder)

    if progress_cb:
        progress_cb("Generating retrieval test queries...")

    queries = _generate_queries(client, doc)
    if not queries:
        return {"score": 50, "details": [], "summary": "Could not generate test queries."}

    results = []
    recall_scores = []
    mrr_scores = []

    for i, query in enumerate(queries):
        if progress_cb:
            progress_cb(f"Evaluating retrieval query {i+1}/{len(queries)}...")

        retrieved_indices = _retrieve(query, index, embedder, k=3)

        # Find the most semantically relevant chunk by re-ranking with cosine sim
        q_emb = embedder.encode([query], normalize_embeddings=True, show_progress_bar=False)
        q_emb = np.array(q_emb, dtype="float32")
        chunk_texts = [c.text for c in doc.chunks]
        chunk_embs = embedder.encode(chunk_texts, normalize_embeddings=True, show_progress_bar=False)
        sims = np.dot(chunk_embs, q_emb.T).flatten()
        best_chunk = int(np.argmax(sims))

        r_at_3 = _recall_at_k(retrieved_indices, best_chunk)
        rr = _reciprocal_rank(retrieved_indices, best_chunk)
        recall_scores.append(r_at_3)
        mrr_scores.append(rr)

        results.append({
            "query": query,
            "retrieved_chunks": retrieved_indices,
            "best_chunk": best_chunk,
            "recall_at_3": r_at_3,
            "reciprocal_rank": rr,
        })

    avg_recall = sum(recall_scores) / len(recall_scores)
    avg_mrr = sum(mrr_scores) / len(mrr_scores)
    # Weighted combination: recall@3 (60%) + MRR (40%), mapped to 0-100
    raw_score = 0.6 * avg_recall + 0.4 * avg_mrr
    score = round(raw_score * 100)

    return {
        "score": score,
        "details": results,
        "summary": f"Recall@3: {avg_recall:.2f} | MRR: {avg_mrr:.2f} across {len(queries)} queries.",
    }