scripts/benchmark_models.py

"""Benchmark every registered model by running a fixed one-turn game.

For each model in `wargame.config.MODEL_REGISTRY` (or a subset via
`--models`), this script:

1. Creates a fresh competitor_disruption game pinned to that model
2. Submits a canned player action
3. Advances one turn (deliberations + adjudication + world update + next briefing)
4. Captures per-model metrics: wall time, total LLM calls, validation retries,
   input/output tokens, Turn 1 briefing headline, the first verdict narrative
5. Writes a markdown report with a summary table + per-model detail blocks

Opt-in — makes real API calls. Not wired into pytest. Expected runtime ~2 min
per Anthropic model, similar for the others.

Usage:
    uv run python scripts/benchmark_models.py
    uv run python scripts/benchmark_models.py --out benchmarks/custom.md
    uv run python scripts/benchmark_models.py --models claude-opus-4-7,gpt-5.4
"""

from __future__ import annotations

import argparse
import sys
import time
import traceback
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path

# Add project root to sys.path so `wargame` is importable when run as a script.
_ROOT = Path(__file__).resolve().parent.parent
if str(_ROOT) not in sys.path:
    sys.path.insert(0, str(_ROOT))

from wargame.agents.base import cache_stats, reset_cache_stats  # noqa: E402
from wargame.config import MODEL_REGISTRY  # noqa: E402
from wargame.engine import loop as engine_loop  # noqa: E402
from wargame.engine.state import ActionSubmission, Anchor, AnchorKind  # noqa: E402


CANNED_ACTION = ActionSubmission(
    actor="player",
    intent=(
        "Announce an EU-made heat-pump bundle with a 10-year service guarantee "
        "and match Aergis on installer payouts."
    ),
    justification=(
        "Stabilize the installer channel before Aergis locks in DACH direct "
        "sales, while framing the move around EU-made provenance to pressure "
        "Aergis's Foreign Subsidies exposure."
    ),
    anchors=[
        Anchor(
            kind=AnchorKind.PRIOR_ANALOGUE,
            detail=(
                "Viessmann's 2023 installer-loyalty program stabilized channel "
                "defection within two quarters."
            ),
        ),
        Anchor(
            kind=AnchorKind.STATED_CAPABILITY,
            detail=(
                "Service arm already runs 8-year service contracts; extension "
                "to 10 is a margin decision, not a capability gap."
            ),
        ),
        Anchor(
            kind=AnchorKind.RESOURCE_COMMITTED,
            detail=(
                "EUR 12M service-reserve tranche earmarked for warranty expansion "
                "in the FY24 budget."
            ),
        ),
    ],
)


@dataclass
class ModelResult:
    model_id: str
    display_name: str
    provider: str
    ok: bool = False
    error: str | None = None
    total_seconds: float = 0.0
    new_game_seconds: float = 0.0
    advance_seconds: float = 0.0
    llm_calls: int = 0
    validation_retries: int = 0
    input_tokens: int = 0
    output_tokens: int = 0
    cache_read_tokens: int = 0
    cache_write_tokens: int = 0
    turn1_headline: str = ""
    first_verdict: str = ""
    notes: list[str] = field(default_factory=list)

    @property
    def mean_call_ms(self) -> float:
        return (self.total_seconds / self.llm_calls * 1000.0) if self.llm_calls else 0.0


def benchmark_model(model_id: str) -> ModelResult:
    spec = MODEL_REGISTRY[model_id]
    result = ModelResult(
        model_id=model_id,
        display_name=spec.display_name,
        provider=spec.provider.value,
    )
    reset_cache_stats()
    t_start = time.perf_counter()
    try:
        t0 = time.perf_counter()
        game = engine_loop.new_game(
            scenario_id="competitor_disruption",
            model=model_id,
        )
        result.new_game_seconds = time.perf_counter() - t0

        engine_loop.act(game, CANNED_ACTION)

        t1 = time.perf_counter()
        engine_loop.advance(game, generate_next_briefing=True)
        result.advance_seconds = time.perf_counter() - t1

        briefing = game.turns[0].briefing
        if briefing:
            result.turn1_headline = briefing.headline
        verdicts = game.turns[0].verdicts
        if verdicts:
            v = verdicts[0]
            result.first_verdict = (
                f"{v.action_actor} → {v.outcome.value} "
                f"(plausibility={v.plausibility}): {v.effect_narrative}"
            )
        result.ok = True
    except Exception as exc:  # noqa: BLE001
        result.error = f"{type(exc).__name__}: {exc}"
        result.notes.append(traceback.format_exc(limit=3))
    finally:
        result.total_seconds = time.perf_counter() - t_start
        stats = cache_stats()
        result.llm_calls = stats["calls"]
        result.validation_retries = stats["validation_retries"]
        result.input_tokens = stats["uncached_input_tokens"]
        result.output_tokens = stats["output_tokens"]
        result.cache_read_tokens = stats["cache_read_tokens"]
        result.cache_write_tokens = stats["cache_write_tokens"]
    return result


