run_eval.py

"""Run BU Bench V2 with the weighted findings judge by default.

uv run python run_eval.py --tasks 5
uv run python run_eval.py --browser local_headless
uv run python run_eval.py --benchmark BU_Bench_V1  # legacy binary judge
"""

import argparse
import asyncio
import base64
import hashlib
import json
import logging
import os
import traceback
from datetime import UTC, datetime
from importlib.metadata import version
from pathlib import Path

import certifi

os.environ.setdefault("SSL_CERT_FILE", certifi.where())
os.environ["BROWSER_USE_SETUP_LOGGING"] = "false"  # Before importing browser_use
logging.basicConfig(level=logging.CRITICAL)

from browser_use import Agent, Browser
from browser_use.llm import ChatBrowserUse, ChatOpenAI
from dotenv import load_dotenv

from browsers import PROVIDERS, get_provider
from evaluation import (
    BENCHMARKS,
    DEFAULT_BENCHMARK,
    create_judge,
    judge_config,
    judge_trace,
    load_tasks,
    validate_findings_task,
)

load_dotenv()

MAX_CONCURRENT = 3
TASK_TIMEOUT = 3600
AGENT_FRAMEWORK_NAME = "BrowserUse"
AGENT_FRAMEWORK_VERSION = version("browser-use")
MODEL_NAME = "bu-2-0"


def create_agent_model(model: str = MODEL_NAME, reasoning: str = "xhigh"):
    if model == "bu-2-0":
        return ChatBrowserUse(model=model)
    if not os.getenv("OPENAI_API_KEY"):
        raise ValueError("OPENAI_API_KEY is required for the selected agent model")
    return ChatOpenAI(
        model=model,
        reasoning_effort=reasoning,
        max_completion_tokens=32768,
        temperature=None,
        frequency_penalty=None,
        timeout=300,
        max_retries=2,
    )


def select_shard(
    tasks: list[dict], shard_index: int | None, shard_count: int | None
) -> list[dict]:
    """Select one contiguous shard while preserving the dataset's task order."""
    if shard_index is None or shard_count is None:
        return tasks
    start = len(tasks) * shard_index // shard_count
    end = len(tasks) * (shard_index + 1) // shard_count
    return tasks[start:end]


def select_tasks(
    tasks: list[dict], task_ids: list[str] | None = None, limit: int | None = None
):
    if task_ids:
        by_id = {task["task_id"]: task for task in tasks}
        unknown = set(task_ids) - by_id.keys()
        if unknown or len(task_ids) != len(set(task_ids)):
            raise ValueError("Task IDs must be unique and present in the dataset")
        tasks = [by_id[task_id] for task_id in task_ids]
    return tasks[:limit] if limit is not None else tasks


def encode_screenshots(paths: list[str], evidence_notes: list[str] | None = None) -> list[str]:
    images = []
    for index, path in enumerate(paths, 1):
        try:
            images.append(base64.b64encode(Path(path).read_bytes()).decode())
        except OSError:
            if evidence_notes is not None:
                evidence_notes.append(f"Capture {index} could not be read from the saved screenshots.")
    return images


def collect_output_files(agent, evidence_notes: list[str] | None = None) -> str:
    """Collect source text of files managed by this agent, including PDF/DOCX."""
    fs = agent.file_system
    contents = []
    for name in fs.list_files():
        try:
            contents.append(f"--- {name} ---\n{fs.get_file(name).read()}")
        except Exception:
            if evidence_notes is not None:
                evidence_notes.append(f"Managed output file {name} could not be read by the collector.")
    return "\n\n".join(contents)


def save_task(run_data_dir: Path | None, task_id: str, payload: dict) -> None:
    if run_data_dir is not None:
        run_data_dir.mkdir(parents=True, exist_ok=True)
        (run_data_dir / f"{task_id}.json").write_text(json.dumps(payload, indent=2))


