models.py

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


class SQLAction(Action):
    action_type: str = Field(
        ...,
        description="Either 'submit_fix' to submit a corrected SQL query, or 'request_hint' to get a hint.",
    )
    sql_query: Optional[str] = Field(
        default=None,
        description="The corrected SQL query. Required when action_type is 'submit_fix'.",
    )


class SQLObservation(Observation):
    # The broken query the agent needs to fix
    broken_query: str = Field(..., description="The SQL query that contains a bug.")
    schema_description: str = Field(..., description="Description of the database schema.")
    task_description: str = Field(..., description="What the correct query should return.")

    # Feedback from the last action
    execution_result: str = Field(
        default="", description="Output or error from executing the submitted query."
    )
    is_correct: bool = Field(
        default=False, description="Whether the last submitted query was correct."
    )
    hint: Optional[str] = Field(
        default=None, description="A hint, if one was requested."
    )

    # Progress info
    steps_taken: int = Field(default=0, description="Number of actions taken so far.")
    max_steps: int = Field(default=5, description="Maximum allowed actions.")
    hints_used: int = Field(default=0, description="Number of hints used so far.")


class SQLState(State):
    challenge_id: str = ""
    broken_query: str = ""
    correct_query: str = ""
    schema_sql: str = ""
    schema_description: str = ""
    task_description: str = ""
    hints: List[str] = Field(default_factory=list)

    steps_taken: int = 0
    max_steps: int = 5
    hints_used: int = 0
    is_resolved: bool = False
    cumulative_reward: float = 0.0