src/graphtrace/runner.py

"""Experiment runner: predictions (resumable), evaluation, failure examples.

Layout of ``runs/<name>/``::

    manifest.json        config snapshot, prompt version, packages, git commit, model specs
    predictions.jsonl    one MethodOutput per (method, item), with full trace
    judgments.jsonl      evaluation-judge labels for non-gold cited chunks (optional)
    per_item.csv         every metric for every (method, item)
    metrics.json         aggregates with bootstrap CIs, per-category tables, paired tests, calibration
    failures.jsonl       automatically selected failure examples per method
"""

from __future__ import annotations

import csv
import hashlib
import json
import math
import threading
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

from .bench import read_items
from .config import ExperimentConfig, ProjectConfig, load_models, make_client, manifest, write_json
from .evaluation.metrics import by_key, calibration, item_metrics
from .evaluation.stats import bootstrap_ci, paired_test
from .graph import KG
from .linking import EntityLinker
from .llm import ModelSpec
from .methods import MethodConfig, Resources, run_method
from .pipeline import load_corpus, read_jsonl
from .prompts import PROMPT_VERSION
from .schemas import BenchItem, MethodOutput

METRICS = [
    "correct",
    "decision_correct",
    "answer_f1",
    "answer_exact",
    "path_valid_fraction",
    "path_all_valid",
    "path_completeness",
    "citation_rate",
    "grounding_precision",
    "false_premise_rejected",
    "conflict_flagged",
    "insufficient_abstained",
    "answered",
    "errored",
    "context_tokens",
    "calls",
    "input_tokens",
    "output_tokens",
    "cost_usd",
    "llm_latency_s",
]


# Headline = categories whose gold does not rely on the result-record agreement rule that the
# GraphTrace verifier also uses. C1/C2 are reported separately (protocol, threats to validity).
HEADLINE_CATEGORIES = ("single_hop", "multi_hop", "false_premise")


def method_hash(method: MethodConfig, base_spec: ModelSpec, project: ProjectConfig, exp: ExperimentConfig) -> str:
    """Identity of everything that determines a method's output for an item. Predictions
    with a different hash are stale and are re-run."""
    payload = {
        "method": method.model_dump(),
        "base_model": base_spec.model_dump(),
        "prompt_version": PROMPT_VERSION,
        "nli": exp.nli_model if exp.nli_verifier else None,
        "min_edge_confidence": project.bench.min_edge_confidence,
    }
    return hashlib.sha256(json.dumps(payload, sort_keys=True, default=str).encode("utf-8")).hexdigest()[:16]


def load_bench(exp: ExperimentConfig) -> list[BenchItem]:
    items = read_items(exp.bench_file)
    unreviewed = [i for i in items if i.review_status not in ("accepted", "edited")]
    if unreviewed and not exp.allow_unreviewed:
        raise ValueError(
            f"{len(unreviewed)} items in {exp.bench_file} are not human-accepted. Freeze the benchmark first "
            "(`graphtrace freeze`) or set allow_unreviewed: true for a smoke test that must not be reported."
        )
    return items[: exp.limit] if exp.limit else items


def resources(project: ProjectConfig, exp: ExperimentConfig, offline: bool = False) -> Resources:
    papers, chunks = load_corpus(project)
    kg = KG.load(project.paths.graph, project.bench.min_edge_confidence)
    nli = None
    if exp.nli_verifier:
        from .verification import NLIVerifier

        nli = NLIVerifier(exp.nli_model)
    return Resources(
        base=make_client(exp.base_model, project, offline=offline),
        kg=kg,
        linker=EntityLinker(kg),
        chunks={c.chunk_id: c for c in chunks},
        titles={d: p.title for d, p in papers.items()},
        nli=nli,
    )


