server/dataset.py

"""Dataset loader — pulls tasks from CentificAIResearch/BA-Agent-Bench HF dataset."""
from __future__ import annotations

import os
import random
from pathlib import Path
from typing import List, Optional

from ba_agent_env.models import GeneratedStory, InputDocument, TaskSample

_DATASET_ID = "CentificAIResearch/BA-Agent-Bench"
_CACHE_FILE = Path(__file__).resolve().parents[1] / "data" / "tasks_cache.parquet"

_TASKS: Optional[List[TaskSample]] = None


def _load_from_hf() -> List[TaskSample]:
    """Pull parquet from HF datasets and convert to TaskSample list."""
    import pyarrow.parquet as pq

    cache = _CACHE_FILE
    cache.parent.mkdir(parents=True, exist_ok=True)
    if not cache.exists():
        from huggingface_hub import hf_hub_download

        downloaded = hf_hub_download(
            repo_id=_DATASET_ID,
            filename="train.parquet",
            repo_type="dataset",
            local_dir=str(cache.parent),
        )
        if downloaded != str(cache):
            try:
                os.replace(downloaded, cache)
            except Exception:
                cache = Path(downloaded)

    table = pq.read_table(str(cache))
    rows = table.to_pylist()
    tasks: List[TaskSample] = []
    for r in rows:
        docs = [
            InputDocument(filename=d.get("filename", ""), content=d.get("content", ""))
            for d in (r.get("input_documents") or [])
        ]
        golden = [
            GeneratedStory(
                story_id=s.get("story_id"),
                title=s.get("title", ""),
                description=s.get("description", ""),
                acceptance_criteria=s.get("acceptance_criteria", ""),
                story_points=s.get("story_points"),
                state=s.get("state"),
            )
            for s in (r.get("golden_stories") or [])
        ]
        tasks.append(
            TaskSample(
                task_id=r.get("task_id", ""),
                title=r.get("title", ""),
                description=r.get("description", ""),
                input_documents=docs,
                golden_stories=golden,
            )
        )
    return tasks


def get_tasks() -> List[TaskSample]:
    global _TASKS
    if _TASKS is None:
        try:
            _TASKS = _load_from_hf()
        except Exception as exc:
            print(f"[BAAgentEnv] Failed to load dataset: {exc}. Falling back to stub task.")
            _TASKS = [
                TaskSample(
                    task_id="STUB-001",
                    title="Stub feature (dataset unavailable)",
                    description="Replace with real task once HF dataset is reachable.",
                    input_documents=[
                        InputDocument(filename="stub.txt", content="Stub document content.")
                    ],
                    golden_stories=[
                        GeneratedStory(
                            story_id="S-001",
                            title="Stub story",
                            description="As a stub user, I want a stub feature, so that I can stub.",
                            acceptance_criteria="Given stub, When stub, Then stub.",
                        )
                    ],
                )
            ]
    return _TASKS


def sample_task(rng: Optional[random.Random] = None) -> TaskSample:
    tasks = get_tasks()
    rng = rng or random
    return rng.choice(tasks)


def get_task_by_id(task_id: str) -> Optional[TaskSample]:
    for t in get_tasks():
        if t.task_id == task_id:
            return t
    return None