commit

environment/env.py

@@ -1,135 +1,61 @@
-"""
-CodeReviewEnv — main environment logic.
-Manages state, episode flow, reward accumulation, and per-step grading.
-"""
-from typing import Optional, List, Dict, Any
-
-from environment.models import (
-    CodeReviewAction,
-    CodeReviewObservation,
-    StepResult,
-    ResetResult,
-    StateResult,
-)
-from environment.tasks import get_task, list_tasks as _list_tasks
+from uuid import uuid4
+from typing import Optional
+from openenv.core.env_server.interfaces import Environment
+from openenv.core.env_server.types import State
+from environment.models import CodeReviewAction, CodeReviewObservation
+from environment.tasks import get_task
 from environment.graders import grade
 
 
-class CodeReviewEnv:
-    """
-    OpenEnv-compliant code-review environment.
+class CodeReviewEnv(Environment):
 
-    Episode flow
-    ────────────
-    reset(task_id) → observation
-    step(action)   → (observation, reward, done, info)   [repeated ≤ max_steps]
-    state()        → current state snapshot
-    """
+    SUPPORTS_CONCURRENT_SESSIONS: bool = False
 
-    def __init__(self) -> None:
-        self._task: Optional[dict] = None
-        self._step_number: int = 0
-        self._total_reward: float = 0.0
-        self._actions_history: List[Dict[str, Any]] = []
-        self._done: bool = False
-        self._initialized: bool = False
+    def __init__(self):
+        self._state = State(episode_id=str(uuid4()), step_count=0)
+        self._task = get_task("easy_syntax")
         self._last_feedback: Optional[str] = None
 
-    # ── public properties ────────────────────────────────────────────────────
-
-    @property
-    def is_initialized(self) -> bool:
-        return self._initialized
-
-    # ── core API ─────────────────────────────────────────────────────────────
-
-    def reset(self, task_id: Optional[str] = None) -> ResetResult:
-        """Start a new episode. Defaults to the easy task."""
-        if task_id is None:
-            task_id = "easy_syntax"
-
-        self._task = get_task(task_id)
-        self._step_number = 0
-        self._total_reward = 0.0
-        self._actions_history = []
-        self._done = False
-        self._initialized = True
+    def reset(self) -> CodeReviewObservation:
+        self._task = get_task("easy_syntax")
+        self._state = State(episode_id=str(uuid4()), step_count=0)
         self._last_feedback = None
-
-        obs = self._make_observation()
-        return ResetResult(observation=obs)
-
-    def step(self, action: CodeReviewAction) -> StepResult:
-        """Process one agent action and return (observation, reward, done, info)."""
-        if not self._initialized or self._done:
-            raise RuntimeError("Call reset() before stepping, or episode is over.")
-
-        self._step_number += 1
-        task_id = self._task["task_id"]
-        ground_truth = self._task["ground_truth"]
-
-        # Grade the action
-        reward, feedback = grade(task_id, action, ground_truth)
-        self._last_feedback = feedback
-        self._total_reward += reward
-
-        # Record history
-        self._actions_history.append(
-            {
-                "step": self._step_number,
-                "num_issues_reported": len(action.identified_issues),
-                "has_fix": action.suggested_fix is not None,
-                "reward": reward,
-            }
+        return CodeReviewObservation(
+            task_id=self._task["task_id"],
+            task_name=self._task["task_name"],
+            difficulty=self._task["difficulty"],
+            language=self._task["language"],
+            code_snippet=self._task["code_snippet"],
+            context=self._task["context"],
+            step_number=0,
+            max_steps=self._task["max_steps"],
+            previous_feedback=None,
+            done=False,
+            reward=0.0,
         )
 
-        # Episode ends when: agent says done, reward is perfect, or max steps reached
+    def step(self, action: CodeReviewAction) -> CodeReviewObservation:
+        self._state.step_count += 1
+        reward, feedback = grade(
+            self._task["task_id"], action, self._task["ground_truth"]
+        )
+        self._last_feedback = feedback
         max_steps = self._task["max_steps"]
-        done = action.done or reward >= 0.95 or self._step_number >= max_steps
-        self._done = done
-
-        obs = self._make_observation()
-
-        return StepResult(
-            observation=obs,
-            reward=reward,
+        done = action.submit or reward >= 0.95 or self._state.step_count >= max_steps
+        return CodeReviewObservation(
+            task_id=self._task["task_id"],
+            task_name=self._task["task_name"],
+            difficulty=self._task["difficulty"],
+            language=self._task["language"],
+            code_snippet=self._task["code_snippet"],
+            context=self._task["context"],
+            step_number=self._state.step_count,
+            max_steps=max_steps,
+            previous_feedback=feedback,
             done=done,
-            info={
-                "feedback": feedback,
-                "step": self._step_number,
-                "total_reward": round(self._total_reward, 4),
-                "cumulative_score": round(
-                    self._total_reward / max(self._step_number, 1), 4
-                ),
-            },
-        )
-
-    def get_state(self) -> StateResult:
-        """Return a snapshot of the current episode state."""
-        return StateResult(
-            task_id=self._task["task_id"] if self._task else "",
-            step_number=self._step_number,
-            total_reward=round(self._total_reward, 4),
-            actions_history=self._actions_history,
-            done=self._done,
-            initialized=self._initialized,
+            reward=reward,
         )
 
-    def list_tasks(self) -> list:
-        return _list_tasks()
-
-    # ── internal helpers ─────────────────────────────────────────────────────
-
-    def _make_observation(self) -> CodeReviewObservation:
-        t = self._task
-        return CodeReviewObservation(
-            task_id=t["task_id"],
-            task_name=t["task_name"],
-            difficulty=t["difficulty"],
-            language=t["language"],
-            code_snippet=t["code_snippet"],
-            context=t["context"],
-            step_number=self._step_number,
-            max_steps=t["max_steps"],
-            previous_feedback=self._last_feedback,
-        )
+    @property
+    def state(self) -> State:
+        return self._state