Tsah00/sql-env

1"""2models.py - Typed Pydantic models for the SQL Query Learning Environment.3 4Defines Action, Observation, and State models compatible with the OpenEnv specification.5Uses Pydantic BaseModel for typed validation. Falls back gracefully if openenv-core6is not installed.7"""8 9from __future__ import annotations10 11from typing import Any, Dict, List, Optional12 13from pydantic import BaseModel, Field14 15# Try to import OpenEnv base classes for full spec compliance16try:17    from openenv.core.env_server.interfaces import (18        Action as _BaseAction,19        Observation as _BaseObservation,20        State as _BaseState,21    )22except ImportError:23    # openenv-core not installed — use plain Pydantic24    _BaseAction = BaseModel25    _BaseObservation = BaseModel26    _BaseState = BaseModel27 28 29class SQLAction(_BaseAction):30    """31    Action: a SQL query submitted by the agent.32 33    Fields:34        query: The SQL query string to execute.35        task_id: Optional task ID to target (defaults to current task).36        difficulty: Task difficulty tier — 'easy', 'medium', or 'hard'.37    """38 39    query: str40    task_id: Optional[str] = None41    difficulty: str = "easy"42 43 44class SQLObservation(_BaseObservation):45    """46    Observation returned after executing a SQL action.47 48    Fields:49        result: Rows returned by the query (list of dicts).50        error: Error message if the query failed, else empty.51        reward: Step reward in [0.0, 1.0].52        done: Whether the episode is finished.53        message: Human-readable grader feedback.54        schema_info: Database schema description.55        task_description: Natural-language task objective.56        expected_columns: Column names the answer should include.57        step_count: Steps taken so far in this episode.58        score_breakdown: Sub-scores (correctness, efficiency, etc.).59    """60 61    result: List[Dict[str, Any]] = Field(default_factory=list)62    error: str = ""63    reward: float = 0.064    done: bool = False65    message: str = ""66    schema_info: str = ""67    task_description: str = ""68    expected_columns: List[str] = Field(default_factory=list)69    step_count: int = 070    score_breakdown: Dict[str, float] = Field(default_factory=dict)71 72 73class SQLState(_BaseState):74    """75    Full state of an ongoing episode.76 77    Fields:78        episode_id: Unique episode identifier.79        step_count: Steps taken so far.80        current_task_id: Active task ID.81        current_difficulty: Active difficulty tier.82        total_reward: Cumulative reward.83        tasks_completed: Count of tasks solved.84        best_reward: Best single-step reward.85        last_query: Most recent SQL query submitted.86        last_error: Most recent error (empty if none).87    """88 89    episode_id: str = ""90    step_count: int = 091    current_task_id: str = ""92    current_difficulty: str = "easy"93    total_reward: float = 0.094    tasks_completed: int = 095    best_reward: float = 0.096    last_query: str = ""97    last_error: str = ""98