server/task_bank.py

"""Dataset-backed task catalog for the Python code-review benchmark."""

from __future__ import annotations

import json
from functools import lru_cache
from pathlib import Path
from typing import Dict, List

try:
    from ..models import CodeReviewSnippet, Difficulty, TaskMetadata
except ImportError:
    from models import CodeReviewSnippet, Difficulty, TaskMetadata  # type: ignore


DATA_DIR = Path(__file__).with_name("data")


TASK_DEFINITIONS: Dict[str, dict[str, object]] = {
    "task_easy": {
        "name": "Style & Convention Review",
        "difficulty": Difficulty.EASY,
        "description": "Find style, naming, formatting, and documentation issues.",
        "filename": "snippets_easy.json",
        "max_steps": 25,
    },
    "task_medium": {
        "name": "Logic Bug Detection",
        "difficulty": Difficulty.MEDIUM,
        "description": "Identify correctness issues in ordinary Python code.",
        "filename": "snippets_medium.json",
        "max_steps": 25,
    },
    "task_hard": {
        "name": "Security Vulnerability Audit",
        "difficulty": Difficulty.HARD,
        "description": "Review web and data-processing code for security flaws.",
        "filename": "snippets_hard.json",
        "max_steps": 25,
    },
}


@lru_cache(maxsize=1)
def load_task_bank() -> Dict[str, List[CodeReviewSnippet]]:
    """Load and validate all snippet JSON files."""

    task_bank: Dict[str, List[CodeReviewSnippet]] = {}
    for task_id, spec in TASK_DEFINITIONS.items():
        raw_items = json.loads((DATA_DIR / str(spec["filename"])).read_text(encoding="utf-8"))
        task_bank[task_id] = [CodeReviewSnippet.model_validate(item) for item in raw_items]
    return task_bank


@lru_cache(maxsize=1)
def load_task_catalog() -> List[TaskMetadata]:
    """Return visible task metadata for `/tasks` and environment resets."""

    task_bank = load_task_bank()
    catalog: List[TaskMetadata] = []
    for task_id, spec in TASK_DEFINITIONS.items():
        catalog.append(
            TaskMetadata(
                task_id=task_id,
                name=str(spec["name"]),
                difficulty=spec["difficulty"],  # type: ignore[arg-type]
                description=str(spec["description"]),
                snippet_count=len(task_bank[task_id]),
                max_steps=int(spec["max_steps"]),
                min_score=0.0,
                max_score=1.0,
            )
        )
    return catalog


def get_task_metadata(task_id: str) -> TaskMetadata:
    """Return task metadata for one family."""

    for task in load_task_catalog():
        if task.task_id == task_id:
            return task
    raise KeyError(f"Unknown task_id: {task_id}")