Spaces:
Sleeping
Sleeping
File size: 2,092 Bytes
798461f | 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 | """
models.py
---------
Action / Observation / State typed models for the SQL Analytics environment.
The agent submits a SQL query (SQLAction) and receives back the query
result rows, any SQL error, the current task description, and reward
signal (SQLObservation). Episode metadata is tracked in SQLState.
"""
from __future__ import annotations
from typing import Any, Dict, List, Optional
from pydantic import Field, field_validator
try:
from openenv.core.env_server.interfaces import Action, Observation, State
except ImportError:
from openenv.core.env_server.types import Action, Observation, State
class SQLAction(Action):
"""
A single SQL query submitted by the agent.
Example::
SQLAction(sql="SELECT COUNT(*) FROM customers WHERE state = 'CA'")
"""
sql: str = Field(
...,
min_length=1,
description="SQL SELECT query to execute against the e-commerce database.",
json_schema_extra={"example": "SELECT COUNT(*) FROM customers WHERE state = 'CA'"},
)
@field_validator("sql")
@classmethod
def strip_whitespace(cls, v: str) -> str:
"""Strip leading/trailing whitespace from the query."""
return v.strip()
class SQLObservation(Observation):
"""
Result of executing a SQL query in the environment.
Fields
------
result : Rows returned by the query (list of dicts), None on error.
error : SQL error message if the query failed, else None.
db_schema : Human-readable database schema (populated on reset).
task_description: Natural-language question the agent must answer.
"""
result: Optional[List[Dict[str, Any]]] = None
error: Optional[str] = None
db_schema: Optional[str] = None
task_description: Optional[str] = None
class SQLState(State):
"""Episode-level state for the SQL Analytics environment."""
task_id: str = ""
task_description: str = ""
difficulty: str = "" # "easy" | "medium" | "hard"
attempts: int = 0
max_attempts: int = 10
solved: bool = False
|