evaluation/eval.py

"""Retrieval evaluation: Accuracy@k and MRR across chunking configurations."""

import itertools
import logging
from dataclasses import dataclass
from pathlib import Path

from ingestion.embedder import Embedder
from ingestion.pipeline import IngestionPipeline
from retrieval.index import VectorIndex
from retrieval.searcher import search

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Placeholder test cases — replace with real questions from your PDFs.
# Each dict needs: question, expected_source (filename), expected_page (1-based).
# ---------------------------------------------------------------------------

PLACEHOLDER_TEST_CASES: list[dict] = [
    {
        "question": "What two components does RAG combine to generate answers?",
        "expected_source": "original_rag_paper.pdf",
        "expected_page": 1,
    },
    {
        "question": "Which dataset is used to evaluate open-domain question answering in the RAG paper?",
        "expected_source": "original_rag_paper.pdf",
        "expected_page": 6,
    },
    {
        "question": "What is the role of the retriever in the RAG architecture?",
        "expected_source": "original_rag_paper.pdf",
        "expected_page": 2,
    },
]


# ---------------------------------------------------------------------------
# Metrics dataclass
# ---------------------------------------------------------------------------

@dataclass
class EvalMetrics:
    accuracy_at_1: float
    accuracy_at_3: float
    accuracy_at_5: float
    mrr: float              # Mean Reciprocal Rank
    n_queries: int

    def __str__(self) -> str:
        return (
            f"Acc@1={self.accuracy_at_1:.2f}  "
            f"Acc@3={self.accuracy_at_3:.2f}  "
            f"Acc@5={self.accuracy_at_5:.2f}  "
            f"MRR={self.mrr:.2f}  "
            f"(n={self.n_queries})"
        )


# ---------------------------------------------------------------------------
# Core evaluation
# ---------------------------------------------------------------------------

def evaluate_retrieval(
    test_cases: list[dict],
    embedder: Embedder,
    index: VectorIndex,
    k: int = 5,
) -> EvalMetrics:
    """Compute Accuracy@1/3/5 and MRR for a set of labelled questions.

    Each test case must have:
        question        (str)  — the query to embed and search
        expected_source (str)  — filename of the expected chunk (e.g. "paper.pdf")
        expected_page   (int)  — 1-based page number of the expected chunk

    A result is considered a hit when both source and page_num match the
    expected values.  MRR is 0 for queries where the expected chunk does not
    appear in the top-k results.

    Args:
        test_cases: List of labelled query dicts.
        embedder:   Embedder used at ingestion time (must be the same model).
        index:      Populated VectorIndex to evaluate against.
        k:          Maximum rank to consider (search retrieves this many results).

    Returns:
        EvalMetrics with aggregated scores.
    """
    if not test_cases:
        raise ValueError("test_cases must be non-empty")

    hits_at_1 = hits_at_3 = hits_at_5 = 0
    reciprocal_ranks: list[float] = []

    for case in test_cases:
        results = search(case["question"], embedder, index, k=k).chunks
        rank = _find_rank(
            results,
            expected_source=case["expected_source"],
            expected_page=int(case["expected_page"]),
        )

        if rank is not None:
            if rank <= 1:
                hits_at_1 += 1
            if rank <= 3:
                hits_at_3 += 1
            if rank <= 5:
                hits_at_5 += 1
            reciprocal_ranks.append(1.0 / rank)
        else:
            reciprocal_ranks.append(0.0)

    n = len(test_cases)
    return EvalMetrics(
        accuracy_at_1=hits_at_1 / n,
        accuracy_at_3=hits_at_3 / n,
        accuracy_at_5=hits_at_5 / n,
        mrr=sum(reciprocal_ranks) / n,
        n_queries=n,
    )


# ---------------------------------------------------------------------------
# Chunking strategy comparison
# ---------------------------------------------------------------------------

