task_loader.py

"""Load Terminal-Bench task data from disk."""

import os
from pathlib import Path
from typing import Any

try:
    import tomli as toml
except ImportError:
    import tomllib as toml

_DEFAULT_TB2_DIR = Path(__file__).parent / "data" / "terminal-bench-2-main"
TB2_DIR = Path(os.getenv("TB2_DIR") or str(_DEFAULT_TB2_DIR))


def load_task(task_name: str) -> dict[str, Any]:
    """
    Load a single task's metadata and instruction.

    Returns:
        {
            "task_name": str,
            "problem_statement": str,
            "metadata": dict (from task.toml),
            "docker_image": str | None (from task.toml environment.docker_image),
        }
    """
    task_dir = TB2_DIR / task_name
    if not task_dir.exists():
        raise FileNotFoundError(f"Task directory not found: {task_dir}")

    instruction_file = task_dir / "instruction.md"
    toml_file = task_dir / "task.toml"

    if not instruction_file.exists():
        raise FileNotFoundError(f"instruction.md not found for task: {task_name}")

    problem_statement = instruction_file.read_text(encoding="utf-8").strip()

    metadata = {}
    docker_image = None
    if toml_file.exists():
        with open(toml_file, "rb") as f:
            metadata = toml.load(f)
            # Extract docker_image from environment section
            if "environment" in metadata and "docker_image" in metadata["environment"]:
                docker_image = metadata["environment"]["docker_image"]

    return {
        "task_name": task_name,
        "problem_statement": problem_statement,
        "metadata": metadata,
        "docker_image": docker_image,
    }