ritvik360/nl2sql-bench

1"""2nl2sql-bench/models.py3======================4Typed contracts for the NL2SQL-Bench OpenEnv environment.5 6Action  : The SQL query the agent submits.7Observation : What the agent sees after each step.8State   : Episode-level metadata (for state() endpoint).9"""10 11from __future__ import annotations12 13from dataclasses import dataclass, field14from typing import Any, Dict, List, Optional15 16from openenv.core.env_server import Action, Observation, State17 18 19# ---------------------------------------------------------------------------20# Action21# ---------------------------------------------------------------------------22 23class NL2SQLAction(Action):24    """A single SQL query submitted by the agent."""25    query: str = ""26 27 28# ---------------------------------------------------------------------------29# Observation30# ---------------------------------------------------------------------------31 32 33class NL2SQLObservation(Observation):34    """35    Everything the agent needs to reason about and iterate its SQL query.36 37    Fields38    ------39    question        : The natural-language question to answer.40    schema_context  : Relevant table/column descriptions as a string block.41    task_name       : Identifier of the current task (easy / medium / hard).42    last_query      : The SQL the agent submitted on the last step (empty on reset).43    last_result     : Up to 10 rows returned by the last query (list of dicts).44    last_error      : SQLite error string if the query failed, else None.45    result_columns  : Column names of last_result rows.46    step            : Current step number (1-indexed).47    max_steps       : Maximum steps allowed per episode.48    done            : True when the episode is over (success or step exhausted).49    reward          : Reward for the most recent action (None on reset).50    score           : Normalised cumulative score so far [0.0, 1.0].51    """52    question: str = ""53    schema_context: str = ""54    task_name: str = ""55    last_query: str = ""56    last_result: List[Dict[str, Any]] = field(default_factory=list)57    last_error: Optional[str] = None58    result_columns: List[str] = field(default_factory=list)59    step: int = 060    max_steps: int = 561    done: bool = False62    reward: Optional[float] = None63    score: float = 0.064 65 66# ---------------------------------------------------------------------------67# State68# ---------------------------------------------------------------------------69 70class NL2SQLState(State):71    """Episode-level state (returned by the /state endpoint)."""72    episode_id: Optional[str] = None73    step_count: int = 074    task_name: str = ""75    task_difficulty: str = ""        # easy | medium | hard76    question: str = ""77    best_reward: float = 0.0         # highest reward seen this episode78    cumulative_reward: float = 0.079    solved: bool = False             # True if exact match was achieved80