Tsah00/sql-env
0
1"""2server/app.py - FastAPI server exposing the SQL Environment over HTTP and WebSocket.3 4Endpoints:5 POST /reset - Start a new episode6 POST /step - Execute a SQL action7 GET /state - Get current environment state8 GET /health - Health check (returns HTTP 200)9 WebSocket /ws - Real-time bidirectional communication10 11Compatible with the OpenEnv client interface.12 13Usage:14 uvicorn server.app:app --host 0.0.0.0 --port 786015"""16 17from __future__ import annotations18 19import asyncio20import json21import logging22import uuid23from contextlib import asynccontextmanager24from typing import Any as _Any25from typing import Any, Dict, Optional26 27from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Body28from fastapi.middleware.cors import CORSMiddleware29from fastapi.responses import HTMLResponse, JSONResponse30from pydantic import BaseModel31 32import sys33import os34sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))35 36from models import SQLAction, SQLObservation, SQLState37from server.sql_environment import SQLEnvironment38 39logging.basicConfig(level=logging.INFO)40logger = logging.getLogger(__name__)41 42# ---------------------------------------------------------------------------43# Pydantic request/response models (for FastAPI schema generation)44# ---------------------------------------------------------------------------45 46class ResetRequest(BaseModel):47 difficulty: str = "easy"48 task_id: Optional[str] = None49 50 51class StepRequest(BaseModel):52 query: str53 task_id: Optional[str] = None54 difficulty: str = "easy"55 56 57# ---------------------------------------------------------------------------58# Session management59# ---------------------------------------------------------------------------60 61# Map episode_id -> SQLEnvironment instance62_sessions: Dict[str, SQLEnvironment] = {}63_default_session_id = "default"64 65 66def _get_or_create_session(session_id: str = _default_session_id) -> SQLEnvironment:67 if session_id not in _sessions:68 _sessions[session_id] = SQLEnvironment()69 return _sessions[session_id]70 71 72# ---------------------------------------------------------------------------73# Application lifecycle74# ---------------------------------------------------------------------------75 76@asynccontextmanager77async def lifespan(app: FastAPI):78 """Startup: initialise default session. Shutdown: close all sessions."""79 env = _get_or_create_session(_default_session_id)80 env.reset()81 logger.info("SQL Environment server started. Default session ready.")82 yield83 for env in _sessions.values():84 env.close()85 logger.info("SQL Environment server stopped.")86 87 88app = FastAPI(89 title="SQL Query Learning Environment",90 description=(91 "An OpenEnv-compatible reinforcement learning environment where "92 "agents learn to write correct and efficient SQL queries."93 ),94 version="1.0.0",95 lifespan=lifespan,96)97 98app.add_middleware(99 CORSMiddleware,100 allow_origins=["*"],101 allow_methods=["*"],102 allow_headers=["*"],103)104 105 106# ---------------------------------------------------------------------------107# HTTP Endpoints108# ---------------------------------------------------------------------------109 110@app.get("/health")111async def health():112 """Health check - required by OpenEnv deployment validation."""113 return JSONResponse({"status": "ok", "environment": "sql_env", "version": "1.0.0"})114 115 116@app.post("/reset")117async def reset(118 request: Optional[ResetRequest] = Body(default=None),119 session_id: str = _default_session_id,120):121 """122 Initialize a new episode.123 124 Returns:125 SQLObservation as JSON with schema_info and first task description.126 """127 env = _get_or_create_session(session_id)128 difficulty = request.difficulty if request else "easy"129 task_id = request.task_id if request else None130 obs: SQLObservation = env.reset(difficulty=difficulty, task_id=task_id)131 return _obs_to_dict(obs)132 133 134@app.post("/step")135async def step(request: StepRequest, session_id: str = _default_session_id):136 """137 Execute a SQL query action.138 139 Returns:140 SQLObservation as JSON with result rows, reward, and feedback.141 """142 env = _get_or_create_session(session_id)143 action = SQLAction(144 query=request.query,145 task_id=request.task_id,146 difficulty=request.difficulty,147 )148 obs: SQLObservation = env.step(action)149 return _obs_to_dict(obs)150 151 152@app.get("/state")153async def get_state(session_id: str = _default_session_id):154 """155 Return current episode state metadata.156 157 Returns:158 SQLState as JSON with episode_id, step_count, total_reward, etc.159 """160 env = _get_or_create_session(session_id)161 return env.state.model_dump()162 163 164@app.get("/tasks")165async def list_tasks():166 """List all available tasks grouped by difficulty."""167 from server.tasks import TASKS168 grouped: Dict[str, list] = {"easy": [], "medium": [], "hard": []}169 for task in TASKS.values():170 entry = {171 "id": task["id"],172 "description": task["description"],173 "expected_columns": task["expected_columns"],174 }175 grouped[task["difficulty"]].append(entry)176 return grouped177 178 179@app.get("/schema")180async def get_schema():181 """Return the database schema description."""182 from server.tasks import SCHEMA_INFO183 return {"schema": SCHEMA_INFO}184 185 186@app.get("/", response_class=HTMLResponse)187async def root():188 """Simple web interface for manual testing."""189 return HTMLResponse(content=_web_ui_html(), status_code=200)190 191 192# ---------------------------------------------------------------------------193# WebSocket Endpoint194# ---------------------------------------------------------------------------195 196@app.websocket("/ws")197async def websocket_endpoint(websocket: WebSocket):198 """199 WebSocket endpoint for real-time agent interaction.200 201 Message format (JSON):202 { "action": "reset", "difficulty": "easy" }203 { "action": "step", "query": "SELECT ...", "difficulty": "easy" }204 { "action": "state" }205 206 Response format (JSON):207 { "type": "observation"|"state"|"error", "data": {...} }208 """209 await websocket.accept()210 session_id = str(uuid.uuid4())211 env = SQLEnvironment()212 _sessions[session_id] = env213 logger.info(f"WebSocket session opened: {session_id}")214 215 try:216 while True:217 raw = await websocket.receive_text()218 try:219 msg = json.loads(raw)220 except json.JSONDecodeError:221 await websocket.send_text(json.dumps({222 "type": "error",223 "data": {"message": "Invalid JSON"},224 }))225 continue226 227 action_type = msg.get("action", "")228 229 if action_type == "reset":230 obs = env.reset(231 difficulty=msg.get("difficulty", "easy"),232 task_id=msg.get("task_id"),233 )234 await websocket.send_text(json.dumps({235 "type": "observation",236 "data": _obs_to_dict(obs),237 }))238 239 elif action_type == "step":240 query = msg.get("query", "")241 if not query:242 await websocket.send_text(json.dumps({243 "type": "error",244 "data": {"message": "Missing 'query' field"},245 }))246 continue247 action = SQLAction(248 query=query,249 task_id=msg.get("task_id"),250 difficulty=msg.get("difficulty", "easy"),251 )252 obs = env.step(action)253 await websocket.send_text(json.dumps({254 "type": "observation",255 "data": _obs_to_dict(obs),256 }))257 258 elif action_type == "state":259 await websocket.send_text(json.dumps({260 "type": "state",261 "data": env.state.model_dump(),262 }))263 264 else:265 await websocket.send_text(json.dumps({266 "type": "error",267 "data": {"message": f"Unknown action: {action_type}"},268 }))269 270 except WebSocketDisconnect:271 logger.info(f"WebSocket session closed: {session_id}")272 finally:273 env.close()274 _sessions.pop(session_id, None)275 276 277# ---------------------------------------------------------------------------278# Helpers279# ---------------------------------------------------------------------------280 281def _obs_to_dict(obs: SQLObservation) -> Dict[str, Any]:282 return obs.model_dump()283 284 285def _web_ui_html() -> str:286 return """287<!DOCTYPE html>288<html lang="en">289<head>290 <meta charset="UTF-8">291 <title>SQL Query Learning Environment</title>292 <style>293 body { font-family: monospace; max-width: 900px; margin: 40px auto; padding: 0 20px; background: #1e1e2e; color: #cdd6f4; }294 h1 { color: #89b4fa; }295 h2 { color: #a6e3a1; font-size: 1em; margin-top: 20px; }296 textarea { width: 100%; height: 120px; background: #313244; color: #cdd6f4; border: 1px solid #585b70; padding: 8px; font-family: monospace; font-size: 13px; }297 button { background: #89b4fa; color: #1e1e2e; border: none; padding: 8px 16px; cursor: pointer; margin: 4px; font-weight: bold; }298 button:hover { background: #74c7ec; }299 select { background: #313244; color: #cdd6f4; border: 1px solid #585b70; padding: 6px; }300 pre { background: #313244; padding: 12px; overflow: auto; font-size: 12px; max-height: 400px; border: 1px solid #585b70; }301 .reward { color: #a6e3a1; font-size: 1.2em; font-weight: bold; }302 .error { color: #f38ba8; }303 .task { background: #2a273f; padding: 10px; border-left: 3px solid #cba6f7; margin: 10px 0; }304 </style>305</head>306<body>307 <h1>SQL Query Learning Environment</h1>308 <p>An OpenEnv RL environment for learning SQL against an e-commerce database.</p>309 310 <h2>1. Start Episode</h2>311 Difficulty:312 <select id="difficulty">313 <option value="easy">Easy</option>314 <option value="medium">Medium</option>315 <option value="hard">Hard</option>316 </select>317 <button onclick="doReset()">Reset / New Episode</button>318 319 <div id="taskBox" class="task" style="display:none">320 <b>Task:</b> <span id="taskDesc"></span><br>321 <b>Expected columns:</b> <span id="taskCols"></span>322 </div>323 324 <h2>2. Submit SQL Query</h2>325 <textarea id="query" placeholder="SELECT name, email FROM customers WHERE country = 'USA'"></textarea>326 <br>327 <button onclick="doStep()">Execute Query</button>328 <button onclick="doState()">Get State</button>329 330 <h2>3. Result</h2>331 <div class="reward">Reward: <span id="reward">-</span></div>332 <div id="message"></div>333 <pre id="output">Response will appear here...</pre>334 335 <script>336 const base = window.location.origin;337 338 async function doReset() {339 const diff = document.getElementById('difficulty').value;340 const res = await fetch(base + '/reset', {341 method: 'POST',342 headers: {'Content-Type': 'application/json'},343 body: JSON.stringify({ difficulty: diff })344 });345 const data = await res.json();346 document.getElementById('output').textContent = JSON.stringify(data, null, 2);347 document.getElementById('reward').textContent = data.reward ?? '-';348 document.getElementById('message').textContent = data.message ?? '';349 document.getElementById('taskDesc').textContent = data.task_description ?? '';350 document.getElementById('taskCols').textContent = (data.expected_columns ?? []).join(', ');351 document.getElementById('taskBox').style.display = 'block';352 }353 354 async function doStep() {355 const query = document.getElementById('query').value;356 const diff = document.getElementById('difficulty').value;357 const res = await fetch(base + '/step', {358 method: 'POST',359 headers: {'Content-Type': 'application/json'},360 body: JSON.stringify({ query, difficulty: diff })361 });362 const data = await res.json();363 document.getElementById('output').textContent = JSON.stringify(data, null, 2);364 document.getElementById('reward').textContent = data.reward ?? '-';365 document.getElementById('message').textContent = data.message ?? '';366 if (data.task_description) {367 document.getElementById('taskDesc').textContent = data.task_description;368 document.getElementById('taskCols').textContent = (data.expected_columns ?? []).join(', ');369 }370 }371 372 async function doState() {373 const res = await fetch(base + '/state');374 const data = await res.json();375 document.getElementById('output').textContent = JSON.stringify(data, null, 2);376 }377 </script>378</body>379</html>380""".strip()381 382 383def main():384 import uvicorn385 uvicorn.run(app, host="0.0.0.0", port=7860)386 387 388if __name__ == "__main__":389 main()390