git-conflict-resolver / environment.py
Bharath Vikas Tadepalli
First commit
6f0f8d0
Raw
History Blame Contribute Delete
7.56 kB
"""
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)