scripts/run_eval.py

"""Run the full eval harness against the RAG system.

Usage:
    uv run python scripts/run_eval.py --groq               # FAISS-only, Groq LLM
    uv run python scripts/run_eval.py --groq --hybrid      # BM25+FAISS fusion
    uv run python scripts/run_eval.py --groq --limit 5     # quick smoke test
    uv run python scripts/run_eval.py --out data/eval_results_hybrid.json --groq --hybrid
"""
from __future__ import annotations

import argparse
import json
import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from rich.console import Console
from rich.progress import Progress, SpinnerColumn, TextColumn, TimeElapsedColumn
from rich.table import Table

from researchpath.embeddings import Embedder
from researchpath.eval import (
    EvalResult,
    evaluate_example,
    load_gold_dataset,
    results_to_json,
    summarize,
)
from researchpath.index import load_index, search
from researchpath.rag import answer as rag_answer, answer_groq
from researchpath.retrieval import HybridRetriever, Reranker

console = Console()
ROOT = Path(__file__).resolve().parents[1]
INDEX_PATH = ROOT / "data" / "index.faiss"
GOLD_PATH = ROOT / "data" / "gold_dataset.json"
DEFAULT_OUT = ROOT / "data" / "eval_results.json"


def main() -> int:
    parser = argparse.ArgumentParser(description="Evaluate baseline RAG against gold dataset.")
    parser.add_argument("--k", type=int, default=5, help="Top-k chunks to retrieve (default: 5).")
    parser.add_argument("--limit", type=int, default=None, help="Evaluate only first N examples (for quick tests).")
    parser.add_argument("--out", type=str, default=str(DEFAULT_OUT), help="Output JSON path.")
    parser.add_argument("--groq", action="store_true", help="Use Groq for RAG generation (avoids Gemini rate limits).")
    parser.add_argument("--hybrid", action="store_true", help="Use BM25+FAISS hybrid retrieval (RRF fusion).")
    parser.add_argument("--rerank", action="store_true", help="Apply cross-encoder reranking on top of retrieval.")
    parser.add_argument("--resume", action="store_true", help="Skip examples already in --out file (for retrying after rate limits).")
    args = parser.parse_args()

    if not INDEX_PATH.exists():
        console.print(f"[red]No index at {INDEX_PATH}. Run scripts/build_index.py first.[/red]")
        return 1
    if not GOLD_PATH.exists():
        console.print(f"[red]No gold dataset at {GOLD_PATH}.[/red]")
        return 1

    examples = load_gold_dataset(GOLD_PATH)
    if args.limit:
        examples = examples[: args.limit]

    completed_ids: set[str] = set()
    if args.resume:
        out_path_check = Path(args.out) if Path(args.out).is_absolute() else ROOT / args.out
        if out_path_check.exists():
            try:
                with open(out_path_check, encoding="utf-8") as f:
                    prior = json.load(f)
                completed_ids = {r["id"] for r in prior.get("results", [])}
                console.print(f"[yellow]Resuming: skipping {len(completed_ids)} already-evaluated examples.[/yellow]")
            except Exception:
                pass
        examples = [e for e in examples if e.id not in completed_ids]

    rag_fn = answer_groq if args.groq else rag_answer
    rag_provider = "Groq / llama-3.3-70b" if args.groq else "Gemini / gemini-2.5-flash-lite"
    if args.rerank:
        retrieval_mode = "BM25+FAISS+CrossEncoder rerank" if args.hybrid else "FAISS+CrossEncoder rerank"
    else:
        retrieval_mode = "BM25+FAISS (RRF)" if args.hybrid else "FAISS dense"

    console.print(f"\n[bold cyan]ResearchPath Eval Harness[/bold cyan]")
    console.print(f"  Gold examples : {len(examples)}")
    console.print(f"  Retrieval     : {retrieval_mode}  k={args.k}")
    console.print(f"  RAG provider  : {rag_provider}")
    console.print(f"  Judge         : Groq / llama-3.3-70b")
    console.print(f"  Index         : {INDEX_PATH.relative_to(ROOT)}")
    console.print()

    console.print("[dim]Loading index and embedder...[/dim]")
    index, chunks = load_index(INDEX_PATH)
    embedder = Embedder()
    hybrid = HybridRetriever(index, chunks, embedder) if args.hybrid or args.rerank else None
    if args.rerank:
        console.print("[dim]Loading cross-encoder reranker (first run downloads ~80MB)...[/dim]")
        reranker = Reranker(hybrid)
    else:
        reranker = None
    console.print("[green]Ready.[/green]\n")

    results: list[EvalResult] = []
    failures: list[str] = []
    out_path = Path(args.out) if Path(args.out).is_absolute() else ROOT / args.out
    out_path.parent.mkdir(parents=True, exist_ok=True)

    prior_results: list[dict] = []
    if args.resume and out_path.exists():
        try:
            with open(out_path, encoding="utf-8") as f:
                prior_results = json.load(f).get("results", [])
        except Exception:
            prior_results = []

    def _save_partial() -> None:
        if not results and not prior_results:
            return
        partial_summary = summarize(results) if results else None
        partial_payload = (
            results_to_json(results, partial_summary)
            if partial_summary
            else {"summary": {}, "results": []}
        )
        # Merge resumed prior results with new ones (prior first, preserves order)
        if prior_results:
            new_ids = {r["id"] for r in partial_payload["results"]}
            merged = [r for r in prior_results if r["id"] not in new_ids] + partial_payload["results"]
            partial_payload["results"] = merged
            partial_payload["summary"]["n"] = len(merged)
        with open(out_path, "w", encoding="utf-8") as f:
            json.dump(partial_payload, f, indent=2, ensure_ascii=False)

    with Progress(
        SpinnerColumn(),
        TextColumn("[progress.description]{task.description}"),
        TimeElapsedColumn(),
        console=console,
    ) as progress:
        task = progress.add_task("Evaluating...", total=len(examples))

        for ex in examples:
            progress.update(task, description=f"[cyan]{ex.id}[/cyan] — {ex.question[:55]}...")
            try:
                t0 = time.time()
                if reranker:
                    hits = reranker.search(ex.question, k=args.k)
                elif hybrid:
                    hits = hybrid.search(ex.question, k=args.k)
                else:
                    hits = search(index, chunks, embedder, ex.question, k=args.k)
                ra = rag_fn(ex.question, hits)
                latency = time.time() - t0
                if not args.groq:
                    time.sleep(3)  # stay under Gemini's 20 RPM free-tier ceiling

                result = evaluate_example(ex, hits, ra, latency)
                results.append(result)
                _save_partial()  # checkpoint after every example so 429s don't lose work

                status = "[green]OK[/green]" if result.answer_correct else "[yellow]MISS[/yellow]"
                recall_str = f"recall={result.retrieval_recall:.0%}"
                progress.console.print(
                    f"  {status}  {ex.id:20s}  {recall_str}  cite={'Y' if result.citation_present else 'N'}  "
                    f"correct={'Y' if result.answer_correct else 'N'}  {latency:.1f}s"
                )
            except Exception as exc:
                failures.append(f"{ex.id}: {exc}")
                progress.console.print(f"  [red]ERR[/red]  {ex.id}: {exc}")

            progress.advance(task)

    if not results:
        console.print("[red]No results collected — check errors above.[/red]")
        return 1

    summary = summarize(results)
    summary.print_table()

    payload = results_to_json(results, summary)
    with open(out_path, "w", encoding="utf-8") as f:
        json.dump(payload, f, indent=2, ensure_ascii=False)
    try:
        display_path = out_path.relative_to(ROOT)
    except ValueError:
        display_path = out_path
    console.print(f"[green]Saved detailed results to {display_path}[/green]")

    if failures:
        console.print(f"\n[red]{len(failures)} example(s) failed:[/red]")
        for msg in failures:
            console.print(f"  {msg}")

    _print_failures_table(results, console)
    return 0


def _print_failures_table(results: list[EvalResult], console: Console) -> None:
    misses = [r for r in results if not r.answer_correct]
    if not misses:
        console.print("[bold green]All examples answered correctly.[/bold green]")
        return

    console.print(f"\n[bold yellow]Incorrect answers ({len(misses)}/{len(results)}):[/bold yellow]")
    table = Table(show_header=True, header_style="bold")
    table.add_column("ID", style="cyan", width=22)
    table.add_column("Difficulty", width=8)
    table.add_column("Recall", width=7)
    table.add_column("Cite", width=5)
    table.add_column("Expected key claim (truncated)", width=55)

    for r in misses:
        table.add_row(
            r.example.id,
            r.example.difficulty,
            f"{r.retrieval_recall:.0%}",
            "Y" if r.citation_present else "N",
            r.example.expected_key_claim[:55],
        )
    console.print(table)


if __name__ == "__main__":
    sys.exit(main())