eval.py

"""Evaluation harness for agents on benchmark task datasets.

Loops over every sample in the dataset, runs an agent function on each task,
then executes the per-task validation function to score the result.

Usage:
    python eval.py --agent bash                                   # bash terminal agent
    python eval.py --agent mcp                                    # MCP agent
    python eval.py --agent playwright                             # playwright browser agent
    python eval.py --agent api_call                               # single api_call tool agent
    python eval.py --agent bash --domain gitlab --no-docs --no-skills
    python eval.py --agent bash --task-name "*.create-*" --max-samples 5
"""

import argparse
import asyncio
import fnmatch
import json
import logging
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List
from datasets import load_dataset

from starshell.config import PROJECT_ROOT
from starshell.eval import sanitize_name, create_agent_factory, run_single_task
from starshell.model import DEFAULT_MODEL_ID
from starshell.data import DOMAIN_DATASETS

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Helpers for structured output directories
# ---------------------------------------------------------------------------
def make_run_dir(
    agent_type: str,
    model_id: str,
    domain: str,
    task_filter: str | None,
) -> Path:
    """Build a timestamped run directory under results/."""
    ts = datetime.now().strftime("%Y%m%d_%H%M%S")
    model_slug = sanitize_name(model_id)
    subset = sanitize_name(task_filter) if task_filter else "all"
    name = f"{ts}_{agent_type}_{model_slug}_{domain}_{subset}"
    run_dir = PROJECT_ROOT / "results" / name
    run_dir.mkdir(parents=True, exist_ok=True)
    return run_dir


# ---------------------------------------------------------------------------
# Main eval loop
# ---------------------------------------------------------------------------
async def evaluate(
    agent_type: str,
    domain: str = "servicenow",
    task_name_filter: str | None = None,
    max_samples: int | None = None,
    include_docs: bool = True,
    include_skills: bool = True,
    model_id: str | None = None,
) -> Dict[str, Any]:
    """Run the full evaluation and return aggregate results."""

    dataset_name = DOMAIN_DATASETS[domain]
    ds = load_dataset(dataset_name, split="train")
    log.info("Loaded %d samples from %s", len(ds), dataset_name)

    # Filter
    samples = list(ds)
    if task_name_filter:
        samples = [s for s in samples if fnmatch.fnmatch(s["task_name"], task_name_filter)]
        log.info("Filtered to %d samples matching %r", len(samples), task_name_filter)
    if max_samples:
        samples = samples[:max_samples]
        log.info("Capped to %d samples", len(samples))

    # Build factory with domain args
    factory = create_agent_factory(
        agent_type, domain,
        model_id=model_id,
        include_docs=include_docs,
        include_skills=include_skills,
    )

    # Resolve effective model for directory naming
    effective_model_id = model_id or DEFAULT_MODEL_ID
    run_dir = make_run_dir(agent_type, effective_model_id, domain, task_name_filter)
    log.info("Run directory: %s", run_dir)

    results = []

    async with factory() as run_agent:
        for i, sample in enumerate(samples):
            task_id = f"{sample['task_name']}.{sample['task_seed']}"
            task_dir = run_dir / task_id
            log.info("[%d/%d] Running: %s", i + 1, len(samples), task_id)

            outcome = await run_single_task(run_agent, sample, task_dir)
            results.append(outcome)

    # --- Aggregate (exclude skipped tasks) ---
    evaluated = [r for r in results if not r.get("skipped")]
    skipped = [r for r in results if r.get("skipped")]
    total = len(evaluated)
    passed = sum(r["score"] for r in evaluated)
    by_task: Dict[str, List[int]] = {}
    for r in evaluated:
        by_task.setdefault(r["task_name"], []).append(r["score"])

    summary = {
        "agent_type": agent_type,
        "model_id": effective_model_id,
        "domain": domain,
        "total": total,
        "passed": passed,
        "skipped": len(skipped),
        "accuracy": round(passed / total, 4) if total else 0,
        "per_task_accuracy": {
            name: round(sum(scores) / len(scores), 4)
            for name, scores in sorted(by_task.items())
        },
        "results": results,
    }

    # Write summary inside the run directory
    with open(run_dir / "summary.json", "w") as f:
        json.dump(summary, f, indent=2)
    log.info("Results written to %s", run_dir / "summary.json")

    log.info(
        "=== DONE [%s]: %d/%d passed (%.1f%%), %d skipped (pre-solved) ===",
        agent_type,
        passed,
        total,
        summary["accuracy"] * 100,
        len(skipped),
    )
    return summary


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------
def main():
    parser = argparse.ArgumentParser(description="Evaluate an agent on benchmark tasks")
    parser.add_argument(
        "--agent",
        choices=["bash", "mcp", "playwright", "mas", "hybrid", "api_call"],
        required=True,
        help="Agent type to evaluate",
    )
    parser.add_argument(
        "--domain",
        default="servicenow",
        choices=["servicenow", "gitlab", "erpnext"],
        help="Domain to target (default: servicenow)",
    )
    parser.add_argument(
        "--task-name",
        default=None,
        help="Glob filter on task_name (e.g. '*.create-*')",
    )
    parser.add_argument(
        "--max-samples",
        type=int,
        default=None,
        help="Max number of samples to evaluate",
    )
    parser.add_argument(
        "--model",
        default=None,
        help=f"LiteLLM model ID (default: {DEFAULT_MODEL_ID})",
    )
    parser.add_argument(
        "--no-docs",
        action="store_true",
        help="Omit docs section from bash agent prompt",
    )
    parser.add_argument(
        "--no-skills",
        action="store_true",
        help="Omit skills section from bash agent prompt",
    )
    args = parser.parse_args()
    asyncio.run(
        evaluate(
            agent_type=args.agent,
            domain=args.domain,
            task_name_filter=args.task_name,
            max_samples=args.max_samples,
            include_docs=not args.no_docs,
            include_skills=not args.no_skills,
            model_id=args.model,
        )
    )


if __name__ == "__main__":
    main()