Spaces:
Sleeping
Sleeping
| """Framework-neutral SpreadsheetBench episode lifecycle: SkillOpt's multi-turn code repair. | |
| Mirrors ``skillopt.envs.spreadsheetbench.codegen_agent`` in ``mode=multi``: every turn the agent | |
| submits one complete script, the harness executes it on the first hidden case, and a failing | |
| execution comes back as feedback for another turn, up to 30 turns within a 600 s task budget. | |
| The episode ends on the first successful execution (evaluation) or, during training, on the first | |
| execution whose first-case output grades correct; a wrong output is fed back as SkillOpt's | |
| cell-level report with expected values hidden. The final script is then graded independently on | |
| every hidden case, as SkillOpt's ``process_one`` does. | |
| """ | |
| from __future__ import annotations | |
| from collections.abc import Callable | |
| from dataclasses import dataclass | |
| from pathlib import Path | |
| import tempfile | |
| import time | |
| from typing import Protocol | |
| from uuid import uuid4 | |
| from spreadsheetbench_env.catalog import SpreadsheetTask, get_task | |
| from spreadsheetbench_env.executor import ExecutionResult, ScriptExecutor | |
| from spreadsheetbench_env.grader import InfraError, compare_workbooks, verification_report | |
| from spreadsheetbench_env.interface import ( | |
| DEADLINE_MARGIN_S, | |
| STEP_LIMIT, | |
| TASK_TIME_BUDGET_S, | |
| build_task_prompt, | |
| eval_feedback, | |
| execution_error_feedback, | |
| extract_script, | |
| preview_workbook, | |
| ) | |
| class CoreObservation: | |
| task_id: int | |
| text: str | |
| goal: str | |
| task_type: str | |
| case_count: int | |
| answer_position: str | |
| success: bool | |
| progress: float | |
| step_limit: int = STEP_LIMIT | |
| class CoreStep: | |
| observation: CoreObservation | |
| reward: float | |
| terminated: bool | |
| truncated: bool | |
| class CoreState: | |
| task_id: int | None | |
| goal: str | None | |
| reward: float | |
| success: bool | |
| progress: float | |
| terminated: bool | |
| truncated: bool | |
| step_count: int | |
| step_limit: int | |
| diagnostics: dict[str, str] | |
| closed: bool | |
| class _SessionExecutor(Protocol): | |
| def root(self) -> Path: ... | |
| def run( | |
| self, | |
| code: str, | |
| input_path: str | Path, | |
| *, | |
| case_name: str, | |
| ) -> ExecutionResult: ... | |
| def remove_case(self, case_name: str) -> None: ... | |
| def close(self) -> None: ... | |
| _TaskLookup = Callable[[int], SpreadsheetTask] | |
| _ExecutorFactory = Callable[[Path], _SessionExecutor] | |
| def _default_executor_factory(root: Path) -> _SessionExecutor: | |
| return ScriptExecutor(root) | |
| class SpreadsheetBenchCore: | |
| """Own one isolated multi-turn SpreadsheetBench episode.""" | |
| def __init__( | |
| self, | |
| data_root: str | Path, | |
| *, | |
| task_lookup: _TaskLookup = get_task, | |
| executor_factory: _ExecutorFactory = _default_executor_factory, | |
| work_root: str | Path | None = None, | |
| clock: Callable[[], float] = time.monotonic, | |
| ) -> None: | |
| self._data_root = Path(data_root).resolve() | |
| if not self._data_root.is_dir(): | |
| raise FileNotFoundError( | |
| f"SpreadsheetBench runtime data root does not exist: {self._data_root}" | |
| ) | |
| self._task_lookup = task_lookup | |
| self._executor_factory = executor_factory | |
| self._work_root = Path(work_root or tempfile.gettempdir()).resolve() | |
| self._clock = clock | |
| self._executor: _SessionExecutor | None = None | |
| self._task: SpreadsheetTask | None = None | |
| self._cases: tuple[tuple[Path, Path], ...] = () | |
| self._prompt = "" | |
| self._preview = "" | |
| self._training_feedback = False | |
| self._deadline = 0.0 | |
| self._last_code = "" | |
| self._reward = 0.0 | |
| self._progress = 0.0 | |
| self._terminated = False | |
| self._truncated = False | |
| self._step_count = 0 | |
| self._diagnostics: dict[str, str] = {} | |
| self._closed = False | |
| def _runtime_path(self, relative: str, *, hidden_answer: bool) -> Path: | |
| candidate = (self._data_root / relative).resolve() | |
| if not candidate.is_relative_to(self._data_root): | |
| raise InfraError(f"runtime path escapes data root: {relative!r}") | |
| if not candidate.is_file(): | |
| kind = "hidden answer" if hidden_answer else "input workbook" | |
| raise InfraError(f"{kind} is missing: {relative}") | |
| return candidate | |
| def reset(self, task_id: int, *, training_feedback: bool) -> CoreObservation: | |
| """Start ``task_id``. ``training_feedback`` enables SkillOpt's training-only gold check: | |
| a successfully executed but wrong first-case output is fed back instead of ending the | |
| episode. Evaluation must pass ``False`` so no grading evidence reaches the agent.""" | |
| self._ensure_open() | |
| if not isinstance(training_feedback, bool): | |
| raise TypeError("training_feedback must be a bool") | |
| if self._executor is not None: | |
| self._executor.close() | |
| task = self._task_lookup(int(task_id)) | |
| cases = tuple( | |
| ( | |
| self._runtime_path(case.input_path, hidden_answer=False), | |
| self._runtime_path(case.answer_path, hidden_answer=True), | |
| ) | |
| for case in task.cases | |
| ) | |
| if not cases: | |
| raise InfraError(f"task {task.task_id} has no runtime cases") | |
| session_root = self._work_root / f"spreadsheetbench-session-{uuid4().hex}" | |
| executor = self._executor_factory(session_root) | |
| try: | |
| preview = preview_workbook(cases[0][0]) | |
| except Exception as error: | |
| executor.close() | |
| raise InfraError( | |
| f"could not preview input workbook for task {task.task_id}: {error}" | |
| ) from error | |
| self._executor = executor | |
| self._task = task | |
| self._cases = cases | |
| self._preview = preview | |
| self._prompt = build_task_prompt( | |
| task.instruction, | |
| task.task_type, | |
| task.answer_position, | |
| preview, | |
| ) | |
| self._training_feedback = training_feedback | |
| self._deadline = self._clock() + TASK_TIME_BUDGET_S | |
| self._last_code = "" | |
| self._reward = 0.0 | |
| self._progress = 0.0 | |
| self._terminated = False | |
| self._truncated = False | |
| self._step_count = 0 | |
| self._diagnostics = {} | |
| return self._observation(self._prompt) | |
| def step(self, action_code: str) -> CoreStep: | |
| self._ensure_open() | |
| if self._task is None or self._executor is None or not self._cases: | |
| raise RuntimeError("SpreadsheetBench core must be reset before step") | |
| if self._terminated or self._truncated: | |
| raise RuntimeError("SpreadsheetBench episode has ended") | |
| turn = self._step_count + 1 | |
| self._step_count = turn | |
| if self._deadline - self._clock() <= DEADLINE_MARGIN_S: | |
| # SkillOpt stops issuing turns here, so the script produced on this turn never | |
| # existed for it: the previous script is the one that gets graded. | |
| return self._finish( | |
| self._last_code, | |
| truncated=True, | |
| preface=( | |
| f"The {TASK_TIME_BUDGET_S:g} s task time budget was exhausted before this " | |
| "turn; the previous script is graded." | |
| ), | |
| ) | |
| code = extract_script(action_code) | |
| self._last_code = code | |
| first_input, first_answer = self._cases[0] | |
| case_name = f"turn-{turn}" | |
| try: | |
| execution = self._executor.run(code, first_input, case_name=case_name) | |
| if execution.ok and execution.output_path is not None: | |
| if not self._training_feedback: | |
| return self._finish(code, truncated=False) | |
| matched, _ = compare_workbooks( | |
| first_answer, execution.output_path, self._task.answer_position | |
| ) | |
| if matched: | |
| return self._finish(code, truncated=False) | |
| feedback = eval_feedback( | |
| verification_report( | |
| execution.output_path, first_answer, self._task.answer_position | |
| ) | |
| ) | |
| else: | |
| feedback = execution_error_feedback(execution.error) | |
| finally: | |
| self._executor.remove_case(case_name) | |
| if turn >= STEP_LIMIT: | |
| return self._finish(code, truncated=True, preface=feedback) | |
| return CoreStep( | |
| observation=self._observation(feedback), | |
| reward=0.0, | |
| terminated=False, | |
| truncated=False, | |
| ) | |
| def _finish(self, code: str, *, truncated: bool, preface: str = "") -> CoreStep: | |
| """Grade ``code`` independently on every hidden case and end the episode.""" | |
| assert self._task is not None and self._executor is not None | |
| summaries: list[str] = [] | |
| verification: list[str] = [] | |
| passed = 0 | |
| fail_reason = "" | |
| for number, (input_path, answer_path) in enumerate(self._cases, start=1): | |
| case_name = f"case-{number}" | |
| try: | |
| execution = self._executor.run(code, input_path, case_name=case_name) | |
| if not execution.ok or execution.output_path is None: | |
| summaries.append(f"Case {number}: execution ERROR") | |
| verification.append( | |
| f"## Execution (case {number})\nERROR: {execution.error}" | |
| ) | |
| if not fail_reason: | |
| final_line = (execution.error.strip().splitlines() or ["unknown"])[-1] | |
| fail_reason = f"exec-error: {final_line[:200]}" | |
| if execution.timed_out: | |
| break | |
| continue | |
| matched, reason = compare_workbooks( | |
| answer_path, | |
| execution.output_path, | |
| self._task.answer_position, | |
| ) | |
| summaries.append( | |
| f"Case {number}: execution ok; grade {'PASS' if matched else 'FAIL'}" | |
| ) | |
| if matched: | |
| passed += 1 | |
| elif not fail_reason: | |
| fail_reason = f"eval-mismatch: {reason[:200]}" | |
| verification.append( | |
| f"## Eval Result (case {number}): {'PASS' if matched else 'FAIL'}\n" | |
| f"{reason}\n\n" | |
| f"{verification_report(execution.output_path, answer_path, self._task.answer_position)}" | |
| ) | |
| finally: | |
| self._executor.remove_case(case_name) | |
| self._terminated = not truncated | |
| self._truncated = truncated | |
| self._progress = passed / len(self._cases) | |
| self._reward = 1.0 if passed == len(self._cases) else 0.0 | |
| self._diagnostics = { | |
| "verification": "[POST-EXECUTION VERIFICATION]\n\n" | |
| + "\n\n---\n\n".join(verification), | |
| "spreadsheet_preview": self._preview, | |
| "fail_reason": "" if self._reward == 1.0 else fail_reason, | |
| "code": code, | |
| "turns": str(self._step_count), | |
| "training_feedback": "true" if self._training_feedback else "false", | |
| } | |
| text = "\n".join(summaries) | |
| if preface: | |
| text = f"{preface}\n\n{text}" | |
| return CoreStep( | |
| observation=self._observation(text), | |
| reward=self._reward, | |
| terminated=self._terminated, | |
| truncated=self._truncated, | |
| ) | |
| def state(self) -> CoreState: | |
| return CoreState( | |
| task_id=None if self._task is None else self._task.task_id, | |
| goal=None if self._task is None else self._task.instruction, | |
| reward=self._reward, | |
| success=self._reward == 1.0 and (self._terminated or self._truncated), | |
| progress=self._progress, | |
| terminated=self._terminated, | |
| truncated=self._truncated, | |
| step_count=self._step_count, | |
| step_limit=STEP_LIMIT, | |
| diagnostics=dict(self._diagnostics), | |
| closed=self._closed, | |
| ) | |
| def close(self) -> None: | |
| if self._closed: | |
| return | |
| if self._executor is not None: | |
| self._executor.close() | |
| self._closed = True | |
| def _observation(self, text: str) -> CoreObservation: | |
| if self._task is None: | |
| raise RuntimeError("SpreadsheetBench core has no current task") | |
| return CoreObservation( | |
| task_id=self._task.task_id, | |
| text=text, | |
| goal=self._task.instruction, | |
| task_type=self._task.task_type, | |
| case_count=len(self._cases), | |
| answer_position=self._task.answer_position, | |
| success=self._reward == 1.0 and (self._terminated or self._truncated), | |
| progress=self._progress, | |
| ) | |
| def _ensure_open(self) -> None: | |
| if self._closed: | |
| raise RuntimeError("SpreadsheetBench core is closed") | |