Blame for scripts/run_eval.py

"""Run one configuration over a dataset → runs/<UTC time>_<config>_<profile>/.

    MODEL_PROFILE=dev python scripts/run_eval.py --config C1 --data data/smoke.jsonl --echo
    MODEL_PROFILE=dev python scripts/run_eval.py --config C2 --data data/smoke.jsonl --limit 3

C1-C5 all run on the one graph of orchestrator/graph.py; a configuration is pipeline.yaml plus the
flags of config/ablations/<config>_*.yaml. `solved` comes only from the acceptance predicate.
`--config PG` is the Phase 4 preview, not an experimental configuration: Planner + Generator
only (plans and sketches, no proof), one shared budget per problem:

    MODEL_PROFILE=dev python scripts/run_eval.py --config PG --data data/dev_preview.jsonl --echo

Outputs: run_config.json (RNF-2, rewritten at the end), trace.jsonl (RNF-5), results.jsonl
(one line per problem), console.log (a copy of everything printed, CLI only). The profile comes from MODEL_PROFILE (models.yaml → default_profile).
Paid profiles need the team's explicit permission (CLAUDE.md §2).

Spend (llm/spend.py, pipeline.yaml → spend): prints the project's spend so far and what is
left before the first problem, and the run's cost at the end. When a paid call could pass
spend.total_usd or spend.run_usd the problem is `budget_stop` and the run stops there
(run_config.json status `budget_stop`); the remaining problems are not run.

FPR_B (C2 only, `--check`, usually through scripts/run_baseline_b.py): every claimed proof is
formalized by baseline_b.check_claim (Autoformalizer–Verifier loop, critic.k_max attempts, its own
budget of the base pipeline.yaml budget.b_max calls, at most k_max × 2 = 6 in practice); the row gets
`check` (status, accepted, attempts, its own `llm`) and the summary `check` (claimed, verified,
not_verified, fpr_b, calls, cost). C2's own `llm` excludes it. FPR_B = not_verified / (verified +
not_verified); `budget_exhausted` counts as not_verified; a row marked `lean_down`, and checks ended
by provider_error / budget_stop / error, are left out (re-run once, docs/analysis_plan.md).

Lean (lean/runner.py): before the first problem the runner must load the base imports and answer
its probes, in the REPL or with the file back-end. If neither works, no problem is run and no LLM
call is made (status `lean_unavailable`, exit code 2); the same stop happens between problems if
both back-ends give up during the run.
"""

from __future__ import annotations

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

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

from formal_mas.agents.generator import run_generator  # noqa: E402
from formal_mas.agents.planner import run_planner  # noqa: E402
from formal_mas.agents.prompting import fragment  # noqa: E402
from formal_mas.baselines.baseline_b import check_claim  # noqa: E402
from formal_mas.config import ABLATIONS, ROOT, apply_ablation, load_models, load_pipeline  # noqa: E402
from formal_mas.lean.runner import LeanRunner  # noqa: E402
from formal_mas.llm.client import LLMClient  # noqa: E402
from formal_mas.llm.pricing import in_peak  # noqa: E402
from formal_mas.llm.spend import SpendGuard  # noqa: E402
from formal_mas.llm.specs import profile_name  # noqa: E402
from formal_mas.lean.syntax_check import SyntaxChecker  # noqa: E402
from formal_mas.orchestrator.graph import run_graph  # noqa: E402
from formal_mas.retrieval.retriever import Retriever  # noqa: E402
from formal_mas.run_config import build_run_config, run_start_info, write_run_config  # noqa: E402
from formal_mas.tracing import Budget, Tracer, iso, utc_now  # noqa: E402

CONFIGS = (*ABLATIONS, "PG")
LOG = logging.getLogger("formal_mas.run_eval")


def load_problems(path: Path, split: str | None, limit: int | None) -> list[dict]:
    rows = [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line.strip()]
    if split:
        rows = [r for r in rows if r.get("split") == split]
    return rows[:limit] if limit else rows


