Blame for run_evaluation.py

"""
run_evaluation.py
=================
Runs inference.py once per scenario across all difficulties and aggregates results.

Usage:
    python run_evaluation.py

Environment variables forwarded to each inference.py subprocess:
    API_BASE_URL, MODEL_NAME, HF_TOKEN, OPENAI_API_KEY, MAX_STEPS

Each subprocess is run as:
    DIFFICULTY=<d> SCENARIO_ID=<id> python inference.py
"""

from __future__ import annotations

import json
import os
import re
import subprocess
import sys
from pathlib import Path
from typing import Optional

try:
    from dotenv import load_dotenv
    load_dotenv()
except ImportError:
    pass

# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------

DIFFICULTIES   = ["easy", "medium", "hard"]
TASKS_DIR      = Path(__file__).parent / "tasks"
PYTHON         = sys.executable          # same interpreter that runs this script
INFERENCE_SCRIPT = Path(__file__).parent / "inference.py"
TIMEOUT        = 180                     # seconds per episode (3 min ceiling)

# Env vars to forward to subprocesses (all LLM + infra config)
FORWARD_VARS = [
    "API_BASE_URL", "MODEL_NAME",
    "HF_TOKEN", "OPENAI_API_KEY",
    "MAX_STEPS", "REDIS_URL",
]


# ---------------------------------------------------------------------------
# Load scenarios
# ---------------------------------------------------------------------------

def load_scenario_ids(difficulty: str) -> list[str]:
    path = TASKS_DIR / difficulty / "scenarios.json"
    with open(path, "r", encoding="utf-8") as f:
        scenarios = json.load(f)
    return [s["id"] for s in scenarios]


# ---------------------------------------------------------------------------
# Parse [END] line
# ---------------------------------------------------------------------------

_END_RE = re.compile(
    r"\[END\]\s+"
    r"success=(?P<success>\w+)\s+"
    r"steps=(?P<steps>\d+)\s+"
    r"score=(?P<score>[\d.]+)"
)


def parse_end_line(output: str) -> Optional[dict]:
    """Extract score/steps/success from [END] line in subprocess stdout."""
    for line in output.splitlines():
        m = _END_RE.search(line)
        if m:
            return {
                "score":   float(m.group("score")),
                "steps":   int(m.group("steps")),
                "success": m.group("success").lower() == "true",
            }
    return None


# ---------------------------------------------------------------------------
# Run one episode
# ---------------------------------------------------------------------------

def run_episode(difficulty: str, scenario_id: str) -> Optional[dict]:
    """
    Invoke inference.py as a subprocess with DIFFICULTY and SCENARIO_ID set.
    Returns parsed result dict or None on failure.
    """
    print(f"  Running {scenario_id}...", flush=True)

    env = {**os.environ}
    env["DIFFICULTY"]  = difficulty
    env["SCENARIO_ID"] = scenario_id

    try:
        proc = subprocess.run(
            [PYTHON, str(INFERENCE_SCRIPT)],
            env=env,
            capture_output=True,
            text=True,
            timeout=TIMEOUT,
        )
    except subprocess.TimeoutExpired:
        print(f"    TIMEOUT after {TIMEOUT}s — skipping", flush=True)
        return None
    except Exception as exc:
        print(f"    ERROR launching subprocess: {exc}", flush=True)
        return None

    # Print [START]/[STEP]/[END] lines live (already captured, echo them)
    for line in proc.stdout.splitlines():
        if line.startswith(("[START]", "[STEP]", "[END]")):
            print(f"    {line}", flush=True)

    if proc.returncode != 0:
        print(f"    FAILED (exit {proc.returncode})", flush=True)
        if proc.stderr:
            # Show last error line only to keep output clean
            last_err = [l for l in proc.stderr.splitlines() if l.strip()]
            if last_err:
                print(f"    stderr: {last_err[-1]}", flush=True)
        return None

    result = parse_end_line(proc.stdout)
    if result is None:
        print(f"    WARNING: could not parse [END] line", flush=True)
    return result


# ---------------------------------------------------------------------------
# Main evaluation loop
# ---------------------------------------------------------------------------

def run_evaluation() -> dict:
    results: dict[str, list[dict]] = {d: [] for d in DIFFICULTIES}

    for difficulty in DIFFICULTIES:
        print(f"\n{'='*40}", flush=True)
        print(f"Difficulty: {difficulty.upper()}", flush=True)
        print(f"{'='*40}", flush=True)

        ids = load_scenario_ids(difficulty)
        for scenario_id in ids:
            ep = run_episode(difficulty, scenario_id)
            if ep is not None:
                ep["scenario_id"] = scenario_id
                results[difficulty].append(ep)
            else:
                results[difficulty].append({
                    "scenario_id": scenario_id,
                    "score":       0.0,
                    "steps":       0,
                    "success":     False,
                    "error":       True,
                })

    return results


# ---------------------------------------------------------------------------
# Metrics + summary
# ---------------------------------------------------------------------------

def compute_metrics(episodes: list[dict]) -> dict:
    if not episodes:
        return {"avg_score": 0.0, "avg_steps": 0.0, "success_rate": 0.0, "n": 0}
    n            = len(episodes)
    avg_score    = sum(e["score"] for e in episodes) / n
    avg_steps    = sum(e["steps"] for e in episodes) / n
    success_rate = sum(1 for e in episodes if e["success"]) / n * 100
    return {"avg_score": avg_score, "avg_steps": avg_steps, "success_rate": success_rate, "n": n}


def print_summary(results: dict[str, list[dict]]) -> None:
    metrics = {d: compute_metrics(results[d]) for d in DIFFICULTIES}

    print("\n" + "=" * 34)
    print("      FINAL RESULTS")
    print("=" * 34)

    for difficulty in DIFFICULTIES:
        m = metrics[difficulty]
        print(f"\n{difficulty}:")
        print(f"  avg_score:    {m['avg_score']:.2f}")
        print(f"  avg_steps:    {m['avg_steps']:.1f}")
        print(f"  success_rate: {m['success_rate']:.0f}%  ({sum(e['success'] for e in results[difficulty])}/{m['n']})")

    # Difficulty tier validation
    easy_s   = metrics["easy"]["avg_score"]
    medium_s = metrics["medium"]["avg_score"]
    hard_s   = metrics["hard"]["avg_score"]

    print()
    if easy_s > medium_s > hard_s:
        print("Difficulty progression validated ✓  (easy > medium > hard)")
    else:
        print(
            "WARNING: Difficulty tiers not well separated "
            f"(easy={easy_s:.3f}, medium={medium_s:.3f}, hard={hard_s:.3f})"
        )
    print()


# ---------------------------------------------------------------------------
# Entrypoint
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    print("Starting evaluation — 15 episodes across easy / medium / hard")
    print(f"Model : {os.getenv('MODEL_NAME', 'gpt-4o')}")
    print(f"API   : {os.getenv('API_BASE_URL', 'https://api.openai.com/v1')}")
    print(f"Timeout per episode: {TIMEOUT}s")

    results = run_evaluation()
    print_summary(results)