File size: 5,456 Bytes
5dd1bb4 9e64e71 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 | """
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
# ---------------------------------------------------------------------------
# Wire types: these cross the HTTP boundary between client and server
# ---------------------------------------------------------------------------
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.
"""
# Inherited from Observation: done (bool), reward (float | None)
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).
"""
# # Inherited from State: episode_id (str | None), step_count (int)
# game_name: str = Field(
# "sql_env", description="Name of the game/environment"
# )
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
|