| import copy |
| import json |
| import random |
| from dataclasses import dataclass |
| from pathlib import Path |
| from typing import Any, Dict, Optional, List |
|
|
|
|
| @dataclass |
| class StepResult: |
| obs: Dict[str, Any] |
| reward: float |
| done: bool |
| info: Dict[str, Any] |
|
|
|
|
| class FlowDebugEnv: |
| """ |
| This is a simple environment made for OpenEnv. |
| Here's what it does: |
| - It gives you information in text or JSON. |
| - You can only do one thing: fix the 'inputs.expression' in a 'Condition_Check' step. |
| - If your fix is exactly right, you win! |
| """ |
| def __init__(self, cases: List[Dict[str, Any]], max_attempts: int = 3, seed: Optional[int] = None): |
| self.cases = cases |
| self.max_attempts = max_attempts |
| self.rng = random.Random(seed) |
| self.current_case: Optional[Dict[str, Any]] = None |
| self.attempts_left = max_attempts |
|
|
| @classmethod |
| def from_json(cls, cases_json_path: str, max_attempts: int = 3, seed: Optional[int] = None): |
| path = Path(cases_json_path) |
| with open(path, "r", encoding="utf-8") as f: |
| cases = json.load(f) |
| return cls(cases=cases, max_attempts=max_attempts, seed=seed) |
|
|
| def reset(self) -> Dict[str, Any]: |
| self.current_case = copy.deepcopy(self.rng.choice(self.cases)) |
| self.attempts_left = self.max_attempts |
| return self._make_observation() |
|
|
| def step(self, action: Dict[str, Any]) -> StepResult: |
| if self.current_case is None: |
| raise RuntimeError("Call reset() before step().") |
|
|
| self.attempts_left -= 1 |
|
|
| if action.get("action") != "patch_step": |
| return self._invalid_action("Unsupported action type") |
|
|
| step_name = action.get("step") |
| field = action.get("field") |
| value = action.get("value") |
|
|
| patched_ok = self._apply_patch(step_name, field, value) |
| if not patched_ok: |
| return self._invalid_action("Patch failed (step/field not found)") |
|
|
| gold = self.current_case["gold_fix"] |
| solved = (step_name == gold["step"] and field == gold["field"] and value == gold["value"]) |
|
|
| if solved: |
| self._mark_success() |
| obs = self._make_observation(run_status="Succeeded", error=None, failed_step=None) |
| return StepResult(obs=obs, reward=1.0, done=True, |
| info={"result": "success", "case_id": self.current_case["case_id"]}) |
|
|
| if self.attempts_left <= 0: |
| obs = self._make_observation() |
| return StepResult(obs=obs, reward=-0.2, done=True, |
| info={"result": "out_of_attempts", "case_id": self.current_case["case_id"]}) |
|
|
| obs = self._make_observation() |
| return StepResult(obs=obs, reward=-0.1, done=False, |
| info={"result": "still_failed", "case_id": self.current_case["case_id"]}) |
|
|
| |
| def _apply_patch(self, step_name: str, field: str, value: str) -> bool: |
| for step in self.current_case["steps"]: |
| if step["name"] == step_name: |
| if field == "inputs.expression": |
| step.setdefault("inputs", {}) |
| step["inputs"]["expression"] = value |
| return True |
| return False |
|
|
| def _mark_success(self): |
| for step in self.current_case["steps"]: |
| step["status"] = "Succeeded" |
|
|
| def _make_observation(self, run_status="Failed", error="keep", failed_step="keep"): |
| if error == "keep": |
| err_obj = self.current_case["error"] |
| else: |
| err_obj = error |
|
|
| if failed_step == "keep": |
| failed = self.current_case["failed_step"] |
| else: |
| failed = failed_step |
|
|
| return { |
| "case_id": self.current_case["case_id"], |
| "run_status": run_status, |
| "failed_step": failed, |
| "error": err_obj, |
| "steps": self.current_case["steps"], |
| "attempts_left": self.attempts_left |
| } |
|
|
| def _invalid_action(self, msg: str) -> StepResult: |
| obs = self._make_observation() |
| done = (self.attempts_left <= 0) |
| return StepResult(obs=obs, reward=-0.1, done=done, |
| info={"result": "invalid_action", "message": msg, "case_id": self.current_case["case_id"]}) |
|
|
