env/models.py

from pydantic import BaseModel, Field
from typing import Optional, List, Any

class Action(BaseModel):
    """What the agent can do each step."""
    sql_query: Optional[str] = Field(
        None,
        description="A SQL SELECT query to execute against the database"
    )
    submit_answer: Optional[str] = Field(
        None,
        description="Final answer to submit. Ends the episode."
    )

    def is_valid(self) -> bool:
        # Exactly one of the two must be set
        return bool(self.sql_query) != bool(self.submit_answer)


class QueryResult(BaseModel):
    """Result of executing a SQL query."""
    columns:    List[str]      = []
    rows:       List[List[Any]] = []
    error:      Optional[str]  = None
    truncated:  bool           = False
    total_rows: int            = 0


class Observation(BaseModel):
    """What the agent sees after each step."""
    schema_summary: str   = Field(..., description="Compact DB schema")
    question:       str   = Field(..., description="Business question to answer")
    last_query:     Optional[str]         = None
    last_result:    Optional[QueryResult] = None
    last_error:     Optional[str]         = None
    step:           int   = 0
    max_steps:      int   = 20
    hints:          List[str] = []
    done:           bool  = False


class StepResult(BaseModel):
    """Full result returned by step()."""
    observation: Observation
    reward:      float = 0.0
    done:        bool  = False
    info:        dict  = {}


class EnvState(BaseModel):
    """Full environment state returned by state()."""
    task_id:      str
    difficulty:   str
    step:         int
    max_steps:    int
    query_history: List[str] = []
    total_reward: float = 0.0
    done:         bool  = False