Spaces:
Sleeping
Sleeping
| import json | |
| from uuid import uuid4 | |
| from openenv.core.env_server.interfaces import Environment | |
| from openenv.core.env_server.types import State | |
| from typing import List, Dict, Optional | |
| try: | |
| from ..models import GitCommitReviewerAction, GitCommitReviewerObservation | |
| except ImportError: | |
| from models import GitCommitReviewerAction, GitCommitReviewerObservation | |
| import tasks | |
| # Global cache to handle stateless server environments | |
| # Shared across instances | |
| GLOBAL_SESSION_STORE = {} | |
| class GitCommitReviewerEnvironment(Environment): | |
| SUPPORTS_CONCURRENT_SESSIONS: bool = True | |
| def __init__(self): | |
| self._state = State(episode_id=str(uuid4()), step_count=0) | |
| self._reset_count = 0 | |
| self.current_task_commits = [] | |
| self.current_idx = 0 | |
| self.total_rewards = 0.0 | |
| self.last_task = "easy" | |
| def reset(self, task: str = "easy") -> GitCommitReviewerObservation: | |
| self.last_task = str(task).lower() | |
| self._state = State(episode_id=self._state.episode_id, step_count=0) # Keep same EP ID if possible or new one | |
| self._state.episode_id = str(uuid4()) # Generate new unique session | |
| self._reset_count += 1 | |
| self.current_task_commits = tasks.get_task_data(self.last_task) | |
| if not self.current_task_commits: | |
| self.current_task_commits = tasks.get_task_data("easy") | |
| self.current_idx = 0 | |
| self.total_rewards = 0.0 | |
| # Store initial state in global cache | |
| GLOBAL_SESSION_STORE[self._state.episode_id] = { | |
| "current_idx": self.current_idx, | |
| "total_rewards": self.total_rewards, | |
| "commits": self.current_task_commits, | |
| "task": self.last_task, | |
| "step_count": self._state.step_count | |
| } | |
| return GitCommitReviewerObservation( | |
| commit_message=self.current_task_commits[0]["commit_message"], | |
| commit_idx=0, | |
| total_commits=len(self.current_task_commits), | |
| done=False, | |
| reward=0.0, | |
| session_id=self._state.episode_id, | |
| metadata={"session_id": self._state.episode_id} | |
| ) | |
| def step(self, action: GitCommitReviewerAction) -> GitCommitReviewerObservation: | |
| # Try to restore from global cache if session matches | |
| sid = getattr(action, "session_id", None) | |
| if sid and sid in GLOBAL_SESSION_STORE: | |
| cached = GLOBAL_SESSION_STORE[sid] | |
| # Only restore if current instance is uninitialized or mismatched | |
| if not self.current_task_commits or self._state.episode_id != sid: | |
| self.current_task_commits = cached["commits"] | |
| self.current_idx = cached["current_idx"] | |
| self.total_rewards = cached["total_rewards"] | |
| self._state.episode_id = sid | |
| self._state.step_count = cached["step_count"] | |
| self.last_task = cached["task"] | |
| self._state.step_count += 1 | |
| if not self.current_task_commits or self.current_idx >= len(self.current_task_commits): | |
| return GitCommitReviewerObservation( | |
| commit_message="DONE", | |
| commit_idx=self.current_idx, | |
| total_commits=len(self.current_task_commits) if self.current_task_commits else 0, | |
| done=True, | |
| reward=0.0, | |
| session_id=self._state.episode_id, | |
| metadata={"session_id": self._state.episode_id} | |
| ) | |
| gt = self.current_task_commits[self.current_idx]["ground_truth"] | |
| # Grading Logic | |
| reward = 0.0 | |
| # 1. Score matching | |
| score_diff = abs(action.score - gt["score"]) | |
| if score_diff <= 0.05: reward += 0.45 | |
| elif score_diff <= 0.15: reward += 0.3 | |
| elif score_diff <= 0.3: reward += 0.15 | |
| # 2. Type Valid matching | |
| if action.type_valid == gt["type_valid"]: reward += 0.2 | |
| # 3. Summary length proximity | |
| length_diff = abs(action.summary_length - gt["summary_length"]) | |
| if length_diff == 0: reward += 0.15 | |
| elif length_diff <= 2: reward += 0.05 | |
| # 4. Suggestion/Issues quality | |
| if gt["issues"] and len(action.issues) > 0: reward += 0.1 | |
| elif not gt["issues"] and len(action.issues) == 0: reward += 0.1 | |
| if action.suggestion and ":" in action.suggestion and len(action.suggestion.split(":")[0]) <= 20: | |
| reward += 0.1 | |
| reward = max(0.01, min(0.99, reward)) # Clamp for OpenEnv safety | |
| self.total_rewards += reward | |
| self.current_idx += 1 | |
| done = (self.current_idx >= len(self.current_task_commits)) | |
| # Update global cache for this session | |
| GLOBAL_SESSION_STORE[self._state.episode_id] = { | |
| "current_idx": self.current_idx, | |
| "total_rewards": self.total_rewards, | |
| "commits": self.current_task_commits, | |
| "task": self.last_task, | |
| "step_count": self._state.step_count | |
| } | |
| return GitCommitReviewerObservation( | |
| commit_message=self.current_task_commits[self.current_idx]["commit_message"] if not done else "DONE", | |
| commit_idx=self.current_idx, | |
| total_commits=len(self.current_task_commits), | |
| done=done, | |
| reward=reward, | |
| session_id=self._state.episode_id, | |
| metadata={ | |
| "step": self._state.step_count, | |
| "total_rewards": self.total_rewards, | |
| "session_id": self._state.episode_id, | |
| "reward": reward | |
| } | |
| ) | |
| def state(self) -> dict: | |
| # Crucial: Check global cache to ensure we return the latest state | |
| # even if this is a newly created instance for a stateless request. | |
| eid = self._state.episode_id | |
| current_rewards = self.total_rewards | |
| current_steps = self._state.step_count | |
| if eid in GLOBAL_SESSION_STORE: | |
| cached = GLOBAL_SESSION_STORE[eid] | |
| current_rewards = max(current_rewards, cached.get("total_rewards", 0)) | |
| current_steps = max(current_steps, cached.get("step_count", 0)) | |
| return { | |
| "episode_id": eid, | |
| "step_count": current_steps, | |
| "total_rewards": current_rewards, | |
| "current_idx": self.current_idx, | |
| "total_commits": len(self.current_task_commits), | |
| "is_done": self.current_idx >= len(self.current_task_commits) if self.current_task_commits else True, | |
| "task": self.last_task | |
| } | |