Shumaque/code-review-agent

1"""Code Review Environment implementation using openenv-core."""2 3from __future__ import annotations4 5from typing import Any, Optional6 7from openenv.core.env_server.interfaces import Environment8from openenv.core.env_server.types import Action, Observation, State9 10from server.models import CodeReviewAction, CodeReviewObservation, CodeReviewState11from graders.grader import grade12from tasks.seeds import TASK_REGISTRY13 14 15MAX_STEPS = 516 17 18class CodeReviewEnv(19    Environment[CodeReviewAction, CodeReviewObservation, CodeReviewState]20):21    """22    OpenEnv-compliant environment for AI-driven pull request review.23 24    The agent receives a code diff and must submit findings (bugs, security25    issues, style violations). Graders compare findings against seeded ground26    truth using F1 over precision/recall.27    """28 29    def __init__(30        self,31        task_name: str = "detect_logic_bug",32        transform: Optional[Any] = None,33        rubric: Optional[Any] = None,34    ) -> None:35        super().__init__(transform=transform, rubric=rubric)36 37        if task_name not in TASK_REGISTRY:38            raise ValueError(39                f"Unknown task '{task_name}'. Valid: {list(TASK_REGISTRY)}"40            )41 42        self._task_name = task_name43        self._seed = TASK_REGISTRY[task_name]44        self._step_count = 045        self._done = False46        self._last_reward: float = 0.047        self._cumulative_reward: float = 0.048        self._last_action: CodeReviewAction | None = None49        self._last_error: str | None = None50 51    def reset(52        self,53        seed: Optional[int] = None,54        episode_id: Optional[str] = None,55        **kwargs: Any,56    ) -> CodeReviewObservation:57        """Reset the environment and return initial observation."""58        self._step_count = 059        self._done = False60        self._last_reward = 0.061        self._cumulative_reward = 0.062        self._last_action = None63        self._last_error = None64 65        return self._build_observation()66 67    def step(68        self,69        action: CodeReviewAction,70        timeout_s: Optional[float] = None,71        **kwargs: Any,72    ) -> CodeReviewObservation:73        """Take a step in the environment."""74        if self._done:75            raise RuntimeError("Episode done. Call reset() first.")76 77        self._step_count += 178        self._last_action = action79        self._last_error = None80 81        try:82            reward = grade(self._task_name, action)83        except Exception as exc:84            self._last_error = str(exc)85            reward_value = 0.086        else:87            reward_value = reward.value88 89        self._last_reward = reward_value90        self._cumulative_reward += reward_value91 92        # Episode ends after first submission (review is single-turn by nature)93        # or if max steps hit94        self._done = True95 96        obs = self._build_observation()97        obs.reward = reward_value98        obs.done = self._done99 100        return obs101 102    @property103    def state(self) -> CodeReviewState:104        """Get the current environment state."""105        return CodeReviewState(106            task_name=self._task_name,107            step_count=self._step_count,108            done=self._done,109            last_reward=self._last_reward,110            cumulative_reward=round(self._cumulative_reward, 4),111            last_action=self._last_action.model_dump() if self._last_action else None,112            last_error=self._last_error,113        )114 115    def _build_observation(self) -> CodeReviewObservation:116        """Build observation from current state."""117        return CodeReviewObservation(118            diff=self._seed.diff,119            file_path=self._seed.file_path,120            task_name=self._task_name,121            step_number=self._step_count,122            pr_context=self._seed.pr_context,123            done=self._done,124            reward=self._last_reward,125        )126 127    def close(self) -> None:128        """Clean up resources used by the environment."""129        # No resources to clean up for this environment130        pass131