Spaces:
Sleeping
Sleeping
| """ | |
| tasks/registry.py | |
| ================= | |
| Deterministic task registry for InformationIntegrityEnv. | |
| Loads pre-built scenario datasets from tasks/{difficulty}/scenarios.json. | |
| NO external API calls — all scenarios are static and reproducible. | |
| Usage | |
| ----- | |
| from tasks.registry import TASKS | |
| task = TASKS["easy"] | |
| scenario = task.get_scenario(0) # by index | |
| scenario = task.get_scenario("easy_001") # by id | |
| all_sc = task.get_all_scenarios() | |
| # Random selection (seeded for reproducibility) | |
| import random | |
| rng = random.Random(42) | |
| scenario = task.sample(rng) | |
| """ | |
| from __future__ import annotations | |
| import json | |
| import os | |
| import random | |
| from typing import Dict, List, Optional | |
| # --------------------------------------------------------------------------- | |
| # Path resolution | |
| # --------------------------------------------------------------------------- | |
| _TASKS_DIR = os.path.dirname(os.path.abspath(__file__)) | |
| def _load_scenarios(difficulty: str) -> List[dict]: | |
| path = os.path.join(_TASKS_DIR, difficulty, "scenarios.json") | |
| if not os.path.exists(path): | |
| raise FileNotFoundError( | |
| f"Scenario dataset not found: {path}\n" | |
| f"Run data_collector.py to build datasets, or check tasks/{difficulty}/scenarios.json" | |
| ) | |
| with open(path, "r", encoding="utf-8") as fh: | |
| data = json.load(fh) | |
| if not isinstance(data, list) or not data: | |
| raise ValueError(f"scenarios.json for '{difficulty}' must be a non-empty JSON array") | |
| return data | |
| # --------------------------------------------------------------------------- | |
| # Task class | |
| # --------------------------------------------------------------------------- | |
| class Task: | |
| """ | |
| Wraps one difficulty-level dataset. | |
| Parameters | |
| ---------- | |
| difficulty : str | |
| One of 'easy', 'medium', 'hard'. | |
| """ | |
| def __init__(self, difficulty: str) -> None: | |
| self.difficulty = difficulty | |
| self._scenarios: List[dict] = _load_scenarios(difficulty) | |
| self._index_by_id: Dict[str, dict] = {s["id"]: s for s in self._scenarios} | |
| # ------------------------------------------------------------------ | |
| # Public API | |
| # ------------------------------------------------------------------ | |
| def get_scenario(self, scenario_id: str | int) -> dict: | |
| """ | |
| Return a scenario by index (int) or scenario id (str). | |
| Returns a fresh copy so the env cannot mutate the dataset. | |
| """ | |
| if isinstance(scenario_id, int): | |
| if scenario_id < 0 or scenario_id >= len(self._scenarios): | |
| raise IndexError( | |
| f"Scenario index {scenario_id} out of range " | |
| f"(dataset has {len(self._scenarios)} scenarios)" | |
| ) | |
| return dict(self._scenarios[scenario_id]) | |
| # String id lookup | |
| if scenario_id not in self._index_by_id: | |
| raise KeyError( | |
| f"Scenario id '{scenario_id}' not found in {self.difficulty} dataset. " | |
| f"Available ids: {list(self._index_by_id.keys())}" | |
| ) | |
| return dict(self._index_by_id[scenario_id]) | |
| def get_all_scenarios(self) -> List[dict]: | |
| """Return all scenarios as a list of dicts (copies).""" | |
| return [dict(s) for s in self._scenarios] | |
| def sample(self, rng: Optional[random.Random] = None) -> dict: | |
| """Return a uniformly random scenario (optionally seeded).""" | |
| rng = rng or random.Random() | |
| return dict(rng.choice(self._scenarios)) | |
| def __len__(self) -> int: | |
| return len(self._scenarios) | |
| def __repr__(self) -> str: | |
| return f"Task(difficulty={self.difficulty!r}, n_scenarios={len(self)})" | |
| # --------------------------------------------------------------------------- | |
| # Lazy-loaded registry singleton | |
| # --------------------------------------------------------------------------- | |
| class _TaskRegistry: | |
| """ | |
| Lazy-loading registry. Datasets are only read from disk on first access, | |
| so importing this module does not cause IO at import time. | |
| """ | |
| _cache: Dict[str, Task] = {} | |
| _SUPPORTED = ("easy", "medium", "hard") | |
| def __getitem__(self, difficulty: str) -> Task: | |
| key = difficulty.lower().strip() | |
| if key not in self._SUPPORTED: | |
| raise KeyError( | |
| f"Unknown difficulty '{difficulty}'. " | |
| f"Supported values: {list(self._SUPPORTED)}" | |
| ) | |
| if key not in self._cache: | |
| self._cache[key] = Task(key) | |
| return self._cache[key] | |
| def __contains__(self, difficulty: str) -> bool: | |
| return difficulty.lower().strip() in self._SUPPORTED | |
| def keys(self): | |
| return self._SUPPORTED | |
| def items(self): | |
| return [(d, self[d]) for d in self._SUPPORTED] | |
| def reload(self) -> None: | |
| """Force reload all datasets from disk (useful for tests).""" | |
| self._cache.clear() | |
| def __repr__(self) -> str: | |
| loaded = list(self._cache.keys()) | |
| return f"TaskRegistry(supported={list(self._SUPPORTED)}, loaded={loaded})" | |
| # Singleton — import and use directly: from tasks.registry import TASKS | |
| TASKS = _TaskRegistry() | |