models.py

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


class DataDetectiveAction(Action):
    """Agent action: run a SQL query or submit a final answer."""

    action_type: str = Field(
        ...,
        description="'query' to execute SQL against the database, or 'answer' to submit findings",
    )
    content: str = Field(
        ...,
        description="SQL query string (for action_type='query') or final analysis text (for action_type='answer')",
    )


class DataDetectiveObservation(Observation):
    """Observation returned after each action."""

    output: str = Field(default="", description="Query results or system feedback")
    task_description: str = Field(default="", description="The investigation task to solve")
    schema_info: str = Field(default="", description="Database schema (provided at reset)")
    step_number: int = Field(default=0, description="Current step in the episode")
    max_steps: int = Field(default=30, description="Maximum steps allowed")
    message: str = Field(default="", description="Status or feedback message")


class DataDetectiveState(State):
    """Internal environment state."""

    task_id: str = Field(default="", description="Current task identifier")
    queries_executed: int = Field(default=0, description="Number of SQL queries run so far")
    max_steps: int = Field(default=30, description="Maximum steps allowed")