| from typing import Dict, Any, Tuple, Optional |
| from environment.models import ( |
| ReviewAction, |
| ReviewState, |
| Observation |
| ) |
| from environment.tasks import TaskDefinitions |
| from environment.graders import TaskGrader, RewardCalculator |
|
|
|
|
| class CodeReviewEnv: |
| |
| def __init__(self): |
| self._state: Optional[ReviewState] = None |
| self.grader: Optional[TaskGrader] = None |
| self.reward_calculator = RewardCalculator() |
| self.max_steps = 50 |
| self.current_task_id: Optional[str] = None |
| |
| def reset(self, task_id: str = None) -> Dict[str, Any]: |
| if task_id is None: |
| task_id = "bug_detection_easy_1" |
| |
| self.current_task_id = task_id |
| task_data = TaskDefinitions.get_task(task_id) |
| |
| code_context = TaskDefinitions.create_code_context(task_data) |
| task_metadata = TaskDefinitions.create_task_metadata(task_data) |
| |
| self._state = ReviewState( |
| code_context=code_context, |
| task_metadata=task_metadata, |
| comments_made=[], |
| suggestions_made=[], |
| current_step=0, |
| is_complete=False, |
| final_decision=None, |
| last_action_valid=True, |
| last_error=None |
| ) |
| |
| self.grader = TaskGrader(task_metadata.expected_issues) |
| self.reward_calculator.reset() |
| |
| return self._get_observation() |
| |
| def step(self, action: Dict[str, Any]) -> Tuple[Dict[str, Any], float, bool, Dict[str, Any]]: |
| if self._state is None: |
| return {}, -0.1, True, {"error": "Environment not initialized. Call reset() first."} |
|
|
| if self._state.is_complete: |
| return self._get_observation(), 0.0, True, {"error": "Episode already complete"} |
| |
| try: |
| review_action = ReviewAction(**action) |
| except Exception as e: |
| self._state.last_action_valid = False |
| self._state.last_error = str(e) |
| return self._get_observation(), -0.1, False, {"error": str(e), "last_action_valid": False} |
| |
| self._state.current_step += 1 |
| self._process_action(review_action) |
|
|
| if review_action.action_type.value == "approve" and not review_action.final_decision: |
| review_action.final_decision = "approved" |
| elif review_action.action_type.value == "request_changes" and not review_action.final_decision: |
| review_action.final_decision = "changes_requested" |
| |
| if self._state.current_step >= self.max_steps: |
| self._state.is_complete = True |
| if not self._state.final_decision: |
| self._state.final_decision = "changes_requested" |
| |
| if review_action.final_decision and not self._state.is_complete: |
| self._state.is_complete = True |
| self._state.final_decision = review_action.final_decision |
| |
| reward = self.reward_calculator.calculate_reward( |
| review_action, |
| self._state.comments_made, |
| self._state.suggestions_made, |
| self._state.final_decision or "changes_requested", |
| self.grader, |
| self._state.last_action_valid, |
| ) |
| |
| info = { |
| "step": self._state.current_step, |
| "last_action_valid": self._state.last_action_valid, |
| "error": self._state.last_error, |
| "task_score": self.get_task_score(), |
| } |
| |
| return self._get_observation(), reward, self._state.is_complete, info |
| |
| def _process_action(self, action: ReviewAction): |
| if self._state is None: |
| return |
|
|
| self._state.last_action_valid = True |
| self._state.last_error = None |
| |
| if action.action_type.value == "add_comment": |
| for comment in action.comments: |
| if comment.line_number <= self._state.code_context.line_count: |
| self._state.comments_made.append(comment) |
| else: |
| self._state.last_action_valid = False |
| self._state.last_error = f"Line {comment.line_number} out of range" |
| |
| elif action.action_type.value == "suggest_fix": |
| for suggestion in action.suggestions: |
| if suggestion.original_line <= self._state.code_context.line_count: |
| self._state.suggestions_made.append(suggestion) |
| else: |
| self._state.last_action_valid = False |
| self._state.last_error = f"Line {suggestion.original_line} out of range" |
| |
| elif action.action_type.value == "mark_as_resolved": |
| for comment in action.comments: |
| for existing_comment in self._state.comments_made: |
| if existing_comment.line_number == comment.line_number: |
| existing_comment.resolved = True |
| |
| def _get_observation(self) -> Dict[str, Any]: |
| if self._state is None: |
| return {} |
|
|
| return Observation( |
| code_diff=self._state.code_context.code_diff, |
| file_context=self._state.code_context.surrounding_code, |
| file_path=self._state.code_context.file_path, |
| language=self._state.code_context.language, |
| task_description=self._state.task_metadata.description, |
| task_difficulty=self._state.task_metadata.difficulty, |
| current_step=self._state.current_step, |
| max_steps=self.max_steps, |
| previous_comments=self._state.comments_made, |
| previous_suggestions=self._state.suggestions_made, |
| review_complete=self._state.is_complete, |
| final_decision_made=self._state.final_decision |
| ).model_dump() |
| |
| def get_task_score(self) -> float: |
| if not self.grader or self._state is None: |
| return 0.0 |
|
|
| return self.grader.compute_score_from_state( |
| comments=self._state.comments_made, |
| suggestions=self._state.suggestions_made, |
| final_decision=self._state.final_decision or "changes_requested", |
| ) |
| |
| def close(self): |
| pass |
| |
| def state(self) -> Dict[str, Any]: |
| if self._state: |
| return self._state.model_dump() |
| return {} |
