Spaces:
Sleeping
Sleeping
File size: 5,210 Bytes
e5e36db | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 | """
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()
|