models.py

"""
Typed Action and Observation models for the SQL Query Environment.

Action: The agent submits a SQL query string and a task_id.
Observation: The environment returns schema info, query results, feedback, and reward.
"""

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


class SQLAction(Action):
    """Action submitted by the agent: a SQL query to execute."""

    task_id: str = Field(
        ...,
        description="ID of the task being attempted (task_1, task_2, task_3)",
    )
    sql_query: str = Field(
        ...,
        description="The SQL query string to execute against the database",
    )


class SQLObservation(Observation):
    """Observation returned to the agent after each step."""

    # Task information
    task_id: str = Field(default="", description="Current task ID")
    task_description: str = Field(
        default="", description="Natural language question the agent must answer"
    )
    difficulty: str = Field(default="", description="easy, medium, or hard")

    # Database schema
    schema_description: str = Field(
        default="", description="SQL CREATE TABLE statements describing the database"
    )

    # Query result / feedback
    query_result: str = Field(
        default="",
        description="Result of the executed SQL query (rows as text), or error message",
    )
    query_error: bool = Field(
        default=False, description="True if the SQL query caused an error"
    )
    feedback: str = Field(
        default="",
        description="Human-readable feedback on the query result",
    )

    # Scoring
    reward: float = Field(default=0.0, description="Score from 0.0 to 1.0")
    done: bool = Field(default=False, description="True if the episode is complete")

    # Metadata
    step_count: int = Field(default=0, description="Number of steps taken so far")
    max_steps: int = Field(
        default=3, description="Maximum steps allowed per task"
    )