supportbench/tasks.py

"""Task loading for SupportBench."""

from __future__ import annotations

import json
import os
from pathlib import Path
from typing import Any, Dict, List, Optional


DATA_DIR = Path(__file__).parent.parent / "data"

_TASK_FILES = {
    "easy": "easy_tasks.json",
    "medium": "medium_tasks.json",
    "hard": "hard_tasks.json",
}

_TASK_ID_MAP: Dict[str, Dict[str, Any]] = {}


def _load_all_tasks() -> None:
    for difficulty, filename in _TASK_FILES.items():
        path = DATA_DIR / filename
        if not path.exists():
            continue
        with open(path) as f:
            tasks: List[Dict[str, Any]] = json.load(f)
        for task in tasks:
            _TASK_ID_MAP[task["task_id"]] = task


_load_all_tasks()


def get_task(task_id: str) -> Dict[str, Any]:
    if task_id not in _TASK_ID_MAP:
        available = list(_TASK_ID_MAP.keys())
        raise ValueError(f"Unknown task_id '{task_id}'. Available: {available}")
    return _TASK_ID_MAP[task_id]


def list_tasks() -> List[Dict[str, Any]]:
    return [
        {
            "task_id": t["task_id"],
            "task_name": t["task_name"],
            "difficulty": t["difficulty"],
        }
        for t in _TASK_ID_MAP.values()
    ]


def get_default_task_id() -> str:
    return "easy_ticket_triage"


AVAILABLE_ACTIONS = [
    "classify_ticket",
    "set_priority",
    "ask_customer",
    "propose_resolution",
    "apply_resolution",
    "escalate",
    "resolve",
]