codefix-env 0.2.1__py3-none-any.whl

codefix_env/models.py

@@ -0,0 +1,234 @@
+"""
+CodeFix-Env Data Models
+=======================
+All Pydantic v2 models for Actions, Observations, State, and Tasks.
+Strongly typed — no `Any`, no raw dicts.
+"""
+
+from __future__ import annotations
+
+import uuid
+from datetime import datetime
+from enum import Enum
+
+from pydantic import BaseModel, Field, field_validator
+
+# ─────────────────────────────────────────────
+# Enums
+# ─────────────────────────────────────────────
+
+
+class ActionType(str, Enum):
+    RUN_TESTS = "run_tests"
+    EDIT_LINE = "edit_line"
+    INSERT_LINE = "insert_line"
+    DELETE_LINE = "delete_line"
+    GET_HINT = "get_hint"
+    SUBMIT_FIX = "submit_fix"
+    RESET = "reset"
+    VIEW_CODE = "view_code"
+
+
+class Difficulty(str, Enum):
+    EASY = "easy"
+    MEDIUM = "medium"
+    HARD = "hard"
+
+
+class BugCategory(str, Enum):
+    SYNTAX = "syntax_error"
+    LOGIC = "logic_bug"
+    OFF_BY_ONE = "off_by_one"
+    TYPE_ERROR = "type_error"
+    MISSING_IMPORT = "missing_import"
+    WRONG_OPERATOR = "wrong_operator"
+    RECURSION_BUG = "recursion_bug"
+    ALGORITHM_BUG = "algorithm_bug"
+    MULTI_FUNCTION = "multi_function"
+    RETURN_BUG = "return_bug"
+    INDEX_ERROR = "index_error"
+    SCOPE_BUG = "scope_bug"
+
+
+class TerminationReason(str, Enum):
+    SOLVED = "solved"
+    MAX_STEPS = "max_steps_reached"
+    SUBMITTED = "submitted"
+    TIMEOUT = "timeout"
+
+
+# ─────────────────────────────────────────────
+# Action
+# ─────────────────────────────────────────────
+
+
+class CodeFixAction(BaseModel):
+    """
+    An action taken by the agent in the CodeFix environment.
+
+    Examples::
+
+        CodeFixAction(action_type=ActionType.EDIT_LINE, line_number=5, new_content="    return x + 1")
+        CodeFixAction(action_type=ActionType.RUN_TESTS)
+        CodeFixAction(action_type=ActionType.SUBMIT_FIX)
+    """
+
+    action_type: ActionType = Field(..., description="Type of action to perform")
+    line_number: int | None = Field(
+        None, ge=1, description="Target line (1-indexed, for edit/insert/delete)"
+    )
+    new_content: str | None = Field(None, description="New line content (for edit/insert)")
+    reasoning: str | None = Field(
+        None, description="Agent's chain-of-thought (logged, not used by env)"
+    )
+    timestamp: datetime = Field(default_factory=datetime.utcnow)
+
+    @field_validator("line_number")
+    @classmethod
+    def line_required_for_edits(cls, v: int | None, info) -> int | None:
+        action = info.data.get("action_type")
+        if (
+            action in (ActionType.EDIT_LINE, ActionType.INSERT_LINE, ActionType.DELETE_LINE)
+            and v is None
+        ):
+            raise ValueError(f"line_number is required for action_type={action}")
+        return v
+
+    @field_validator("new_content")
+    @classmethod
+    def content_required_for_edits(cls, v: str | None, info) -> str | None:
+        action = info.data.get("action_type")
+        if action in (ActionType.EDIT_LINE, ActionType.INSERT_LINE) and v is None:
+            raise ValueError(f"new_content is required for action_type={action}")
+        return v
+
+    model_config = {"use_enum_values": True}
+
+
+# ─────────────────────────────────────────────
+# Observation
+# ─────────────────────────────────────────────
+
+
+class TestResult(BaseModel):
+    """Result of a single test case."""
+
+    name: str
+    passed: bool
+    output: str = ""
+    error: str = ""
+    runtime_ms: float = 0.0
+
+
+class CodeFixObservation(BaseModel):
+    """
+    Full observation returned after every env.step() or env.reset().
+    Contains everything the agent needs to make the next decision.
+    """
+
+    # Code state
+    current_code: str = Field(..., description="Current (possibly edited) source code")
+    original_code: str = Field(..., description="Original buggy code (never changes)")
+    diff: str = Field("", description="Unified diff: original → current")
+
+    # Test feedback
+    test_results: list[TestResult] = Field(default_factory=list)
+    test_output: str = Field("", description="Raw stdout/stderr from test runner")
+    tests_passed: int = 0
+    tests_total: int = 0
+    all_tests_pass: bool = False
+
+    # Scoring
+    score_so_far: float = Field(0.0, ge=0.0, le=1.0, description="Cumulative normalised score")
+    shaped_reward: float = Field(0.0, description="Dense reward for this step")
+
+    # Episode metadata
+    step_count: int = 0
+    max_steps: int = 20
+    steps_remaining: int = 20
+    done: bool = False
+    termination_reason: TerminationReason | None = None
+
+    # Task info
+    task_id: str = ""
+    difficulty: Difficulty = Difficulty.MEDIUM
+    bug_category: BugCategory = BugCategory.LOGIC
+    hint_available: bool = True
+    hints_used: int = 0
+
+    # Agent feedback
+    feedback: str = ""
+    error_message: str = ""
+
+
+# ─────────────────────────────────────────────
+# Episode State
+# ─────────────────────────────────────────────
+
+
+class CodeFixState(BaseModel):
+    """Internal episode state (returned by env.state())."""
+
+    episode_id: str = Field(default_factory=lambda: str(uuid.uuid4()))
+    task_id: str = ""
+    step_count: int = 0
+    max_steps: int = 20
+    total_reward: float = 0.0
+    action_history: list[dict] = Field(default_factory=list)
+    start_time: datetime = Field(default_factory=datetime.utcnow)
+    done: bool = False
+    solved: bool = False
+    hints_used: int = 0
+
+
+# ─────────────────────────────────────────────
+# Step Result (wraps observation)
+# ─────────────────────────────────────────────
+
+
+class StepResult(BaseModel):
+    """Full result returned by env.step()."""
+
+    observation: CodeFixObservation
+    reward: float
+    done: bool
+    truncated: bool = False
+    info: dict = Field(default_factory=dict)
+
+
+# ─────────────────────────────────────────────
+# Task Definition
+# ─────────────────────────────────────────────
+
+
+class TestCase(BaseModel):
+    """A single test case for a task."""
+
+    name: str
+    code: str = Field(..., description="Test code (function call + assertion)")
+    expected: str = ""
+    timeout_s: float = 5.0
+
+
+class Task(BaseModel):
+    """A complete debugging task."""
+
+    id: str
+    title: str
+    description: str
+    buggy_code: str
+    solution_code: str
+    test_cases: list[TestCase]
+    difficulty: Difficulty
+    bug_category: BugCategory
+    tags: list[str] = Field(default_factory=list)
+    hints: list[str] = Field(default_factory=list)
+    max_steps: int = 20
+    hint_penalty: float = 0.05
+    time_penalty_per_step: float = 0.01
+    author: str = "codefix-env"
+    created_at: datetime = Field(default_factory=datetime.utcnow)
+
+    @property
+    def num_tests(self) -> int:
+        return len(self.test_cases)

