baseline_eval.py

"""
Baseline evaluation harness — run a frontier LLM against the full benchmark
(every task in CentificAIResearch/BA-Agent-Bench) and emit a leaderboard row.

Usage:
    OPENROUTER_API_KEY=... python baseline_eval.py --model openai/gpt-4o
"""
from __future__ import annotations

import argparse
import json
import os
import statistics
import time
from typing import Any, Dict, List

from client import BAAgentClient
from train import _openrouter_policy, run_episode, _stub_policy


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--server", default="http://localhost:8000")
    parser.add_argument("--model", default="openai/gpt-4o")
    parser.add_argument("--output", default="baseline_results.json")
    args = parser.parse_args()

    env = BAAgentClient(args.server)
    tasks = env.tasks()
    print(f"Evaluating {args.model} on {len(tasks)} tasks ...")

    rows: List[Dict[str, Any]] = []
    for t in tasks:
        try:
            t0 = time.time()
            out = run_episode(env, _stub_policy, args.model if os.environ.get("OPENROUTER_API_KEY") else None)
            out["dt_s"] = round(time.time() - t0, 1)
        except Exception as exc:
            out = {"task_id": t["task_id"], "error": str(exc)}
        rows.append(out)
        c = out.get("composite_0_to_100")
        print(f"  {out.get('task_id','?'):<10} composite={c if c is not None else '-':>6}/100  reward={out.get('episode_reward','-'):+}")

    composites = [r["composite_0_to_100"] for r in rows if r.get("composite_0_to_100") is not None]
    summary = {
        "model": args.model,
        "n_tasks": len(rows),
        "mean_composite_0_to_100": round(statistics.mean(composites), 2) if composites else None,
        "median_composite_0_to_100": round(statistics.median(composites), 2) if composites else None,
        "stdev_composite_0_to_100": round(statistics.pstdev(composites), 2) if len(composites) > 1 else 0.0,
    }
    print(json.dumps(summary, indent=2))

    with open(args.output, "w") as f:
        json.dump({"summary": summary, "rows": rows}, f, indent=2)
    print(f"\nWrote {args.output}")


if __name__ == "__main__":
    main()