def write_report(results: list[ModelResult], out_path: Path) -> None:
    out_path.parent.mkdir(parents=True, exist_ok=True)
    lines: list[str] = []
    lines.append("# Model Benchmark — competitor_disruption, one-turn run")
    lines.append("")
    lines.append(
        f"Generated: **{datetime.now(tz=timezone.utc).isoformat(timespec='seconds')}**"
    )
    lines.append("")
    lines.append("Each model ran a fresh game: Turn 1 briefing, canned player action, full")
    lines.append("deliberation + adjudication + world update + Turn 2 briefing. Token counts")
    lines.append("are **uncached** input; cache reads are shown separately.")
    lines.append("")
    lines.append("## Summary")
    lines.append("")
    lines.append(
        "| Model | Provider | Status | Total (s) | New-game (s) | Advance (s) | LLM calls | Retries | Mean call (ms) | In tok | Out tok | Cache read |"
    )
    lines.append(
        "|---|---|---|---:|---:|---:|---:|---:|---:|---:|---:|---:|"
    )
    for r in results:
        status = "ok" if r.ok else f"FAIL ({r.error[:40] if r.error else '?'})"
        lines.append(
            f"| `{r.model_id}` | {r.provider} | {status} | {r.total_seconds:.1f} | "
            f"{r.new_game_seconds:.1f} | {r.advance_seconds:.1f} | {r.llm_calls} | "
            f"{r.validation_retries} | {r.mean_call_ms:.0f} | {r.input_tokens} | "
            f"{r.output_tokens} | {r.cache_read_tokens} |"
        )
    lines.append("")
    lines.append("## Per-model detail")
    lines.append("")
    for r in results:
        lines.append(f"### `{r.model_id}` — {r.display_name}")
        lines.append("")
        if not r.ok:
            lines.append(f"**Failed:** `{r.error}`")
            if r.notes:
                lines.append("")
                lines.append("```")
                lines.append(r.notes[0].strip())
                lines.append("```")
            lines.append("")
            continue
        lines.append(f"- Provider: `{r.provider}`")
        lines.append(f"- Total wall time: **{r.total_seconds:.1f}s** "
                     f"(new_game {r.new_game_seconds:.1f}s + advance {r.advance_seconds:.1f}s)")
        lines.append(f"- LLM calls: **{r.llm_calls}**, validation retries: **{r.validation_retries}**")
        lines.append(f"- Mean call latency: **{r.mean_call_ms:.0f} ms**")
        lines.append(f"- Tokens: {r.input_tokens} in / {r.output_tokens} out "
                     f"(cache read {r.cache_read_tokens}, write {r.cache_write_tokens})")
        lines.append("")
        if r.turn1_headline:
            lines.append(f"**Turn 1 headline:** {r.turn1_headline}")
            lines.append("")
        if r.first_verdict:
            lines.append(f"**First verdict:** {r.first_verdict}")
            lines.append("")
    out_path.write_text("\n".join(lines), encoding="utf-8")


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument(
        "--out",
        type=Path,
        default=None,
        help="Output markdown path. Defaults to benchmarks/<timestamp>.md.",
    )
    ap.add_argument(
        "--models",
        type=str,
        default=None,
        help="Comma-separated model ids. Defaults to every registered model.",
    )
    args = ap.parse_args()

    if args.models:
        requested = [m.strip() for m in args.models.split(",") if m.strip()]
    else:
        requested = list(MODEL_REGISTRY.keys())

    unknown = [m for m in requested if m not in MODEL_REGISTRY]
    if unknown:
        print(f"Unknown models: {', '.join(unknown)}", file=sys.stderr)
        return 2

    if args.out is None:
        ts = datetime.now(tz=timezone.utc).strftime("%Y%m%dT%H%M%SZ")
        args.out = Path("benchmarks") / f"{ts}.md"

    results: list[ModelResult] = []
    for model_id in requested:
        print(f"[benchmark] {model_id}...", flush=True)
        r = benchmark_model(model_id)
        status = "ok" if r.ok else "FAIL"
        print(
            f"  -> {status} in {r.total_seconds:.1f}s, "
            f"{r.llm_calls} calls, {r.validation_retries} retries",
            flush=True,
        )
        if not r.ok:
            print(f"  error: {r.error}", flush=True)
        results.append(r)

    write_report(results, args.out)
    print(f"[benchmark] wrote {args.out}")
    return 0


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