evaluation.py

"""Concurrent dry-run evaluation with atomic per-task checkpoints."""

from __future__ import annotations

import time
from collections.abc import Callable
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Any

from agent import GaiaAgent
from cache import ResultStore
from gaia_client import GaiaClient, GaiaClientError, GaiaNetworkError

Progress = Callable[[int, int, str], None]


def run_evaluation(
    client: GaiaClient,
    agent: GaiaAgent,
    result_store: ResultStore,
    progress: Progress | None = None,
    *,
    force: bool = False,
    task_ids: set[str] | None = None,
) -> list[dict[str, Any]]:
    tasks = client.get_questions()
    if task_ids is not None:
        tasks = [task for task in tasks if str(task["task_id"]) in task_ids]
        missing = task_ids - {str(task["task_id"]) for task in tasks}
        if missing:
            raise ValueError(f"Unknown task IDs: {', '.join(sorted(missing))}")
    existing = result_store.load()
    completed = 0

    def run_one(task: dict[str, Any]) -> dict[str, Any]:
        task_id = str(task["task_id"])
        if not force and existing.get(task_id, {}).get("status") == "ok":
            return existing[task_id]
        started = time.perf_counter()
        row: dict[str, Any] = {
            "task_id": task_id,
            "question": str(task["question"]),
            "answer": "",
            "task_type": "",
            "status": "agent_error",
            "duration_seconds": 0.0,
            "error": "",
            "evidence": [],
            "confidence": None,
        }
        try:
            attachment = client.download_attachment(task)
            solved = agent.solve_detailed(task, attachment, force=force)
            row.update(
                answer=solved.answer,
                task_type=solved.task_type,
                evidence=solved.evidence,
                confidence=solved.confidence,
                status="ok",
            )
        except GaiaNetworkError as exc:
            row.update(status="network_error", error=f"{type(exc).__name__}: {exc}")
        except GaiaClientError as exc:
            row.update(status="api_error", error=f"{type(exc).__name__}: {exc}")
        except Exception as exc:
            row.update(status="agent_error", error=f"{type(exc).__name__}: {exc}")
        row["duration_seconds"] = round(time.perf_counter() - started, 3)
        result_store.save_result(task_id, row)
        return row

    workers = min(agent.settings.worker_limit, max(1, len(tasks)))
    with ThreadPoolExecutor(max_workers=workers, thread_name_prefix="gaia") as executor:
        futures = {executor.submit(run_one, task): task for task in tasks}
        for future in as_completed(futures):
            completed += 1
            task_id = str(futures[future]["task_id"])
            future.result()  # run_one isolates expected task failures
            if progress:
                progress(completed, len(tasks), task_id)
    ordered_ids = [str(task["task_id"]) for task in tasks]
    return result_store.ordered(ordered_ids)


def submission_answers(
    results: list[dict[str, Any]], expected_count: int | None = 20
) -> list[dict[str, str]]:
    if not results:
        raise ValueError("No dry-run results exist")
    if expected_count is not None and len(results) != expected_count:
        raise ValueError(
            f"Refusing submission: expected exactly {expected_count} tasks, found {len(results)}"
        )
    task_ids = [str(row.get("task_id", "")) for row in results]
    if len(set(task_ids)) != len(task_ids):
        raise ValueError("Refusing submission: duplicate task IDs")
    failures = [row for row in results if row.get("status") != "ok"]
    if failures:
        raise ValueError(f"Refusing submission: {len(failures)} task(s) are incomplete")
    empty = [row for row in results if not str(row.get("answer", "")).strip()]
    if empty:
        raise ValueError(f"Refusing submission: {len(empty)} answer(s) are empty")
    return [
        {"task_id": str(row["task_id"]), "submitted_answer": str(row["answer"])}
        for row in results
    ]