async def create_browser(
    browser_provider, timeout_seconds: int = TASK_TIMEOUT
) -> Browser:
    if browser_provider is None:
        return Browser(
            use_cloud=True, cloud_timeout=max(1, (timeout_seconds + 59) // 60)
        )
    cdp_url = await browser_provider.connect()
    if cdp_url is None:
        return Browser(headless=getattr(browser_provider, "HEADLESS", True))
    return Browser(cdp_url=cdp_url)


async def run_task(
    task: dict,
    semaphore: asyncio.Semaphore,
    browser_provider=None,
    llm=None,
    run_data_dir: Path | None = None,
    *,
    benchmark: str = DEFAULT_BENCHMARK,
    judge_llm=None,
    task_timeout: int = TASK_TIMEOUT,
    max_steps: int = 100,
) -> dict:
    """Execute and judge one task, retaining evidence if judging fails."""
    async with semaphore:
        task_id = task.get("task_id", "unknown")
        browser = None
        phase = "setup"
        result = {
            "task_id": task_id,
            "score": None,
            "raw_score": None,
            "steps": 0,
            "duration": 0,
            "cost": 0,
        }
        artifact = {"benchmark": benchmark}
        try:
            if benchmark == DEFAULT_BENCHMARK:
                validate_findings_task(task)
            judge_llm = judge_llm or create_judge(benchmark)
            artifact["judge_config"] = judge_config(benchmark, judge_llm)
            print(f"Running task: {task_id}")
            browser = await create_browser(
                browser_provider, timeout_seconds=task_timeout
            )
            agent = Agent(
                task=task["confirmed_task"],
                llm=llm or ChatBrowserUse(model=MODEL_NAME),
                browser=browser,
                use_judge=False,  # Only the benchmark judge evaluates the run.
            )
            phase = "execution"
            try:
                history = await asyncio.wait_for(
                    agent.run(max_steps=max_steps), timeout=task_timeout
                )
            except asyncio.TimeoutError:
                # Preserve and judge partial work instead of discarding its evidence.
                history = agent.history
                result["execution_error"] = f"Task timed out after {task_timeout}s"
            except Exception as exc:
                history = agent.history
                result["execution_error"] = f"{type(exc).__name__}: {exc}"
            phase = "evidence"
            result.update(
                steps=history.number_of_steps(),
                duration=history.total_duration_seconds(),
                cost=history.usage.total_cost if history.usage else 0,
            )
            evidence_notes = []
            trace = {
                "agent_task": task["confirmed_task"],
                "final_result": history.final_result()
                or "Agent did not return a result",
                "agent_steps": history.agent_steps(),
                "ground_truth": task.get("answer"),
                "screenshots_b64": encode_screenshots(
                    [p for p in history.screenshot_paths() if p is not None], evidence_notes
                ),
                "output_files_text": collect_output_files(agent, evidence_notes),
                "evidence_notes": evidence_notes,
            }
            artifact.update(
                agent_trace=trace,
                metrics={key: result[key] for key in ("steps", "duration", "cost")},
                # Retained locally to allow re-judging the exact task version.
                task=task,
            )
            result["status"] = "awaiting_judge"
            save_task(run_data_dir, task_id, {**artifact, **result})
            phase = "judge"
            result.update(await judge_trace(task, trace, judge_llm, benchmark))
            if result.get("evidence_clipped_sections"):
                sections = ", ".join(result["evidence_clipped_sections"])
                print(f"Task {task_id}: warning: clipped {sections}")
            if result.get("evidence_missing_items"):
                print(f"Task {task_id}: warning: missing evidence for {result['evidence_missing_items']}; item credit follows the rubric")
            image_info = result.get("screenshot_evidence", {})
            if image_info.get("reencoded") or image_info.get("omitted"):
                print(f"Task {task_id}: screenshot preparation: {len(image_info['reencoded'])} reencoded, {len(image_info['resized'])} resized, {len(image_info['omitted'])} omitted; see evidence notes")
            result["status"] = "judged"
            print("Task", task_id, "raw=", result["raw_score"], "final=", result["score"])
        except Exception as exc:
            # A broken judge/configuration is not an agent score of zero.
            result.update(
                status=f"{phase}_error",
                error=f"{type(exc).__name__}: {exc}",
                traceback=traceback.format_exc(),
            )
            print(f"Task {task_id}: {result['status']}: {result['error']}")
        finally:
            cleanup_errors = []
            if browser is not None:
                try:
                    await browser.stop()
                except Exception as exc:
                    cleanup_errors.append(f"browser: {exc}")
            if browser_provider is not None:
                try:
                    await browser_provider.disconnect()
                except Exception as exc:
                    cleanup_errors.append(f"provider: {exc}")
            if cleanup_errors:
                result["cleanup_errors"] = cleanup_errors
        save_task(run_data_dir, task_id, {**artifact, **result})
        return result


def summarize_results(results: list[dict]) -> dict:
    scored = [r for r in results if r["score"] is not None]
    complete = len(scored) == len(results)
    mean = sum(r["score"] for r in scored) / len(scored) if scored else None
    raw_mean = sum(r["raw_score"] for r in scored) / len(scored) if scored else None
    return {
        "tasks_completed": len(results),
        "tasks_scored": len(scored),
        "tasks_unscored": len(results) - len(scored),
        "tasks_successful": sum(r["score"] == 1 for r in scored),
        "judge_errors": sum(r["status"] == "judge_error" for r in results),
        "evidence_incomplete": sum(
            r["status"] == "evidence_incomplete" for r in results
        ),
        "tasks_with_clipped_evidence": sum(
            bool(r.get("evidence_clipped_sections")) for r in results
        ),
        "tasks_with_missing_evidence": sum(
            bool(r.get("evidence_missing_items")) for r in results
        ),
        "execution_errors": sum(
            r["status"] == "execution_error" or "execution_error" in r for r in results
        ),
        "mean_score": mean if complete else None,
        "mean_raw_score": raw_mean if complete else None,
        "mean_score_scored_tasks": mean,
        "mean_raw_score_scored_tasks": raw_mean,
        "rh_zeroed": sum(bool(r.get("rh_zeroed")) for r in results),
        "total_steps": sum(r["steps"] for r in results),
        "total_duration": sum(r["duration"] for r in results),
        "total_cost": sum(r["cost"] for r in results),
    }


def parse_args(argv=None):
    parser = argparse.ArgumentParser(
        description="Run all 200 V2 tasks with BrowserCode and the findings judge (default)"
    )
    parser.add_argument(
        "--executor",
        choices=["bcode", "browser-use"],
        default=None,
        help="Default: bcode for V2; browser-use for V1/Stealth",
    )
    parser.add_argument("--bcode-bin", default=str(Path.home() / ".bcode/bin/bcode"))
    parser.add_argument("--bcode-version", default="0.1.20")
    parser.add_argument("--chrome-bin", help="Local Chrome/Chromium executable")
    parser.add_argument("--fetch-use", action=argparse.BooleanOptionalAction, default=None,
                        help="Browser Use fetch service (default: on with Cloud, off with local Chrome)")
    parser.add_argument(
        "--output-dir", type=Path, default=Path(__file__).parent / "run_data"
    )
    parser.add_argument(
        "--check",
        action="store_true",
        help="Validate BrowserCode configuration without running tasks",
    )
    parser.add_argument(
        "--browser",
        default="browser-use-cloud",
        choices=["browser-use-cloud"] + PROVIDERS,
    )
    parser.add_argument("--benchmark", default=DEFAULT_BENCHMARK, choices=BENCHMARKS)
    parser.add_argument(
        "--tasks", type=int, default=None, help="Number of tasks (default: all)"
    )
    parser.add_argument(
        "--judge-model",
        help="Override judge model (default: gpt-5.6-luna for V2; gemini-2.5-flash for V1)",
    )
    parser.add_argument(
        "--judge-reasoning",
        default="xhigh",
        choices=["none", "minimal", "low", "medium", "high", "xhigh"],
        help="Reasoning effort for the V2 OpenAI judge",
    )
    parser.add_argument(
        "--model",
        default=None,
        help="Executor model; separate from --judge-model",
    )
    parser.add_argument(
        "--agent-reasoning",
        default=None,
        choices=["none", "minimal", "low", "medium", "high", "xhigh", "max"],
    )
    parser.add_argument(
        "--task-ids", nargs="+", help="Exact task IDs, in execution order"
    )
    parser.add_argument(
        "--shard-index", type=int, help="Zero-based contiguous shard index"
    )
    parser.add_argument("--shard-count", type=int, help="Total contiguous shards")
    parser.add_argument("--task-timeout", type=int, default=None)
    parser.add_argument("--max-steps", type=int, default=None)
    parser.add_argument("--concurrency", "--parallel", type=int, default=None)
    args = parser.parse_args(argv)
    args.executor = args.executor or (
        "bcode" if args.benchmark == DEFAULT_BENCHMARK else "browser-use"
    )
    if args.concurrency is None:
        args.concurrency = 100 if args.executor == "bcode" else MAX_CONCURRENT
    if (args.shard_index is None) != (args.shard_count is None):
        parser.error("--shard-index and --shard-count must be provided together")
    if args.shard_count is not None and (
        args.shard_count < 1 or not 0 <= args.shard_index < args.shard_count
    ):
        parser.error("Shard index must be between zero and shard-count - 1")
    if args.task_ids and args.shard_count is not None:
        parser.error("Use either --task-ids or --shard-index/--shard-count")
    if args.executor == "bcode":
        if args.benchmark != DEFAULT_BENCHMARK or args.browser not in ("browser-use-cloud", "local_headless", "local_headful"):
            parser.error(
                "BrowserCode supports V2 with Cloud or local Chrome; select --executor browser-use for other benchmarks/browsers"
            )
        if args.max_steps is not None:
            parser.error(
                "--max-steps applies to --executor browser-use; use --task-timeout for BrowserCode"
            )
        args.fetch_use = args.fetch_use if args.fetch_use is not None else args.browser == "browser-use-cloud"
        from bcode_runner import DEFAULT_MODEL, resolve_model

        try:
            args.model, args.agent_reasoning = resolve_model(
                args.model or DEFAULT_MODEL, args.agent_reasoning
            )
        except ValueError as exc:
            parser.error(str(exc))
        args.task_timeout = args.task_timeout if args.task_timeout is not None else 3600
        if args.browser == "browser-use-cloud" and args.task_timeout > 14400:
            parser.error("BrowserCode Cloud task timeout cannot exceed 14400 seconds")
    else:
        if args.shard_count is not None:
            parser.error(
                "Sharding is currently supported only by the BrowserCode executor"
            )
        if args.check:
            parser.error("--check applies to the BrowserCode executor")
        args.model = args.model or MODEL_NAME
        args.agent_reasoning = args.agent_reasoning or "xhigh"
        args.task_timeout = (
            args.task_timeout if args.task_timeout is not None else TASK_TIMEOUT
        )
        args.max_steps = args.max_steps if args.max_steps is not None else 100
    args.max_steps = args.max_steps if args.max_steps is not None else 100
    if min(args.task_timeout, args.max_steps, args.concurrency) < 1:
        parser.error("Timeout, max steps, and concurrency must be positive")
    if args.tasks is not None and args.tasks < 1:
        parser.error("--tasks must be positive")
    return args


async def main(argv=None):
    args = parse_args(argv)
    if args.executor == "bcode":
        from bcode_eval import main as run_bcode

        return await run_bcode(args)
    tasks = select_tasks(load_tasks(args.benchmark), args.task_ids, args.tasks)
    # Fail on missing judge credentials before launching any browser sessions.
    judge_llm = create_judge(args.benchmark, args.judge_model, args.judge_reasoning)
    config = judge_config(args.benchmark, judge_llm)
    agent_llm = create_agent_model(args.model, args.agent_reasoning)
    browser_provider = (
        None if args.browser == "browser-use-cloud" else get_provider(args.browser)
    )
    run_start = datetime.now(UTC).strftime("%Y%m%d_%H%M%S_%f")
    run_key = f"{args.benchmark}_{AGENT_FRAMEWORK_NAME}_{AGENT_FRAMEWORK_VERSION}_browser_{args.browser}_model_{args.model}"
    root = Path(__file__).parent
    run_data_dir = root / "run_data" / f"{run_key}_start_at_{run_start}"
    results_file = root / "results" / f"{run_key}.json"
    run_config = {
        "benchmark": args.benchmark,
        "runner_source_sha256": hashlib.sha256(Path(__file__).read_bytes()).hexdigest(),
        "dataset_sha256": hashlib.sha256(
            (root / f"{args.benchmark}.enc").read_bytes()
        ).hexdigest(),
        "judge": config,
        "agent_model": args.model,
        "agent_reasoning": getattr(agent_llm, "reasoning_effort", None),
        "agent_max_completion_tokens": getattr(
            agent_llm, "max_completion_tokens", None
        ),
        "task_ids": [t["task_id"] for t in tasks],
        "agent_framework_version": AGENT_FRAMEWORK_VERSION,
        "task_timeout_seconds": args.task_timeout,
        "max_steps": args.max_steps,
        "max_concurrent": args.concurrency,
    }
    run_data_dir.mkdir(parents=True)
    (run_data_dir / "config.json").write_text(json.dumps(run_config, indent=2))
    print(
        f"{args.benchmark}: {len(tasks)} tasks, judge={config['type']} / {config['model']} / {config['reasoning_effort']}"
    )
    sem = asyncio.Semaphore(args.concurrency)
    results = await asyncio.gather(
        *[
            run_task(
                t,
                sem,
                browser_provider=browser_provider,
                run_data_dir=run_data_dir,
                benchmark=args.benchmark,
                judge_llm=judge_llm,
                llm=agent_llm,
                task_timeout=args.task_timeout,
                max_steps=args.max_steps,
            )
            for t in tasks
        ]
    )
    summary = summarize_results(results)
    results_file.parent.mkdir(parents=True, exist_ok=True)
    runs = json.loads(results_file.read_text()) if results_file.exists() else []
    runs.append(
        {
            "run_start": run_start,
            "config": run_config,
            **summary,
            "task_results": results,
        }
    )
    results_file.write_text(json.dumps(runs, indent=2))
    if summary["mean_score"] is not None:
        print(
            f"Mean weighted score: {summary['mean_score']:.2%} (raw: {summary['mean_raw_score']:.2%})"
        )
    else:
        print(
            f"Incomplete evaluation: {summary['tasks_unscored']} unscored tasks; see {run_data_dir}"
        )
    if summary["tasks_with_clipped_evidence"]:
        print(
            f"Evidence warning: {summary['tasks_with_clipped_evidence']} tasks "
            "had clipped evidence; see their findings and clipped sections."
        )
    print(f"Results: {results_file}")
    if summary["tasks_unscored"]:
        raise SystemExit(1)


if __name__ == "__main__":
    asyncio.run(main())