models.py

"""
Pydantic models for the Quant Research Environment.

Defines the Action, Observation, and State types used across
the server, client, and inference script.
"""

from typing import List, Literal, Optional

from openenv.core.env_server.types import Action, Observation, State
from pydantic import Field

ActionType = Literal[
    "explore_data",
    "submit_code",
    "run_backtest",
    "submit_research",
    "submit_final",
]


class QuantAction(Action):
    """Agent action in the quant research environment."""

    action_type: ActionType = Field(
        default="explore_data",
        description=(
            "One of: explore_data, submit_code, run_backtest, "
            "submit_research, submit_final"
        ),
    )
    content: str = Field(
        default="",
        description="Python code, pandas query, or markdown text",
    )


class QuantObservation(Observation):
    """Environment response after each action."""

    task_id: str = Field(default="", description="easy | medium | hard")
    task_description: str = Field(default="", description="Full task spec")
    data_schema: str = Field(default="", description="Column descriptions")
    data_preview: str = Field(default="", description="First rows as string")
    function_signature: str = Field(default="", description="Expected function template")

    execution_output: str = Field(default="", description="stdout/stderr from code execution")
    backtest_results: str = Field(default="", description="JSON metrics from backtest")
    constraint_violations: List[str] = Field(default_factory=list, description="Violation messages")
    steps_remaining: int = Field(default=0, description="Steps left in episode")
    phase: str = Field(default="", description="explore | code | backtest | done")
    message: str = Field(default="", description="Human-readable status")


class QuantState(State):
    """Episode metadata."""

    task_id: str = Field(default="", description="Current task")
    current_phase: str = Field(default="", description="Current phase")
    has_submitted_code: bool = Field(default=False)
    backtest_count: int = Field(default=0)
    best_score: float = Field(default=0.0)
    steps_remaining: int = Field(default=0)