"""Genuine OpenEnv adapter around one SpreadsheetBench core instance.""" from __future__ import annotations from collections.abc import Callable import os from pathlib import Path from typing import Any from uuid import uuid4 from openenv.core.env_server.interfaces import Environment from openenv.core.env_server.types import EnvironmentMetadata from spreadsheetbench_env.catalog import get_split from spreadsheetbench_env.core import CoreObservation, SpreadsheetBenchCore from spreadsheetbench_env.executor import ExecutorLimits, ScriptExecutor from spreadsheetbench_env.interface import environment_description from spreadsheetbench_env.models import ( SpreadsheetBenchAction, SpreadsheetBenchObservation, SpreadsheetBenchState, ) _CoreFactory = Callable[[], SpreadsheetBenchCore] def _optional_integer(name: str) -> int | None: raw = os.getenv(name) if raw is None or not raw.strip(): return None try: return int(raw) except ValueError as error: raise ValueError(f"{name} must be an integer, got {raw!r}") from error def _positive_number(name: str, default: str) -> float: raw = os.getenv(name, default) try: value = float(raw) except ValueError as error: raise ValueError(f"{name} must be numeric, got {raw!r}") from error if value <= 0: raise ValueError(f"{name} must be positive") return value def _default_core_factory() -> SpreadsheetBenchCore: data_root = Path(os.getenv("SPREADSHEETBENCH_DATA", "/opt/spreadsheetbench")) work_root = Path(os.getenv("SPREADSHEETBENCH_WORK_ROOT", "/tmp")) limits = ExecutorLimits( timeout_s=_positive_number("SCRIPT_TIMEOUT_S", "120"), cpu_seconds=int(_positive_number("SCRIPT_CPU_SECONDS", "120")), memory_bytes=int(_positive_number("SCRIPT_MEMORY_BYTES", str(2 * 1024**3))), file_bytes=int(_positive_number("SCRIPT_FILE_BYTES", str(256 * 1024**2))), process_count=int(_positive_number("SCRIPT_PROCESS_COUNT", "64")), open_files=int(_positive_number("SCRIPT_OPEN_FILES", "128")), run_uid_base=_optional_integer("SPREADSHEETBENCH_RUNNER_UID_BASE"), run_gid=_optional_integer("SPREADSHEETBENCH_RUNNER_GID"), ) return SpreadsheetBenchCore( data_root, work_root=work_root, executor_factory=lambda root: ScriptExecutor(root, limits=limits), ) class SpreadsheetBenchEnvironment( Environment[ SpreadsheetBenchAction, SpreadsheetBenchObservation, SpreadsheetBenchState, ] ): """One connection-owned, typed SpreadsheetBench episode.""" SUPPORTS_CONCURRENT_SESSIONS = True REQUIRES_SINGLE_THREAD_EXECUTOR = False def __init__( self, core_factory: _CoreFactory = _default_core_factory, *, default_task_id: int | None = None, ) -> None: super().__init__() self._core_factory = core_factory self._core: SpreadsheetBenchCore | None = None self._default_task_id = default_task_id self._episode_id: str | None = None self._closed = False def _get_core(self) -> SpreadsheetBenchCore: if self._closed: raise RuntimeError("SpreadsheetBench environment is closed") if self._core is None: self._core = self._core_factory() return self._core def get_metadata(self) -> EnvironmentMetadata: return EnvironmentMetadata( name="SpreadsheetBench", description=environment_description(), version="0.1.0", documentation_url="https://arxiv.org/abs/2406.14991", ) def list_splits(self) -> list[dict[str, Any]]: return [ { "name": name, "role": get_split(name).role, "num_tasks": len(get_split(name).tasks), "source_revision": get_split(name).source_revision, "sha256": get_split(name).sha256, } for name in ("development", "test") ] def list_tasks(self, split: str) -> list[dict[str, Any]]: task_split = get_split(split) return [ { "id": task.task_id, "source_id": task.source_id, "instruction": task.instruction, "task_type": task.task_type, "split": task_split.role, "source_revision": task_split.source_revision, "catalog_sha256": task_split.sha256, } for task in task_split.tasks ] def num_tasks(self, split: str) -> int: return len(get_split(split).tasks) def get_task(self, split: str, index: int) -> dict[str, Any]: tasks = self.list_tasks(split) try: return tasks[index] except IndexError as error: raise IndexError( f"SpreadsheetBench {split} task index out of range: {index}" ) from error def get_task_range( self, split: str, start: int | None = None, stop: int | None = None, ) -> list[dict[str, Any]]: return self.list_tasks(split)[slice(start, stop)] def reset( self, seed: int | None = None, episode_id: str | None = None, *, task_id: int | None = None, training_feedback: bool | None = None, **kwargs: Any, ) -> SpreadsheetBenchObservation: del seed, kwargs resolved = self._default_task_id if task_id is None else int(task_id) if resolved is None: raise ValueError("SpreadsheetBench reset requires task_id") if not isinstance(training_feedback, bool): # Chosen per run by the client: SkillOpt training rollouts receive the gold-based # wrong-output feedback, evaluation rollouts never do. raise ValueError("SpreadsheetBench reset requires training_feedback (true or false)") core_observation = self._get_core().reset(resolved, training_feedback=training_feedback) self._episode_id = episode_id or str(uuid4()) return self._observation( core_observation, reward=0.0, terminated=False, truncated=False ) def step( self, action: SpreadsheetBenchAction, timeout_s: float | None = None, **kwargs: Any, ) -> SpreadsheetBenchObservation: del timeout_s, kwargs transition = self._get_core().step(action.code) return self._observation( transition.observation, reward=transition.reward, terminated=transition.terminated, truncated=transition.truncated, ) @property def state(self) -> SpreadsheetBenchState: if self._core is None: return SpreadsheetBenchState( episode_id=self._episode_id, step_count=0, closed=self._closed, ) state = self._core.state() return SpreadsheetBenchState( episode_id=self._episode_id, step_count=state.step_count, task_id=state.task_id, goal=state.goal, reward=state.reward, success=state.success, progress=state.progress, terminated=state.terminated, truncated=state.truncated, step_limit=state.step_limit, diagnostics=state.diagnostics, closed=state.closed, ) def close(self) -> None: if self._closed: return if self._core is not None: self._core.close() self._closed = True def _observation( self, observation: CoreObservation, *, reward: float, terminated: bool, truncated: bool, ) -> SpreadsheetBenchObservation: return SpreadsheetBenchObservation( task_id=observation.task_id, text=observation.text, goal=observation.goal, task_type=observation.task_type, case_count=observation.case_count, answer_position=observation.answer_position, success=observation.success, progress=observation.progress, terminated=terminated, truncated=truncated, done=terminated or truncated, reward=reward, step_limit=observation.step_limit, )