| """ |
| SQLEnv Pydantic models — the data contracts between client and server. |
| |
| These models define the typed interface for the SQLEnv RL environment, |
| following the OpenEnv pattern (see OpenEnv Tutorial for reference): |
| |
| Action — what the agent sends each step |
| Observation — what the agent receives back |
| State — episode metadata (exposed via the state endpoint) |
| |
| RL terminology — state vs observation |
| ───────────────────────────────────── |
| In RL theory: |
| |
| State (s) A COMPLETE description of the world. Nothing is hidden. |
| Observation (o) A PARTIAL description of a state, which may omit info. |
| |
| In SQLEnv these map to: |
| |
| EpisodeContext The full RL state (s). Lives on the server only. |
| Contains gold answers, reward accumulators, DB |
| connection, full query history — everything needed |
| to advance the simulation and compute rewards. |
| |
| SQLObservation The observation (o). Sent to the agent over the wire. |
| Contains the question, truncated results, revealed |
| schema, budget, and action history. The agent NEVER |
| sees the gold answer, progress scores, or full DB. |
| |
| SQLState OpenEnv's "State" base class — lightweight episode |
| metadata (episode_id, step_count). This is NOT the |
| RL state; it is a convenience for logging/debugging. |
| |
| This separation is what makes SQLEnv a POMDP: the agent must act under |
| uncertainty, which is what makes exploration necessary and learnable. |
| """ |
|
|
| import sqlite3 |
| from dataclasses import dataclass, field as dataclass_field |
|
|
| from openenv.core.env_server.interfaces import Message |
| from openenv.core.env_server.types import Action, Observation, State |
| from pydantic import Field |
|
|
| |
| |
| |
|
|
|
|
| class SQLAction(Action): |
| """What the agent sends each step. |
| |
| The action space is intentionally small and structured so agents can |
| explicitly control the environment loop. |
| """ |
|
|
| action_type: str = Field( |
| ..., |
| description="One of: DESCRIBE, SAMPLE, QUERY, ANSWER", |
| ) |
| argument: str = Field( |
| ..., |
| description=( |
| "Table name (DESCRIBE/SAMPLE), SQL string (QUERY), " |
| "or answer value (ANSWER)." |
| ), |
| ) |
|
|
|
|
| class SQLObservation(Observation): |
| """What the agent receives after each step. |
| |
| This is the agent's PARTIAL view of the world. Key design choices: |
| |
| - schema_info starts with table names only; columns are revealed |
| incrementally as the agent DESCRIBEs tables. |
| - result is always a truncated string, never raw data. The agent sees |
| what a human analyst would see in a terminal — at most N rows of |
| formatted text. This keeps the observation bounded and forces the |
| agent to reason about what it sees rather than brute-force scanning. |
| - action_history gives the agent memory of its own trajectory without |
| the server needing to re-send full results from prior steps. |
| """ |
|
|
| |
| question: str = Field(..., description="The NL question to answer") |
| schema_info: str = Field(..., description="Known schema information") |
| result: str = Field(default="", description="Result of the last action") |
| error: str = Field(default="", description="Error message if action failed") |
| step_count: int = Field(default=0, description="Current step number") |
| budget_remaining: int = Field(default=0, description="Steps remaining") |
| action_history: list[str] = Field( |
| default_factory=list, |
| description="Summary of previous actions", |
| ) |
|
|
|
|
| class SQLState(State): |
| """Episode metadata exposed via GET /state. |
| |
| This is the minimal public state — enough for logging and debugging, |
| but NOT the full internal bookkeeping (see EpisodeContext below). |
| """ |
|
|
| |
| |
| |
| |
| history_messages: list[Message] = Field(default_factory=list) |
| current_action_type: str = Field( |
| default="QUERY", |
| description="Current action type: DESCRIBE, SAMPLE, QUERY, or ANSWER", |
| ) |
|
|
|
|
| @dataclass |
| class QuestionRecord: |
| """One question from the Spider dataset.""" |
|
|
| question_id: str |
| question_text: str |
| database_name: str |
| gold_sql: str |
| gold_answer: str |
| answer_type: str |
| difficulty: str |
| tables_involved: list[str] |
|
|
|
|
| @dataclass |
| class EpisodeContext: |
| """Per-episode server-side state (never sent to agent).""" |
|
|
| episode_id: str |
| db_connection: sqlite3.Connection |
| question_record: QuestionRecord |
| step_count: int = 0 |
| budget: int = 15 |
| described_tables: set[str] = dataclass_field(default_factory=set) |
| action_log: list[str] = dataclass_field(default_factory=list) |
| done: bool = False |
| gold_answer: str | None = None |
| gold_rows: list[tuple] = dataclass_field(default_factory=list) |
| query_hashes: set[str] = dataclass_field(default_factory=set) |
| previous_progress: float = 0.0 |
|
|
