data_studio/task_extensions.py

from __future__ import annotations

from typing import Any


TASK_EXTENSION_REQUIREMENTS = frozenset(
    {
        "rubric",
        "failure_taxonomy",
        "sampling_policy",
        "cleaner_policy",
        "audit_policy",
        "export_schema",
    }
)


def validate_task_extension(task: str, artifacts: dict[str, Any]) -> dict[str, Any]:
    missing = sorted(requirement for requirement in TASK_EXTENSION_REQUIREMENTS if not artifacts.get(requirement))
    if missing:
        raise ValueError(f"Task extension {task} is blocked; missing artifacts: {missing}")
    return {"task": task, "status": "ready_for_controlled_pilot", "artifacts": artifacts}