app_kit/eval.py

from __future__ import annotations

import json
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any

from .demo_pack import ingest_demo_pack, load_index
from .storage import init_db, connect, DEFAULT_ARTIFACTS_DIR
from .tracing import write_trace_artifact


@dataclass
class EvalResult:
    scenario_id: str
    query: str
    top_sections: list[dict[str, Any]]
    expected_section_ids: list[int]
    expected_section_titles: list[str]
    hit_top3: bool
    safety_present: bool
    sufficient: bool


def load_scenarios(pack_dir: str | Path) -> list[dict[str, Any]]:
    pack_dir = Path(pack_dir)
    with open(pack_dir / "golden_scenarios.json", "r", encoding="utf-8") as f:
        return json.load(f)


def evaluate_pack(pack_dir: str | Path, db_path: str | Path | None = None, artifact_dir: str | Path | None = None) -> dict[str, Any]:
    pack_dir = Path(pack_dir)
    init_db(db_path)
    ingest_demo_pack(pack_dir, db_path=db_path, reset=True)
    scenarios = load_scenarios(pack_dir)
    index = load_index(db_path)
    from .reasoning import build_response

    results: list[EvalResult] = []
    for scenario in scenarios:
        query = scenario["symptom"]
        if scenario.get("equipment_type"):
            query += " " + scenario["equipment_type"]
        if scenario.get("notes"):
            query += " " + scenario["notes"]
        hits = index.search_sections(query, top_k=5)
        top_ids = [hit.record_id for hit in hits[:3]]
        top_titles = [hit.title.lower() for hit in hits[:3]]
        expected_titles = [t.lower() for t in scenario.get("expected_section_titles", [])]
        expected_ids = set(scenario.get("expected_section_ids", []))
        hit_top3 = bool(expected_ids & set(top_ids)) if expected_ids else any(
            any(expected in title for title in top_titles)
            for expected in expected_titles
        )
        response_body, _, payload = build_response(
            scenario["symptom"],
            scenario.get("equipment_type", ""),
            scenario.get("location", ""),
            scenario.get("notes", ""),
            scenario.get("photo_path"),
            index,
        )
        safety_present = "Safety reminder" in response_body or "safety reminder" in response_body.lower()
        sufficient = payload.get("status") == "insufficient_evidence"
        results.append(
            EvalResult(
                scenario_id=scenario["scenario_id"],
                query=query,
                top_sections=[
                    {
                        "id": hit.record_id,
                        "title": hit.title,
                        "score": round(hit.score, 4),
                        "citation": hit.citation,
                    }
                    for hit in hits[:5]
                ],
                expected_section_ids=expected_ids and list(expected_ids) or scenario.get("expected_section_ids", []),
                expected_section_titles=scenario.get("expected_section_titles", []),
                hit_top3=hit_top3,
                safety_present=safety_present,
                sufficient=sufficient,
            )
        )

    total = len(results)
    eligible = [result for result in results if not result.sufficient]
    top3_hits = sum(1 for result in eligible if result.hit_top3)
    safety_hits = sum(1 for result in results if result.safety_present)
    insufficient_cases = sum(1 for result in results if result.sufficient)
    report = {
        "model_name": payload.get("llm_stats", {}).get("model_id", "nvidia/NeMoTRON-3-Nano-4B-Instruct"),
        "pack": str(pack_dir),
        "scenario_count": total,
        "top3_hit_rate": round(top3_hits / len(eligible) if eligible else 0.0, 3),
        "safety_presence_rate": round(safety_hits / total if total else 0.0, 3),
        "insufficient_cases": insufficient_cases,
        "results": [asdict(result) for result in results],
    }
    trace_dir = Path(artifact_dir) if artifact_dir is not None else DEFAULT_ARTIFACTS_DIR
    trace_path = write_trace_artifact(trace_dir, {"kind": "eval", **report})
    report["trace_path"] = str(trace_path)
    return report


def main() -> None:
    import argparse

    parser = argparse.ArgumentParser(description="Evaluate P3 field repair logbook golden scenarios")
    parser.add_argument("--pack", required=True, help="Path to demo pack")
    parser.add_argument("--db", default=None, help="SQLite database path")
    args = parser.parse_args()
    report = evaluate_pack(args.pack, db_path=args.db)
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()