YashR05/pullrequest-arena

1import json2import random3import uuid4import os5from openenv.core.env_server.interfaces import Environment6 7try:8    from ..models import ReviewAction, PRObservation9    from .graders import route_grader, _clip10except ImportError:11    from models import ReviewAction, PRObservation12    from server.graders import route_grader, _clip13 14class PullRequestEnvironment(Environment):15    def __init__(self):16        self.current_task = None17        self.episode_id = None18        self.step_count = 019        20        # Determine path to tasks.json (typically in parent directory from server/)21        base_dir = os.path.dirname(__file__)22        tasks_path = os.path.join(base_dir, "..", "tasks.json")23        if not os.path.exists(tasks_path):24            tasks_path = os.path.join(base_dir, "tasks.json")25            if not os.path.exists(tasks_path):26                # Fallback to CWD27                tasks_path = "tasks.json"28                29        with open(tasks_path, "r") as f:30            self.tasks = json.load(f)31 32    def reset(self, task_id=None):33        if task_id is not None:34            # Find the specific task by id string35            target = str(task_id)36            self.current_task = next((t for t in self.tasks if str(t["id"]) == target), None)37            if not self.current_task:38                self.current_task = random.choice(self.tasks)39        else:40            self.current_task = random.choice(self.tasks)41            42        self.episode_id = str(uuid.uuid4())43        self.step_count = 044        45        return PRObservation(46            pr_title=self.current_task.get("pr_title", ""),47            pr_description=self.current_task.get("pr_description", ""),48            files_changed=self.current_task.get("files_changed", []),49            code_diff=self.current_task.get("code_diff", ""),50            language=self.current_task.get("language", "python"),51            tests_passed=self.current_task.get("tests_passed", False),52            ci_logs=self.current_task.get("ci_logs", ""),53            repository_context=self.current_task.get("repository_context", ""),54            test_results=self.current_task.get("test_results", None),55            repo_tree=self.current_task.get("repo_tree", None),56            previous_comments=self.current_task.get("previous_comments", None),57            review_status=self.current_task.get("review_status", "pending"),58            task_id=str(self.current_task["id"]),59            difficulty=self.current_task.get("difficulty", "easy"),60            feedback="",61            done=False,62            reward=0.0163        )64 65    def step(self, action: ReviewAction):66        reward = _clip(route_grader(action, self.current_task))67        self.step_count += 168        69        feedback = f"Review action '{action.type}' submitted. Evaluated reward: {reward:.2f}."70        71        return PRObservation(72            pr_title=self.current_task.get("pr_title", ""),73            pr_description=self.current_task.get("pr_description", ""),74            files_changed=self.current_task.get("files_changed", []),75            code_diff=self.current_task.get("code_diff", ""),76            language=self.current_task.get("language", "python"),77            tests_passed=self.current_task.get("tests_passed", False),78            ci_logs=self.current_task.get("ci_logs", ""),79            repository_context=self.current_task.get("repository_context", ""),80            test_results=self.current_task.get("test_results", None),81            repo_tree=self.current_task.get("repo_tree", None),82            previous_comments=self.current_task.get("previous_comments", None),83            review_status=self.current_task.get("review_status", "pending"),84            task_id=str(self.current_task["id"]),85            difficulty=self.current_task.get("difficulty", "easy"),86            feedback=feedback,87            done=True,88            reward=reward89        )90 91    @property92    def state(self):93        return {94            "episode_id": self.episode_id,95            "step_count": self.step_count,96            "current_task_id": str(self.current_task["id"]) if self.current_task else None,97            "review_status": "done" if self.step_count > 0 else "pending"98        }99