def summarize_check(rows: list[dict]) -> dict | None:
    """FPR_B (Eq. 16) over the claims that were checked to the end (verified or not_verified)."""
    checks = [r["check"] for r in rows if "check" in r]
    if not checks:
        return None
    status: dict[str, int] = {}
    for r in rows:  # Lean lost during the row: its check may have failed for Lean, not for the proof
        if "check" in r:
            key = "lean_down" if r.get("lean_down") and r["check"]["status"] != "not_claimed" else r["check"]["status"]
            status[key] = status.get(key, 0) + 1
    verified = status.get("verified", 0)
    not_verified = status.get("not_verified", 0) + status.get("budget_exhausted", 0)  # the budget is the method's
    done = verified + not_verified
    return {"claimed": sum(1 for r in rows if r.get("claimed_solved")), "status": status, "verified": verified,
            "not_verified": not_verified, "fpr_b": round(not_verified / done, 6) if done else None,
            "llm_calls": sum(c["llm"]["llm_calls"] for c in checks if "llm" in c),
            "cost_usd": round(sum(c["llm"]["cost_usd"] for c in checks if "llm" in c), 6),
            "cost_usd_nominal": round(sum(c["llm"]["cost_usd_nominal"] for c in checks if "llm" in c), 6)}


def run_check(p: dict, res: dict, client, runner, retriever, tracer: Tracer, pipeline_cfg: dict, syntax, b_max: int,
              rag: bool) -> dict:
    """`--check`: formalize C2's claimed proof (nothing to check if it claimed nothing)."""
    if not res.get("claimed_solved"):
        return {"status": "not_claimed", "accepted": False}
    solution = res.get("solution") or ""
    if res.get("final_answer"):
        solution += fragment("final_answer", answer=res["final_answer"])
    budget = Budget(b_max)
    tracer.event("claim_check_start", node="baseline_b_check", b_max=b_max)
    out = check_claim(p, solution, client, runner, retriever, budget, tracer,
                      attempts=int((pipeline_cfg.get("critic") or {}).get("k_max", 3)), syntax=syntax, rag=rag,
                      k_retry=int((pipeline_cfg.get("retrieval") or {}).get("k_retry", 32)))
    return {**out, "llm": budget.summary()}


def summarize(rows: list[dict]) -> dict:
    status: dict[str, int] = {}
    for r in rows:
        status[r["status"]] = status.get(r["status"], 0) + 1
    check = summarize_check(rows)
    return {"problems": len(rows), "status": status, **({"check": check} if check else {}),
            "solved": sum(1 for r in rows if r.get("solved")),
            "claimed_solved": sum(1 for r in rows if r.get("claimed_solved")),
            "llm_calls": sum(r["llm"]["llm_calls"] for r in rows),
            "length_finishes": sum(r["llm"]["length_finishes"] for r in rows),
            "cost_usd": round(sum(r["llm"]["cost_usd"] for r in rows), 6),
            "cost_usd_nominal": round(sum(r["llm"]["cost_usd_nominal"] for r in rows), 6)}


def run_planner_generator(p: dict, client, runner, retriever, budget: Budget, tracer: Tracer,
                          pipeline_cfg: dict, syntax=None) -> dict:
    """Phase 4 preview: plan, then sketch every subgoal (status `budget_stop` if either stopped)."""
    pc, gc = pipeline_cfg.get("planner") or {}, pipeline_cfg.get("generator") or {}
    plan = run_planner(p, client, runner, retriever, budget, tracer, n_max=int(pc.get("n_max", 4)),
                       retrieve=bool(pc.get("retrieve", True)), check_skeleton=bool(pc.get("check_skeleton", True)),
                       refute=pc.get("refute"), closing_auto=pc.get("closing_auto"), syntax=syntax)
    gen = {"status": "not_run", "subgoals": [], "detail": None}
    if plan["plan"] is not None and plan["status"] != "budget_stop":
        gen = run_generator(p, plan["plan"], client, retriever, budget, tracer, samples=int(gc.get("samples", 3)),
                            retrieve=bool(gc.get("retrieve", True)))
    if "budget_stop" in (plan["status"], gen["status"]):
        status = "budget_stop"
    elif plan["plan"] is None:
        status = plan["status"]
    else:
        status = f"plan_{plan['plan']['source']}/sketch_{gen['status']}"
    return {"status": status, "solved": None, "planner": plan, "generator": gen}


