models.py

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


class SQLAction(Action):
    corrected_query: str


class SQLObservation(Observation):
    broken_query: str = ""
    table_schema: List[str] = []        # renamed from "schema" — parent class has that name
    description: str = ""
    expected_output: Optional[list] = None
    actual_output: Optional[list] = None
    score: float = 0.0
    error_message: Optional[str] = None
    # done is inherited from Observation (default False)
    # reward is inherited from Observation (default None)
    attempts_left: int = 3


class SQLState(State):
    task_id: str = ""
    difficulty: str = ""
    # step_count is inherited from State (default 0)
    # episode_id is inherited from State (default None)