forgeenv-space/forgeenv/tasks/models.py

"""Core data models for ForgeEnv tasks and execution results.

These are framework-internal dataclasses (not Pydantic) used throughout the
simulation, verifier, and primitive layers. The OpenEnv-facing Pydantic
models live in `forgeenv.env.actions` / `forgeenv.env.observations`.
"""
from __future__ import annotations

from dataclasses import dataclass, field
from typing import Optional


@dataclass
class Task:
    """A HuggingFace training script with execution metadata."""

    task_id: str
    description: str
    script_content: str
    difficulty: str  # "easy", "medium", "hard"
    category: str = "general"
    expected_loss_range: tuple[float, float] = (0.0, 5.0)
    expected_accuracy_range: tuple[float, float] = (0.0, 1.0)
    checkpoint_output_path: str = "/tmp/forge_output/checkpoint"


@dataclass
class ExecutionResult:
    """Result of executing a Python script in the sandbox."""

    exit_code: int
    stdout: str
    stderr: str
    wall_time_ms: int
    checkpoint_exists: bool = False
    peak_memory_mb: float = 0.0
    script_content: str = ""


@dataclass
class ValidationResult:
    """Result of AST validation on a script."""

    is_valid: bool
    violations: list[str] = field(default_factory=list)