codefix_env/rewards.py

@@ -0,0 +1,157 @@
+"""
+Reward Shaping Pipeline
+========================
+High-level reward computation that combines:
+1. Rule-based rewards (always active)
+2. Neural reward model (optional, loaded from checkpoint)
+3. Reward normalisation across an episode
+
+Used by the environment server on every step.
+
+NOTE: torch and RewardMLP are imported lazily inside __init__, only when
+a neural_model_path is actually provided. This keeps importing
+codefix_env fast (no torch load) for the common case of rule-based-only
+reward, which matters a lot for sandboxed worker processes that re-import
+the whole package on every spawn (see utils/sandbox.py).
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+from pathlib import Path
+
+from codefix_env.models import ActionType, CodeFixObservation, Task
+from codefix_env.utils.metrics import (
+    EpisodeMetrics,
+    ScoringConfig,
+    compute_diff_score,
+    compute_final_score,
+    compute_shaped_reward,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class RewardPipeline:
+    """
+    Combines rule-based reward with an optional neural reward model.
+
+    Usage::
+
+        pipeline = RewardPipeline(cfg=ScoringConfig())
+        reward = pipeline.step_reward(prev_obs, curr_obs, action_type, task)
+        final  = pipeline.episode_reward(curr_obs, task)
+    """
+
+    def __init__(
+        self,
+        cfg: ScoringConfig | None = None,
+        neural_model_path: str | None = None,
+        neural_weight: float = 0.3,
+    ):
+        self.cfg = cfg if cfg is not None else ScoringConfig()
+        self.neural_weight = neural_weight
+        self._neural_model = None
+
+        # Load pre-trained neural reward model if provided.
+        # torch + RewardMLP are imported here, not at module top, so the
+        # common no-neural-model path never pays torch's import cost.
+        if neural_model_path and Path(neural_model_path).exists():
+            try:
+                import torch
+
+                from codefix_env.utils.reward_model import RewardMLP
+
+                self._neural_model = RewardMLP()
+                state = torch.load(neural_model_path, map_location="cpu", weights_only=True)
+                self._neural_model.load_state_dict(state)
+                self._neural_model.eval()
+                logger.info("Neural reward model loaded from %s", neural_model_path)
+            except Exception as e:  # noqa: BLE001
+                logger.warning("Failed to load neural reward model: %s", e)
+                self._neural_model = None
+
+    @property
+    def uses_neural(self) -> bool:
+        return self._neural_model is not None
+
+    def step_reward(
+        self,
+        prev_obs: CodeFixObservation,
+        curr_obs: CodeFixObservation,
+        action_type: ActionType,
+        task: Task,
+    ) -> float:
+        """
+        Compute dense reward for a single step.
+        Blend of rule-based + optional neural signal.
+        """
+        rule_reward = compute_shaped_reward(
+            prev_passed=prev_obs.tests_passed,
+            curr_passed=curr_obs.tests_passed,
+            tests_total=curr_obs.tests_total,
+            step_count=curr_obs.step_count,
+            hints_used=curr_obs.hints_used,
+            action_type=action_type.value if hasattr(action_type, "value") else str(action_type),
+            cfg=self.cfg,
+        )
+
+        if self._neural_model is None:
+            return rule_reward
+
+        # Neural component
+        diff_score = compute_diff_score(curr_obs.original_code, curr_obs.current_code)
+        action_id = list(ActionType).index(action_type) if action_type in list(ActionType) else 0
+        neural_reward = self._neural_model.predict(
+            tests_passed_ratio=curr_obs.tests_passed / max(curr_obs.tests_total, 1),
+            step_ratio=curr_obs.step_count / max(task.max_steps, 1),
+            hints_used=curr_obs.hints_used,
+            diff_score=diff_score,
+            prev_score=prev_obs.score_so_far,
+            action_type_id=action_id,
+            code_len_ratio=len(curr_obs.current_code) / max(len(curr_obs.original_code), 1),
+            test_count=curr_obs.tests_total,
+        )
+
+        # Weighted blend: mostly rules, neural provides nuance
+        blended = (1.0 - self.neural_weight) * rule_reward + self.neural_weight * neural_reward
+        return float(blended)
+
+    def episode_reward(
+        self,
+        final_obs: CodeFixObservation,
+        task: Task,
+    ) -> float:
+        """Final episode reward on termination."""
+        return compute_final_score(
+            tests_passed=final_obs.tests_passed,
+            tests_total=final_obs.tests_total,
+            step_count=final_obs.step_count,
+            hints_used=final_obs.hints_used,
+            cfg=self.cfg,
+        )
+
+    def build_metrics(
+        self,
+        task: Task,
+        final_obs: CodeFixObservation,
+        total_reward: float,
+    ) -> EpisodeMetrics:
+        return EpisodeMetrics(
+            task_id=task.id,
+            solved=final_obs.all_tests_pass,
+            final_score=final_obs.score_so_far,
+            total_steps=final_obs.step_count,
+            hints_used=final_obs.hints_used,
+            tests_passed=final_obs.tests_passed,
+            tests_total=final_obs.tests_total,
+            total_reward=total_reward,
+        )
+
+
+# Singleton default pipeline (no neural model unless env var is set)
+default_pipeline = RewardPipeline(
+    cfg=ScoringConfig(),
+    neural_model_path=os.environ.get("CODEFIX_REWARD_MODEL_PATH"),
+)

codefix_env/tasks/__init__.py

@@ -0,0 +1,75 @@
+"""
+Task Registry — Central access point for all CodeFix tasks.
+"""
+
+from __future__ import annotations
+
+import random
+
+from codefix_env.models import Difficulty, Task
+from codefix_env.tasks.easy import EASY_TASKS
+from codefix_env.tasks.hard import HARD_TASKS
+from codefix_env.tasks.medium import MEDIUM_TASKS
+
+# ── Master registry ──────────────────────────────────────────────────────────
+ALL_TASKS: dict[str, Task] = {task.id: task for task in EASY_TASKS + MEDIUM_TASKS + HARD_TASKS}
+
+TASKS_BY_DIFFICULTY: dict[Difficulty, list[Task]] = {
+    Difficulty.EASY: EASY_TASKS,
+    Difficulty.MEDIUM: MEDIUM_TASKS,
+    Difficulty.HARD: HARD_TASKS,
+}
+
+
+def load_task(task_id: str) -> Task:
+    """Load a task by its exact ID."""
+    if task_id not in ALL_TASKS:
+        raise KeyError(f"Task '{task_id}' not found. Use list_tasks() to see available tasks.")
+    return ALL_TASKS[task_id]
+
+
+def random_task(
+    difficulty: Difficulty | None = None,
+    exclude: list[str] | None = None,
+) -> Task:
+    """
+    Return a random task, optionally filtered by difficulty and
+    excluding already-seen task IDs.
+    """
+    pool = list(ALL_TASKS.values())
+    if difficulty:
+        pool = [t for t in pool if t.difficulty == difficulty]
+    if exclude:
+        pool = [t for t in pool if t.id not in exclude]
+    if not pool:
+        raise ValueError("No tasks available after applying filters.")
+    return random.choice(pool)
+
+
+def list_tasks(difficulty: Difficulty | None = None) -> list[Task]:
+    """List all tasks, optionally filtered by difficulty."""
+    if difficulty:
+        return TASKS_BY_DIFFICULTY.get(difficulty, [])
+    return list(ALL_TASKS.values())
+
+
+def task_count() -> dict[str, int]:
+    """Return count of tasks by difficulty."""
+    return {
+        "easy": len(EASY_TASKS),
+        "medium": len(MEDIUM_TASKS),
+        "hard": len(HARD_TASKS),
+        "total": len(ALL_TASKS),
+    }
+
+
+__all__ = [
+    "ALL_TASKS",
+    "EASY_TASKS",
+    "MEDIUM_TASKS",
+    "HARD_TASKS",
+    "load_task",
+    "random_task",
+    "list_tasks",
+    "task_count",
+]