|
|
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"]}) |
|
|
|
