cactus183/patchbench-dev

1import json2import random3from pathlib import Path4from typing import Any5 6from patchbench.grader import grade_patch7from patchbench.models import PatchBenchAction, PatchBenchObservation8 9TASKS_DIR = Path(__file__).parent / "tasks"10 11 12class PatchBenchEnv:13    """OpenEnv-compliant environment for bug-fixing tasks."""14 15    def __init__(self, tasks_dir: Path | None = None):16        self.tasks_dir = tasks_dir or TASKS_DIR17        self._tasks: dict[str, dict[str, Any]] = {}18        self._load_tasks()19 20        self.current_task_id: str | None = None21        self.current_code: str = ""22        self.step_number: int = 023        self.last_reward: float = 0.024        self.done: bool = False25        self._rng = random.Random()26 27    def _load_tasks(self) -> None:28        """Load all tasks from disk into memory."""29        if not self.tasks_dir.exists():30            raise RuntimeError(f"Tasks directory not found: {self.tasks_dir}")31 32        for difficulty_dir in sorted(self.tasks_dir.iterdir()):33            if not difficulty_dir.is_dir():34                continue35            for task_dir in sorted(difficulty_dir.iterdir()):36                if not task_dir.is_dir():37                    continue38                task_json_path = task_dir / "task.json"39                buggy_path = task_dir / "buggy_code.py"40                test_path = task_dir / "test_code.py"41                if not (task_json_path.exists() and buggy_path.exists() and test_path.exists()):42                    continue43                with open(task_json_path) as f:44                    task_meta = json.load(f)45                task_id = task_meta["task_id"]46                self._tasks[task_id] = {47                    "meta": task_meta,48                    "buggy_code": buggy_path.read_text(),49                    "test_code": test_path.read_text(),50                    "baseline_passing": set(task_meta.get("baseline_passing_tests", [])),51                    "baseline_failing": set(task_meta.get("baseline_failing_tests", [])),52                }53 54        if not self._tasks:55            raise RuntimeError(f"No tasks found in {self.tasks_dir}")56 57    @property58    def task_ids(self) -> list[str]:59        return sorted(self._tasks.keys())60 61    def reset(62        self,63        seed: int | None = None,64        task_id: str | None = None,65        **kwargs,66    ) -> PatchBenchObservation:67        if seed is not None:68            self._rng.seed(seed)69 70        if task_id is None:71            task_id = self._rng.choice(self.task_ids)72        if task_id not in self._tasks:73            raise ValueError(f"Unknown task_id: {task_id}. Available: {self.task_ids}")74 75        self.current_task_id = task_id76        task = self._tasks[task_id]77        self.current_code = task["buggy_code"]78        self.step_number = 079        self.last_reward = 0.080        self.done = False81 82        initial_failing = "\n".join(sorted(task["baseline_failing"])) or "(no baseline failures recorded)"83 84        return PatchBenchObservation(85            task_description=task["meta"]["description"],86            buggy_code=self.current_code,87            failing_tests=f"FAILING TESTS:\n{initial_failing}",88            step_number=0,89            max_steps=task["meta"]["max_steps"],90            reward=0.0,91            done=False,92            info={"task_id": task_id, "difficulty": task["meta"]["difficulty"]},93        )94 95    def step(self, action: PatchBenchAction, **kwargs) -> PatchBenchObservation:96        if self.current_task_id is None:97            raise RuntimeError("Must call reset() before step()")98        if self.done:99            task = self._tasks[self.current_task_id]100            return PatchBenchObservation(101                task_description=task["meta"]["description"],102                buggy_code=self.current_code,103                failing_tests="(episode already done)",104                step_number=self.step_number,105                max_steps=task["meta"]["max_steps"],106                reward=0.0,107                done=True,108                info={"terminal": True},109            )110 111        self.step_number += 1112        task = self._tasks[self.current_task_id]113 114        reward, info = grade_patch(115            patched_code=action.patched_code,116            test_code=task["test_code"],117            baseline_passing=task["baseline_passing"],118            baseline_failing=task["baseline_failing"],119        )120 121        if info.get("is_valid_python"):122            self.current_code = action.patched_code123 124        self.last_reward = reward125        terminal = info.get("all_tests_pass", False) or self.step_number >= task["meta"]["max_steps"]126        self.done = terminal127 128        failing_summary = (129            "All tests passing." if info.get("all_tests_pass")130            else f"{info.get('tests_failing', '?')} test(s) still failing. {info.get('newly_passing', 0)} newly passing, {info.get('regressions', 0)} regressions."131        )132 133        info["task_id"] = self.current_task_id134        info["difficulty"] = task["meta"]["difficulty"]135 136        return PatchBenchObservation(137            task_description=task["meta"]["description"],138            buggy_code=self.current_code,139            failing_tests=failing_summary,140            step_number=self.step_number,141            max_steps=task["meta"]["max_steps"],142            reward=reward,143            done=terminal,144            info=info,145        )146 147    @property148    def state(self) -> dict[str, Any]:149        return {150            "current_task_id": self.current_task_id,151            "step_number": self.step_number,152            "current_code": self.current_code,153            "last_reward": self.last_reward,154            "done": self.done,155            "available_tasks": self.task_ids,156        }157 158    def close(self) -> None:159        """Cleanup hook for OpenEnv compatibility."""160        pass161