| from __future__ import annotations |
| from typing import Optional |
|
|
| from environment.models import ( |
| Action, Observation, EnvironmentState, |
| StepResponse, ResetResponse, TaskId, ReviewComment, |
| ) |
| from environment.reward import compute_reward |
| from graders.base import BaseGrader |
| from tasks.registry import build_episode |
|
|
|
|
| class CodeSentinelEnv: |
|
|
| def __init__(self): |
| self._task_id: Optional[TaskId] = None |
| self._obs: Optional[Observation] = None |
| self._grader: Optional[BaseGrader] = None |
| self._prev_reward: float = 0.0 |
| self._done: bool = False |
| self._step: int = 0 |
| self._max_steps: int = 20 |
| self._cumulative_reward: float = 0.0 |
|
|
| def reset( |
| self, |
| task_id: TaskId = TaskId.bug_detection, |
| seed: int = 42, |
| ) -> ResetResponse: |
| sample, grader, obs = build_episode(task_id, seed=seed) |
| self._task_id = task_id |
| self._obs = obs |
| self._grader = grader |
| self._prev_reward = 0.0 |
| self._done = False |
| self._step = 0 |
| self._max_steps = obs.max_steps |
| self._cumulative_reward = 0.0 |
| return ResetResponse( |
| observation=obs, |
| info={"task_id": task_id, "seed": seed, "max_steps": obs.max_steps}, |
| ) |
|
|
| def step(self, action: Action) -> StepResponse: |
| if self._obs is None: |
| raise RuntimeError("Call reset() before step().") |
| if self._done: |
| raise RuntimeError("Episode is done. Call reset() first.") |
|
|
| self._step += 1 |
|
|
| terminal_actions = {"approve", "request_changes", "submit_report"} |
| episode_done = ( |
| action.action_type.value in terminal_actions |
| or self._step >= self._max_steps |
| ) |
|
|
| breakdown = self._grader.score_action(action) |
| reward = compute_reward( |
| grader=self._grader, |
| breakdown=breakdown, |
| prev_reward=self._prev_reward, |
| done=episode_done, |
| step=self._step, |
| ) |
|
|
| self._prev_reward = reward.value |
| self._cumulative_reward += reward.value |
| self._done = episode_done |
|
|
| self._obs.history.append(ReviewComment( |
| step = self._step, |
| action = action.action_type, |
| line = action.line, |
| message = action.message or str(action.report or action.reason or ""), |
| severity = action.severity, |
| )) |
| self._obs.step = self._step |
| self._obs.done = episode_done |
|
|
| return StepResponse( |
| observation=self._obs, |
| reward=reward, |
| done=episode_done, |
| info={ |
| "step": self._step, |
| "cumulative_reward": self._cumulative_reward, |
| "final_score": self._grader.final_score() if episode_done else None, |
| }, |
| ) |
|
|
| def state(self) -> EnvironmentState: |
| if self._task_id is None: |
| raise RuntimeError("Call reset() first.") |
| g = self._grader |
| return EnvironmentState( |
| task_id = self._task_id, |
| step = self._step, |
| max_steps = self._max_steps, |
| done = self._done, |
| ground_truth_count = len(g.ground_truth) if isinstance(g.ground_truth, list) else sum(len(v) for v in g.ground_truth.values()), |
| found_count = len(g.found_ids), |
| false_positive_count = g.false_positives, |
| cumulative_reward = self._cumulative_reward, |
| current_score = g.final_score(), |
| ) |
|
|
