NILESH1003/codeguardian-env

1import random2import uuid3from typing import Any, Dict, Optional4 5from models import CodeAction, CodeObservation, EnvState, IssueFound6from server.grader import (7    LOW_REWARD,8    MIN_STRICT_VALUE,9    clamp_strict,10    compute_step_reward,11    evaluate_score,12)13from server.tasks import TASKS14 15 16def clamp_score(value: float) -> float:17    """Clamp any float to the strict open interval (0, 1)."""18    return clamp_strict(value)19 20 21class CodeGuardianEnvironment:22    def __init__(self):23        self.current_task = None24        self.episode_id = ""25        self.step_count = 026        self.total_reward = MIN_STRICT_VALUE27        self.actions_history = []28        self.issues_found = []29        self.done = False30 31    def reset(self, task_id: Optional[str] = None) -> CodeObservation:32        if task_id:33            self.current_task = next((t for t in TASKS if t["task_id"] == task_id), None)34        else:35            self.current_task = random.choice(TASKS)36 37        if not self.current_task:38            self.current_task = TASKS[0]39 40        self.episode_id = str(uuid.uuid4())41        self.step_count = 042        self.total_reward = MIN_STRICT_VALUE43        self.actions_history = []44        self.issues_found = []45        self.done = False46 47        return self.get_observation()48 49    def step(self, action: CodeAction) -> Dict[str, Any]:50        if self.done:51            return {52                "observation": self.get_observation(),53                "reward": clamp_score(LOW_REWARD),54                "done": True,55                "info": {"error": "Environment is already done."},56            }57 58        reward_detail = compute_step_reward(action, self.current_task, self.actions_history)59 60        self.actions_history.append(action)61        if action.action in ["flag_bug", "suggest_fix"] and action.line is not None:62            self.issues_found.append(63                IssueFound(64                    line=action.line,65                    bug_type=action.bug_type or "unknown",66                    comment=action.comment,67                )68            )69 70        self.step_count += 171        step_reward = clamp_score(reward_detail.step_reward)72        self.total_reward = clamp_score(self.total_reward + step_reward)73 74        if action.action in ["approve", "reject"] or self.step_count >= self.current_task["max_steps"]:75            self.done = True76 77        info = {78            "step_reward": step_reward,79            "reason": reward_detail.reason,80            "partial": reward_detail.partial,81            "episode_score": None,82            "score": None,83        }84 85        if self.done:86            final_score = clamp_score(evaluate_score(self.current_task, self.actions_history))87            info["episode_score"] = final_score88            info["score"] = final_score89 90        return {91            "observation": self.get_observation(),92            "reward": step_reward,93            "done": self.done,94            "info": info,95        }96 97    def get_observation(self) -> CodeObservation:98        return CodeObservation(99            code=self.current_task["code"] if self.current_task else "",100            filename=self.current_task["filename"] if self.current_task else "",101            language="python",102            task_id=self.current_task["task_id"] if self.current_task else "",103            task_difficulty=self.current_task["difficulty"] if self.current_task else "easy",104            issues_found=self.issues_found,105            step_count=self.step_count,106            done=self.done,107        )108 109    def state(self) -> EnvState:110        if not self.current_task:111            return EnvState(112                episode_id="",113                task_id="",114                step_count=0,115                total_reward=MIN_STRICT_VALUE,116                done=False,117                task_difficulty="easy",118            )119        return EnvState(120            episode_id=self.episode_id,121            task_id=self.current_task["task_id"],122            step_count=self.step_count,123            total_reward=clamp_score(self.total_reward),124            done=self.done,125            task_difficulty=self.current_task["difficulty"],126        )127