albert-einstein-09/codedark

1"""2CodeDark Data Models3 4Pydantic models for Action, Observation, and State following OpenEnv spec.5"""6 7from pydantic import BaseModel, Field8from typing import Optional, List, Any, Literal9 10 11class CodeDarkAction(BaseModel):12    """13    Action for CodeDark environment.14 15    Agents send actions with a tool name and arguments.16 17    Tools available:18    - run_python: Execute Python code with pandas/numpy19    - read_notes: Read all saved notes20    - save_note: Save a note for later recall21    - clarify: Ask clarifying question (max 2 per episode)22    - submit_answer: Submit final answer (ends episode)23    """24 25    tool: Literal["run_python", "read_notes", "save_note", "clarify", "submit_answer"]26    args: str = ""  # Tool-specific arguments27 28    model_config = {29        "json_schema_extra": {30            "examples": [31                {"tool": "run_python", "args": "result = df['y'].mean() * 100"},32                {"tool": "read_notes", "args": ""},33                {"tool": "save_note", "args": "Average subscription rate is 11.26%"},34                {"tool": "clarify", "args": "What does Q1 mean in this context?"},35                {"tool": "submit_answer", "args": "11.26"},36            ]37        }38    }39 40 41class CodeDarkObservation(BaseModel):42    """43    Observation returned after each action.44 45    Contains execution results, environment state, and episode info.46    Reward is only populated when done=True.47    """48 49    # Execution results50    stdout: str = ""51    stderr: str = ""52    exit_code: int = 053 54    # Turn tracking55    turn: int = 056    max_turns: int = 1057 58    # Persistent state59    notes: List[str] = Field(default_factory=list)60 61    # Task info62    task_id: str = ""63    question: str = ""64    difficulty: str = ""  # L4, L5, L665    dataset: str = ""  # bank, road66 67    # Episode status68    done: bool = False69    submitted: bool = False70 71    # Reward components (only set when done=True)72    reward: Optional[float] = None73    correctness: Optional[float] = None74    efficiency: Optional[float] = None75 76    # Additional metadata77    metadata: dict = Field(default_factory=dict)78 79    model_config = {80        "json_schema_extra": {81            "examples": [82                {83                    "stdout": "run_python Result:\n(45211, 17)",84                    "stderr": "",85                    "exit_code": 0,86                    "turn": 1,87                    "max_turns": 10,88                    "notes": [],89                    "task_id": "bank_hard_001",90                    "question": "What's the subscription rate for month='may'?",91                    "difficulty": "L5",92                    "dataset": "bank",93                    "done": False,94                    "submitted": False,95                    "reward": None,96                }97            ]98        }99    }100 101 102class CodeDarkState(BaseModel):103    """104    Internal state for CodeDark environment.105 106    Tracks episode progress, accumulated notes, and submission status.107    """108 109    episode_id: str = ""110    step_count: int = 0111 112    # Task info113    task_id: str = ""114    dataset: str = ""115 116    # Accumulated state117    notes: List[str] = Field(default_factory=list)118    turn_count: int = 0119    error_count: int = 0120    clarify_count: int = 0121 122    # Submission123    submitted: bool = False124    submitted_answer: Optional[Any] = None125 126    # For scoring127    expected_answer: Optional[Any] = None128    tolerance: float = 0.01129 130 131class ResetRequest(BaseModel):132    """Request body for /reset endpoint."""133 134    task_id: Optional[str] = None135    seed: Optional[int] = None136 137 138class StepRequest(BaseModel):139    """Request body for /step endpoint."""140 141    tool: str142    args: str = ""143 144 145class HealthResponse(BaseModel):146    """Response for /health endpoint."""147 148    status: str = "healthy"149    environment: str = "codedark"150    version: str = "0.1.0"151