def compare_chunking_strategies(
    pdf_paths: list[str | Path],
    test_cases: list[dict] | None = None,
    chunk_sizes: list[int] | None = None,
    overlaps: list[int] | None = None,
) -> list[dict]:
    """Re-ingest PDFs under every (chunk_size, overlap) combination and compare retrieval metrics.

    Builds a fresh VectorIndex for each configuration so results are
    independent.  The Embedder is created once and shared across all runs.
    Only the recursive_character strategy is evaluated — varying chunk size
    and overlap is the most practically useful axis for that splitter.

    Args:
        pdf_paths:   List of PDF paths to ingest for each configuration.
        test_cases:  Labelled queries (defaults to PLACEHOLDER_TEST_CASES).
        chunk_sizes: Character lengths to try (default: [300, 500, 800]).
        overlaps:    Overlap values to try (default: [0, 50]).

    Returns:
        List of result dicts, each with keys: chunk_size, overlap, and the
        four metric fields.  Also prints a formatted comparison table.
    """
    if test_cases is None:
        test_cases = PLACEHOLDER_TEST_CASES
    if chunk_sizes is None:
        chunk_sizes = [300, 500, 800]
    if overlaps is None:
        overlaps = [0, 50]

    pdf_paths = [Path(p) for p in pdf_paths]
    missing = [p for p in pdf_paths if not p.exists()]
    if missing:
        raise FileNotFoundError(f"PDF(s) not found: {missing}")

    print("Loading embedder (shared across all runs)...")
    embedder = Embedder()

    configs = list(itertools.product(chunk_sizes, overlaps))
    rows: list[dict] = []

    for chunk_size, overlap in configs:
        label = f"chunk={chunk_size}, overlap={overlap}"
        print(f"\nIngesting with {label}...")

        index = VectorIndex(dimension=embedder.dimension)
        pipeline = IngestionPipeline(
            embedder=embedder,
            index=index,
            strategy="recursive_character",
            chunk_size=chunk_size,
            overlap=overlap,
        )

        for pdf_path in pdf_paths:
            result = pipeline.ingest_pdf(pdf_path)
            if result.error:
                logger.warning("Skipped %s: %s", pdf_path.name, result.error)
            else:
                print(f"  {pdf_path.name}: {result.chunks} chunks")

        metrics = evaluate_retrieval(test_cases, embedder, index, k=5)
        rows.append({
            "chunk_size": chunk_size,
            "overlap": overlap,
            "acc@1": metrics.accuracy_at_1,
            "acc@3": metrics.accuracy_at_3,
            "acc@5": metrics.accuracy_at_5,
            "mrr": metrics.mrr,
        })

    _print_table(rows)
    return rows


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _find_rank(results, expected_source: str, expected_page: int) -> int | None:
    """Return the 1-based rank of the first matching result, or None."""
    for rank, result in enumerate(results, start=1):
        meta = result.metadata
        source_match = Path(meta.get("source", "")).name == Path(expected_source).name
        page_match = meta.get("page_num") == expected_page
        if source_match and page_match:
            return rank
    return None


def _print_table(rows: list[dict]) -> None:
    col_w = [10, 9, 8, 8, 8, 8]
    headers = ["chunk_size", "overlap", "Acc@1", "Acc@3", "Acc@5", "MRR"]
    divider = "-" * sum(col_w)

    print(f"\n{'Chunking strategy comparison (recursive_character)':^{sum(col_w)}}")
    print(divider)
    print("".join(h.ljust(w) for h, w in zip(headers, col_w)))
    print(divider)

    for row in rows:
        values = [
            str(row["chunk_size"]),
            str(row["overlap"]),
            f"{row['acc@1']:.2f}",
            f"{row['acc@3']:.2f}",
            f"{row['acc@5']:.2f}",
            f"{row['mrr']:.2f}",
        ]
        print("".join(v.ljust(w) for v, w in zip(values, col_w)))

    print(divider)