""" Git Conflict Resolution — OpenEnv Environment Implements the full OpenEnv spec: step() / reset() / state() API with typed Pydantic models. """ from __future__ import annotations from typing import Any from pydantic import BaseModel, Field from tasks.scenarios import TASKS from graders.grader import grade # --------------------------------------------------------------------------- # Typed Pydantic models (OpenEnv spec requirement) # --------------------------------------------------------------------------- class Observation(BaseModel): task_id: str = Field(description="Current task identifier: easy / medium / hard") task_description: str = Field(description="Natural language description of what to resolve") conflicted_files: dict[str, str] = Field( description="Dict of filename -> file content with Git conflict markers" ) test_cases_hint: str = Field( description="Natural language hint about what the tests check" ) attempts_remaining: int = Field(description="How many attempts are left this episode") class Action(BaseModel): resolved_files: dict[str, str] = Field( description=( "Dict of filename -> resolved file content. " "Must NOT contain conflict markers (<<<<<<<, =======, >>>>>>>)." ) ) class RewardInfo(BaseModel): score: float = Field(description="Score between 0.0 and 1.0") feedback: str = Field(description="Detailed feedback on what passed or failed") tests_passed: int = Field(description="Number of test cases passed") conflict_markers_removed: bool = Field( description="True if all conflict markers were removed" ) # --------------------------------------------------------------------------- # Main environment # --------------------------------------------------------------------------- class GitConflictEnv: """ OpenEnv-compliant environment for Git merge conflict resolution. An agent is shown files with Git conflict markers and must resolve them correctly so all unit tests pass. Reward shaping: - 0.0 : conflict markers still present - 0.0 : syntax error in resolved file - 0.25 : partial — some tests pass - 0.5 : partial — half tests pass - 0.75 : most tests pass - 1.0 : all tests pass (perfect resolution) """ TASK_ORDER = ["easy", "medium", "hard"] MAX_ATTEMPTS = 3 def __init__(self, task_id: str = "easy"): if task_id not in TASKS: raise ValueError(f"task_id must be one of {list(TASKS.keys())}") self._task_id = task_id self._task = TASKS[task_id] self._attempts = 0 self._done = False self._last_score = 0.0 self._last_feedback = "" # ------------------------------------------------------------------ # OpenEnv core API # ------------------------------------------------------------------ def reset(self) -> Observation: """Reset the environment and return the initial observation.""" self._attempts = 0 self._done = False self._last_score = 0.0 self._last_feedback = "" return self._make_observation() def step(self, action: Action) -> tuple[Observation, float, bool, dict]: """ Take one step: agent submits resolved files. Returns (observation, reward, done, info). """ if self._done: raise RuntimeError("Episode is done. Call reset() to start a new episode.") self._attempts += 1 # Grade the submission score, feedback = grade(self._task_id, action.resolved_files) self._last_score = score self._last_feedback = feedback # Reward shaping: partial credit + bonus for finishing fast reward = score if score == 1.0 and self._attempts == 1: reward = 1.0 # perfect first try elif score == 1.0: reward = 0.9 # perfect but took multiple attempts else: reward = score # partial credit always given # Done if: perfect score OR out of attempts done = (score == 1.0) or (self._attempts >= self.MAX_ATTEMPTS) self._done = done info = RewardInfo( score=score, feedback=feedback, tests_passed=self._count_passed(score), conflict_markers_removed=self._no_markers(action.resolved_files), ).model_dump() obs = self._make_observation() return obs, round(reward, 4), done, info def state(self) -> dict[str, Any]: """Return the full current state of the environment.""" return { "task_id": self._task_id, "attempts": self._attempts, "max_attempts": self.MAX_ATTEMPTS, "done": self._done, "last_score": self._last_score, "last_feedback": self._last_feedback, "conflicted_files": self._task["conflicted_files"], } # ------------------------------------------------------------------ # Task listing helpers (required by OpenEnv spec) # ------------------------------------------------------------------ def list_tasks(self) -> list[dict]: """List all available tasks with metadata.""" return [ { "id": t["id"], "description": t["description"], "num_files": len(t["conflicted_files"]), "difficulty": t["id"], } for t in TASKS.values() ] def get_task_schema(self, task_id: str) -> dict: """Return the action schema for a given task.""" task = TASKS.get(task_id) if not task: raise ValueError(f"Unknown task: {task_id}") return { "action_fields": list(Action.model_fields.keys()), "files_to_resolve": list(task["conflicted_files"].keys()), } # ------------------------------------------------------------------ # Internal helpers # ------------------------------------------------------------------ def _make_observation(self) -> Observation: hints = { "easy": "Both discount conditions (order > 100 AND is_premium) must apply together.", "medium": "Validation must raise ValueError, loyalty_points = int(total * 10), logging must be present.", "hard": "User needs both role and session_token. All three files must be consistent.", } return Observation( task_id=self._task_id, task_description=self._task["description"], conflicted_files=self._task["conflicted_files"], test_cases_hint=hints.get(self._task_id, "Resolve all conflicts correctly."), attempts_remaining=self.MAX_ATTEMPTS - self._attempts, ) def _count_passed(self, score: float) -> int: total = len(TASKS[self._task_id]["test_cases"]) return round(score * total) def _no_markers(self, resolved_files: dict) -> bool: markers = ["<<<<<<<", "=======", ">>>>>>>"] for content in resolved_files.values(): if any(m in content for m in markers): return False return True # --------------------------------------------------------------------------- # Convenience factory used by baseline and HF Space # --------------------------------------------------------------------------- def make_env(task_id: str = "easy") -> GitConflictEnv: return GitConflictEnv(task_id=task_id)