eval/retrieval_eval.py

"""
eval/retrieval_eval.py — Retrieval evaluation: Precision@K, Recall@K, MRR

Compares vector-only, BM25-only, hybrid RRF, and hybrid+reranker
using labeled test cases.
"""

import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent))

from loguru import logger
from retrieval.embedder import Embedder
from retrieval.vectorstore import VectorStore
from retrieval.bm25_index import BM25Index
from retrieval.reranker import Reranker
from retrieval.hybrid import reciprocal_rank_fusion, retrieve


def precision_at_k(retrieved_ids: list[str], relevant_ids: set[str], k: int) -> float:
    top_k = retrieved_ids[:k]
    hits = sum(1 for rid in top_k if rid[:8] in relevant_ids or rid in relevant_ids)
    return hits / k if k > 0 else 0.0


def recall_at_k(retrieved_ids: list[str], relevant_ids: set[str], k: int) -> float:
    top_k = retrieved_ids[:k]
    hits = sum(1 for rid in top_k if rid[:8] in relevant_ids or rid in relevant_ids)
    return hits / len(relevant_ids) if relevant_ids else 0.0


def mrr(retrieved_ids: list[str], relevant_ids: set[str]) -> float:
    for i, rid in enumerate(retrieved_ids):
        if rid[:8] in relevant_ids or rid in relevant_ids:
            return 1.0 / (i + 1)
    return 0.0


def evaluate_retrieval(
    test_cases: list[dict],
    vectorstore: VectorStore,
    bm25_index: BM25Index,
    reranker: Reranker,
) -> dict:
    """
    Run retrieval evaluation across all test cases.

    Parameters
    ----------
    test_cases : list[dict]
        Each case: {"query": str, "relevant_chunk_ids": list[str]}

    Returns
    -------
    dict
        Results table for each method.
    """
    methods = {
        "Vector only": [],
        "BM25 only": [],
        "Hybrid RRF": [],
        "Hybrid + Reranker": [],
    }

    for case in test_cases:
        query = case["query"]
        relevant = set(case["relevant_chunk_ids"])

        # Vector only
        vec_results = vectorstore.query(query, top_k=20)
        vec_ids = [r.chunk_id for r in vec_results]
        methods["Vector only"].append({"ids": vec_ids, "relevant": relevant})

        # BM25 only
        bm25_results = bm25_index.query(query, top_k=20)
        bm25_ids = [r.chunk_id for r in bm25_results]
        methods["BM25 only"].append({"ids": bm25_ids, "relevant": relevant})

        # Hybrid RRF
        fused = reciprocal_rank_fusion(vec_results, bm25_results)
        fused_ids = [r.chunk_id for r in fused]
        methods["Hybrid RRF"].append({"ids": fused_ids, "relevant": relevant})

        # Hybrid + Reranker
        reranked = reranker.rerank(query, fused[:20], top_k=6)
        reranked_ids = [r.chunk_id for r in reranked]
        methods["Hybrid + Reranker"].append({"ids": reranked_ids, "relevant": relevant})

    # Aggregate metrics
    results = {}
    for method, cases in methods.items():
        p3 = sum(precision_at_k(c["ids"], c["relevant"], 3) for c in cases) / len(cases)
        p6 = sum(precision_at_k(c["ids"], c["relevant"], 6) for c in cases) / len(cases)
        r6 = sum(recall_at_k(c["ids"], c["relevant"], 6) for c in cases) / len(cases)
        m = sum(mrr(c["ids"], c["relevant"]) for c in cases) / len(cases)

        results[method] = {
            "P@3": round(p3, 3),
            "P@6": round(p6, 3),
            "R@6": round(r6, 3),
            "MRR": round(m, 3),
        }
        logger.info(f"{method}: P@3={p3:.3f} P@6={p6:.3f} R@6={r6:.3f} MRR={m:.3f}")

    return results


if __name__ == "__main__":
    # Example usage with placeholder test cases
    print("Run this after ingesting sample documents to get real metrics.")
    print("See eval/feedback_eval.py for the full evaluation pipeline.")