def spend_line(info: dict) -> str:
    return (f"spend: project ${info['total_spent_usd']:.4f} of ${info['total_usd_cap']:.2f} "
            f"(${info['total_remaining_usd']:.4f} left); run cap ${info['run_usd_cap']:.2f}")


def main(argv: list[str] | None = None, *, models_cfg: dict | None = None, pipeline_cfg: dict | None = None,
         client_kwargs: dict | None = None, runner_factory=LeanRunner, retriever_factory=Retriever,
         runs_dir: Path | None = None, on_run_dir=None) -> int:
    """Keyword arguments are for offline tests (fake transport, clock and Lean; temporary folders);
    `on_run_dir(run_dir)` is called once the folder exists (the CLI copies the console there)."""
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--config", required=True, choices=sorted(CONFIGS))
    ap.add_argument("--data", type=Path, default=ROOT / "data" / "eval.jsonl")
    ap.add_argument("--split", default="dev", help="dev | test | all")
    ap.add_argument("--limit", type=int, default=None)
    ap.add_argument("--seed", type=int, default=None, help="default: pipeline.yaml → run.seed")
    ap.add_argument("--run-name", default=None, help="suffix for the run folder")
    ap.add_argument("--echo", action="store_true", help="print one line per problem")
    ap.add_argument("--run-usd", type=float, default=None,
                    help="lower the per-run spend cap for this run (never above pipeline.yaml → spend.run_usd)")
    ap.add_argument("--no-prepass", action="store_true",
                    help="C3-C5: turn the automation pre-pass off (prepass.enabled: false, recorded in run_config.json)")
    ap.add_argument("--check", action="store_true",
                    help="C2 only: formalize every claimed proof to measure FPR_B (Eq. 16); see the docstring")
    args = ap.parse_args(argv)
    start = run_start_info()  # git state and time at launch; run_config.json is rewritten at the end
    if args.check and args.config != "C2":
        ap.error("--check is only for C2 (Baseline B)")

    if not args.data.exists():
        ap.error(f"{args.data} does not exist (data/eval.jsonl arrives in Phase 7; try --data data/smoke.jsonl)")
    models_cfg = models_cfg if models_cfg is not None else load_models()
    pipeline_cfg = pipeline_cfg if pipeline_cfg is not None else load_pipeline()
    check_b_max = int((pipeline_cfg.get("budget") or {}).get("b_max", 0))  # C3-C5's B_max, before any ablation
    if args.config in ABLATIONS:
        pipeline_cfg = apply_ablation(pipeline_cfg, args.config)
    if args.no_prepass:
        pipeline_cfg = {**pipeline_cfg, "prepass": {**(pipeline_cfg.get("prepass") or {}), "enabled": False}}
    profile = profile_name(models_cfg)
    split = None if args.split == "all" else args.split
    problems = load_problems(args.data, split, args.limit)
    seed = args.seed if args.seed is not None else int((pipeline_cfg.get("run") or {}).get("seed", 0))
    b_max = int((pipeline_cfg.get("budget") or {})["b_max"])
    flow = pipeline_cfg.get("flow", "full") if args.config in ABLATIONS else "planner_generator"
    attempts = int(((pipeline_cfg.get("baselines") or {}).get(flow) or {}).get("attempts", 1))

    started = utc_now()
    name = args.run_name or ("check" if args.check else None)
    run_id = f"{started:%Y%m%dT%H%M%SZ}_{args.config}_{profile}" + (f"_{name}" if name else "")
    run_dir = (runs_dir or ROOT / "runs") / run_id
    run_dir.mkdir(parents=True, exist_ok=False)
    if on_run_dir is not None:
        on_run_dir(run_dir)
    console = logging.StreamHandler(sys.stderr)  # warnings only; INFO events go to the trace
    console.setLevel(logging.WARNING)
    console.setFormatter(logging.Formatter("%(levelname)s %(name)s: %(message)s"))
    logging.getLogger().addHandler(console)
    tracer = Tracer(run_dir / "trace.jsonl")
    tracer.attach_logging()
    spend = SpendGuard.from_config(pipeline_cfg, run_id=run_id)
    if args.run_usd is not None and spend.run_usd is not None:
        spend.run_usd = min(spend.run_usd, args.run_usd)
    if spend.stopped:
        LOG.warning("spend guard stopped before the run: %s; every paid call will be refused", spend.stopped)
    elif spend.configured:
        print(spend_line(spend.info()), flush=True)
    else:
        LOG.warning("pipeline.yaml has no `spend` section (or no `spend.ledger`): every paid call will be refused")
    client = LLMClient(profile, models_cfg=models_cfg, pipeline_cfg=pipeline_cfg, tracer=tracer, spend=spend,
                       **(client_kwargs or {}))
    peak = in_peak(client.priced_models(), started)
    if peak:
        tracer.event("run_peak_warning", level="warning", models=peak, started=iso(started))
        LOG.warning("starting inside a peak-price window for %s (models.yaml → pricing.peak_windows_utc); "
                    "off-peak is cheaper", peak)
    runner = runner_factory() if flow != "baseline_b" or args.check else None
    use_rag = ((flow in ("full", "planner_generator") or args.check)
               and (pipeline_cfg.get("retrieval") or {}).get("enabled", True))
    retriever = retriever_factory() if use_rag else None
    lean_cfg = pipeline_cfg.get("lean") or {}
    syntax = (SyntaxChecker(runner, float(lean_cfg.get("syntax_timeout_s", 30)))
              if runner is not None and lean_cfg.get("syntax_check", False) else None)

    def config(**extra) -> dict:
        return build_run_config(run_id=run_id, config_name=args.config, seed=seed, dataset=args.data, split=split,
                                llm_info=client.info(), lean_run_info=runner.run_info() if runner else None,
                                pipeline_cfg=pipeline_cfg, argv=sys.argv, start=start,
                                extra={"spend": spend.info(), "flow": flow, **({"retriever": {"mode": retriever.mode,
                                                                                "warnings": retriever.warnings}}
                                                                 if retriever is not None else {}),
                                       **({"check": {"attempts": int((pipeline_cfg.get("critic") or {}).get("k_max", 3)),
                                                     "b_max": check_b_max, "retrieval": use_rag}} if args.check else {}),
                                       **extra})

    write_run_config(run_dir, config(status="running", problems=[p["id"] for p in problems]))
    tracer.event("run_start", run_id=run_id, config=args.config, profile=profile, problems=len(problems),
                 b_max=b_max, attempts=attempts, seed=seed, spend=spend.info())
    rows = []
    lean_down = None  # set when Lean cannot check anything: the run stops (no LLM call is made)
    try:
        # import Mathlib once (~33 s warm, ~22 min cold on the HDD) and check it really loaded
        # (runner probes); if neither the REPL nor the file back-end passes, nothing is run
        if runner and not runner.ensure_available():
            lean_down = "before the first problem"
            tracer.event("run_lean_unavailable", level="error", when=lean_down, warnings=runner.warnings[-5:],
                         not_run=len(problems))
            LOG.error("Lean is not available (neither the REPL nor `lake env lean` loaded %s with the probes); "
                      "no problem was run and no LLM call was made. See lean_runner_warning in %s",
                      "+".join(runner.base_imports), run_dir / "trace.jsonl")
            problems_to_run = []
        else:
            problems_to_run = problems
        with (run_dir / "results.jsonl").open("a", encoding="utf-8") as results:
            for p in problems_to_run:
                with tracer.context(problem_id=p["id"]):
                    budget = Budget(b_max)
                    t0 = time.perf_counter()
                    tracer.event("problem_start", config=args.config)
                    if flow == "planner_generator":
                        res = run_planner_generator(p, client, runner, retriever, budget, tracer, pipeline_cfg,
                                                    syntax=syntax)
                    else:
                        res = run_graph(p, client, runner, retriever, budget, tracer, pipeline_cfg, syntax=syntax)
                    if args.check and res.get("status") != "budget_stop":
                        res["check"] = run_check(p, res, client, runner, retriever, tracer, pipeline_cfg, syntax,
                                                 check_b_max, use_rag)
                    row = {"problem_id": p["id"], "config": args.config, "profile": profile, "split": p.get("split"),
                           "area": p.get("area"), "difficulty": p.get("difficulty"), **res,
                           "llm": budget.summary(), "wall_s": round(time.perf_counter() - t0, 2)}
                    # Lean still usable after this problem? (instant if the REPL is alive; restarts it
                    # with the probes otherwise). If not, this row is marked: its failures may be Lean's.
                    lean_lost = runner is not None and not runner.ensure_available()
                    if lean_lost:
                        row["lean_down"] = True
                    tracer.event("problem_end", status=row["status"], solved=row.get("solved"),
                                 claimed_solved=row.get("claimed_solved"), llm=row["llm"], wall_s=row["wall_s"],
                                 **({"lean_down": True} if lean_lost else {}))
                results.write(json.dumps(row, ensure_ascii=False) + "\n")
                results.flush()
                rows.append(row)
                if args.echo:
                    stage = f" stage={row['stage']}" if row.get("stage") else ""
                    if "check" in row:
                        stage += (f" check={row['check']['status']}"
                                  + (f" check_calls={row['check']['llm']['llm_calls']}" if "llm" in row["check"] else ""))
                    print(f"{p['id']:<28} {row['status']:<17}{stage} calls={row['llm']['llm_calls']:<3} "
                          f"cost=${row['llm']['cost_usd']:.4f} length={row['llm']['length_finishes']} "
                          f"{row['wall_s']:.1f}s", flush=True)
                if row["status"] == "budget_stop" or (row.get("check") or {}).get("status") == "budget_stop":
                    tracer.event("run_budget_stop", level="error", problem_id=p["id"], reason=spend.stopped,
                                 not_run=len(problems) - len(rows))
                    break
                if lean_lost:  # neither back-end loads Mathlib any more: stop before more LLM calls
                    lean_down = f"after {p['id']}"
                    tracer.event("run_lean_unavailable", level="error", when=lean_down,
                                 warnings=runner.warnings[-5:], not_run=len(problems) - len(rows))
                    LOG.error("Lean became unavailable after %s; the remaining %d problems were not run",
                              p["id"], len(problems) - len(rows))
                    break
    finally:
        summary = summarize(rows)
        summary["spend"] = spend.info()
        tracer.event("run_end", **summary)
        if lean_down:
            status = "lean_unavailable"
        elif spend.stopped:
            status = "budget_stop"
        else:
            status = "finished" if len(rows) == len(problems) else "interrupted"
        write_run_config(run_dir, config(status=status, problems=[p["id"] for p in problems],
                                         finished_utc=iso(utc_now()), summary=summary))
        if runner:
            runner.close()
        tracer.close()
    info = summary["spend"]
    print(f"run cost: ${info['run_spent_usd']:.6f} in {info['run_paid_calls']} paid calls"
          + (f"; BUDGET STOP: {spend.stopped}" if spend.stopped else ""), flush=True)
    if spend.configured and info["total_spent_usd"] is not None:
        print(spend_line(info), flush=True)
    shown = run_dir.relative_to(ROOT) if run_dir.is_relative_to(ROOT) else run_dir
    print(json.dumps({"run_dir": shown.as_posix(), "status": status, **summary}, ensure_ascii=False))
    return 2 if lean_down else 0


class _Tee:
    """Writes to the console and to runs/<run>/console.log (long runs are not piped through `tail`)."""

    def __init__(self, stream, log):
        self.stream, self.log = stream, log

    def write(self, text: str) -> int:
        self.log.write(text)
        self.log.flush()
        return self.stream.write(text)

    def flush(self) -> None:
        self.stream.flush()
        self.log.flush()

    def __getattr__(self, name):
        return getattr(self.stream, name)


def copy_console(run_dir: Path) -> None:
    log = (run_dir / "console.log").open("a", encoding="utf-8")
    sys.stdout, sys.stderr = _Tee(sys.stdout, log), _Tee(sys.stderr, log)


if __name__ == "__main__":
    raise SystemExit(main(on_run_dir=copy_console))