def run_experiment(project: ProjectConfig, exp: ExperimentConfig, offline: bool = False, res: Resources | None = None) -> Path:
    run_dir = Path(project.paths.runs) / exp.name
    run_dir.mkdir(parents=True, exist_ok=True)
    items = load_bench(exp)
    res = res or resources(project, exp, offline)
    models = load_models(project.models_file) if Path(project.models_file).exists() else {}
    hashes = {m.name: method_hash(m, res.base.spec, project, exp) for m in exp.methods}
    write_json(
        run_dir / "manifest.json",
        manifest(
            {
                "experiment": exp.model_dump(),
                "model_specs": {k: v.model_dump() for k, v in models.items() if k in (exp.base_model, exp.judge_model)},
                "n_items": len(items),
                "method_config_hashes": hashes,
                "smoke_test_unreviewed_items": exp.allow_unreviewed,
            }
        ),
    )
    pred_path = run_dir / "predictions.jsonl"
    # Completed = same method name AND same config hash AND no error.
    done = {(r["method"], r["item_id"], r.get("config_hash")) for r in read_jsonl(pred_path) if not r.get("error")}
    todo = [(m, it) for m in exp.methods for it in items if (m.name, it.item_id, hashes[m.name]) not in done]
    lock = threading.Lock()

    def one(t: tuple[MethodConfig, BenchItem]) -> MethodOutput:
        return run_method(t[0], t[1], res).model_copy(update={"config_hash": hashes[t[0].name]})

    with pred_path.open("a", encoding="utf-8") as f, ThreadPoolExecutor(max_workers=exp.workers) as pool:
        for out in pool.map(one, todo):
            with lock:
                f.write(out.model_dump_json() + "\n")
                f.flush()
    return run_dir


def latest_predictions(run_dir: Path, hashes: dict[str, str] | None = None) -> list[MethodOutput]:
    """Latest prediction per (method, item); with ``hashes``, only rows of the current config."""
    latest: dict[tuple[str, str], dict] = {}
    for r in read_jsonl(run_dir / "predictions.jsonl"):
        if hashes is not None and r.get("config_hash") != hashes.get(r["method"]):
            continue
        latest[(r["method"], r["item_id"])] = r  # later lines (successful re-runs) win
    return [MethodOutput(**r) for r in latest.values()]


def _agg(rows: list[dict], seed: int = 0) -> dict:
    out = {"n": len(rows), "errors": sum(1 for r in rows if r.get("errored"))}
    for m in METRICS:
        vals = [r[m] for r in rows if r.get(m) is not None and not (isinstance(r[m], float) and math.isnan(r[m]))]
        mean, lo, hi = bootstrap_ci(vals, seed=seed)
        out[m] = {"mean": mean, "low": lo, "high": hi, "n": len(vals)}
    return out


def evaluate_run(project: ProjectConfig, exp: ExperimentConfig, run_judge: bool = True, offline: bool = False) -> dict:
    run_dir = Path(project.paths.runs) / exp.name
    items = {i.item_id: i for i in load_bench(exp)}
    kg = KG.load(project.paths.graph, project.bench.min_edge_confidence)
    linker = EntityLinker(kg)
    man_path = run_dir / "manifest.json"
    hashes = json.loads(man_path.read_text(encoding="utf-8")).get("method_config_hashes") if man_path.exists() else None
    outputs = [o for o in latest_predictions(run_dir, hashes) if o.item_id in items]

    judged: dict[tuple[str, str], set[str]] = defaultdict(set)
    jpath = run_dir / "judgments.jsonl"
    reused = "; reused existing judgments.jsonl" if jpath.exists() else ""
    if not run_judge:
        judge_status = "skipped (--no-judge)" + reused
    elif not exp.judge_model:
        judge_status = "skipped (no judge_model configured)" + reused
    else:
        judge_status = f"ran ({exp.judge_model})"
    if run_judge and exp.judge_model:
        from .evaluation.judge import LLMJudge, judge_outputs

        _, chunks = load_corpus(project)
        rows = judge_outputs(outputs, items, {c.chunk_id: c for c in chunks}, LLMJudge(make_client(exp.judge_model, project, offline=offline)))
        with jpath.open("w", encoding="utf-8") as f:
            for r in rows:
                f.write(json.dumps(r, ensure_ascii=False) + "\n")
    for r in read_jsonl(jpath):
        if r.get("label") == "supported":
            judged[(r["method"], r["item_id"])].add(r["chunk_id"])

    rows = [item_metrics(o, items[o.item_id], kg, linker, judged.get((o.method, o.item_id)), exp.min_support) for o in outputs]
    rows.sort(key=lambda r: (r["method"], r["item_id"]))
    with (run_dir / "per_item.csv").open("w", newline="", encoding="utf-8") as f:
        w = csv.DictWriter(f, fieldnames=list(rows[0].keys()) if rows else ["item_id"])
        w.writeheader()
        w.writerows(rows)

    per_method = by_key(rows, "method")
    metrics: dict = {
        "methods": {},
        "headline": {},
        "conflict_insufficient": {},
        "by_category": {},
        "calibration": {},
        "paired_vs_reference": {},
        "headline_categories": list(HEADLINE_CATEGORIES),
        "judge_used": bool(jpath.exists()),
        "judge_status": judge_status,
    }
    for m, rs in per_method.items():
        metrics["methods"][m] = _agg(rs)
        metrics["headline"][m] = _agg([r for r in rs if r["category"] in HEADLINE_CATEGORIES])
        metrics["conflict_insufficient"][m] = _agg([r for r in rs if r["category"] not in HEADLINE_CATEGORIES])
        metrics["by_category"][m] = {c: _agg(crs) for c, crs in by_key(rs, "category").items()}
        metrics["calibration"][m] = calibration(rs)
    ref = per_method.get(exp.reference_method, [])
    ref_by_item = {r["item_id"]: r for r in ref}
    for m, rs in per_method.items():
        if m == exp.reference_method:
            continue
        # Paired tests on the headline categories only.
        pair = [(r, ref_by_item[r["item_id"]]) for r in rs if r["item_id"] in ref_by_item and r["category"] in HEADLINE_CATEGORIES]
        metrics["paired_vs_reference"][m] = {
            k: paired_test([a[k] for a, _ in pair], [b[k] for _, b in pair])
            for k in (
                "correct",
                "path_valid_fraction",
                "path_all_valid",
                "citation_rate",
                "grounding_precision",
                "false_premise_rejected",
                "answer_f1",
            )
        }
    write_json(run_dir / "metrics.json", metrics)
    write_failures(run_dir, outputs, items, rows)
    return metrics


def write_failures(run_dir: Path, outputs: list[MethodOutput], items: dict[str, BenchItem], rows: list[dict], per_kind: int = 5) -> None:
    """Pick concrete failure cases per method, deterministically (sorted by item ID)."""
    out_by = {(o.method, o.item_id): o for o in outputs}
    kinds = {
        "accepted_false_premise": lambda r: r["category"] == "false_premise" and r["decision"] == "answer",
        "missed_conflict": lambda r: r["template_id"] == "C1" and r["decision"] != "flag_conflict",
        "wrong_answer_with_invalid_path": lambda r: r["gold_decision"] == "answer" and r["correct"] == 0 and r["path_all_valid"] == 0,
        "correct_answer_invalid_path": lambda r: r["correct"] == 1 and r["path_all_valid"] == 0,
        "over_abstention": lambda r: r["gold_decision"] == "answer" and r["decision"] == "abstain",
        "error": lambda r: bool(r["error"]),
    }
    with (run_dir / "failures.jsonl").open("w", encoding="utf-8") as f:
        for method, rs in sorted(by_key(rows, "method").items()):
            for kind, pred in kinds.items():
                for r in [r for r in rs if pred(r)][:per_kind]:
                    o = out_by[(method, r["item_id"])]
                    it = items[r["item_id"]]
                    f.write(
                        json.dumps(
                            {
                                "method": method,
                                "kind": kind,
                                "item_id": it.item_id,
                                "question": it.question,
                                "gold_decision": it.gold_decision,
                                "gold_answer": it.gold_answer,
                                "decision": o.decision,
                                "answer": o.answer,
                                "confidence": o.confidence,
                                "graph_path": o.graph_path,
                                "source_chunk_ids": o.source_chunk_ids,
                                "verifier": o.trace.get("verifier"),
                            },
                            ensure_ascii=False,
                        )
                        + "\n"
                    )


def estimate_cost(exp: ExperimentConfig, project: ProjectConfig, n_items: int, avg_output_tokens: int = 350) -> dict:
    """Upper-bound style estimate: one answer call per (method, item) using the full budget,
    plus a retry for half of GraphTrace's items, plus ~600 tokens of instructions per call."""
    spec = load_models(project.models_file)[exp.base_model]
    calls = 0
    tokens_in = 0
    for m in exp.methods:
        per = 1.5 if m.kind == "graphtrace" else 1.0
        calls += int(n_items * per)
        tokens_in += int(n_items * per * (m.context_budget + 600))
    tokens_out = calls * avg_output_tokens
    cost = (tokens_in * spec.price_in_per_mtok + tokens_out * spec.price_out_per_mtok) / 1e6
    return {"model": spec.model, "calls": calls, "input_tokens": tokens_in, "output_tokens": tokens_out, "usd": round(cost, 2)}