""" FastAPI server exposing the OpenEnv HTTP API. Endpoints: POST /reset, POST /step, GET /state Also includes: GET /tasks (list available tasks), GET /health """ import asyncio import time import statistics from typing import Dict, Optional from contextlib import asynccontextmanager from fastapi import FastAPI, HTTPException, Header from fastapi.middleware.cors import CORSMiddleware from pydantic import BaseModel from .models import SQLDebugAction, SQLDebugObservation, EpisodeState from .env import SQLDebugEnv, TASKS # Session management: one env instance per session # For HF Space: allow up to 64 concurrent sessions MAX_SESSIONS = 64 _sessions: Dict[str, SQLDebugEnv] = {} _session_lock = asyncio.Lock() @asynccontextmanager async def lifespan(app: FastAPI): yield # Cleanup all sessions on shutdown for env in _sessions.values(): env.close() app = FastAPI( title="SQL Debug Environment", description="OpenEnv-compliant SQL query debugging environment for RL agent training.", version="0.1.0", lifespan=lifespan ) app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"], ) @app.get("/") async def root(): return { "name": "sql-debug-env", "status": "ok", "message": "Use /health, /tasks, /reset, /step, /state, /benchmark", } @app.get("/favicon.ico", status_code=204) async def favicon(): return None class ResetRequest(BaseModel): task_id: Optional[str] = "easy_syntax_fix" class StepRequest(BaseModel): action: SQLDebugAction async def get_or_create_session(session_id: str, task_id: str = "easy_syntax_fix") -> SQLDebugEnv: async with _session_lock: if session_id not in _sessions: if len(_sessions) >= MAX_SESSIONS: # Evict oldest session oldest = next(iter(_sessions)) _sessions[oldest].close() del _sessions[oldest] _sessions[session_id] = SQLDebugEnv(task_id=task_id) return _sessions[session_id] @app.get("/health") async def health(): return {"status": "ok", "sessions_active": len(_sessions)} @app.get("/tasks") async def list_tasks(): """List all available tasks with metadata.""" return { "tasks": [task.to_dict() for task in TASKS.values()] } def _stats(values: list[float]) -> Dict[str, float]: ordered = sorted(values) n = len(ordered) p95_idx = max(0, int(n * 0.95) - 1) return { "avg_ms": round(statistics.mean(ordered), 3), "p50_ms": round(statistics.median(ordered), 3), "p95_ms": round(ordered[p95_idx], 3), "n": n, } @app.get("/benchmark") async def benchmark(runs: int = 20): """ Real-time benchmark endpoint (fresh measurements on every call). Safe to call from dashboards/web pages for live verification. """ runs = max(1, min(runs, 100)) health_times: list[float] = [] tasks_times: list[float] = [] reset_times: list[float] = [] step_times: list[float] = [] bench_env = SQLDebugEnv(task_id="easy_syntax_fix") try: for _ in range(runs): t0 = time.perf_counter() _ = {"status": "ok", "sessions_active": len(_sessions)} health_times.append((time.perf_counter() - t0) * 1000) t0 = time.perf_counter() _ = [task.to_dict() for task in TASKS.values()] tasks_times.append((time.perf_counter() - t0) * 1000) t0 = time.perf_counter() await bench_env.reset() reset_times.append((time.perf_counter() - t0) * 1000) t0 = time.perf_counter() await bench_env.step(SQLDebugAction(action_type="inspect_schema")) step_times.append((time.perf_counter() - t0) * 1000) finally: bench_env.close() return { "benchmark": { "runs": runs, "task_id": "easy_syntax_fix", "timestamp_epoch_ms": int(time.time() * 1000), "results": { "health": _stats(health_times), "tasks": _stats(tasks_times), "reset": _stats(reset_times), "step_inspect_schema": _stats(step_times), }, } } @app.post("/reset") async def reset( request: ResetRequest = ResetRequest(), x_session_id: Optional[str] = Header(default=None) ): """ Reset the environment for a new episode. Returns initial observation with task description and broken query. """ session_id = x_session_id or "default" task_id = request.task_id or "easy_syntax_fix" if task_id not in TASKS: raise HTTPException(status_code=400, detail=f"Unknown task_id: {task_id}. Valid: {list(TASKS.keys())}") # Always create fresh env on reset async with _session_lock: if session_id in _sessions: _sessions[session_id].close() _sessions[session_id] = SQLDebugEnv(task_id=task_id) env = _sessions[session_id] observation, info = await env.reset() return { "observation": observation.model_dump(), "info": info, "reward": None, "done": False } @app.post("/step") async def step( request: StepRequest, x_session_id: Optional[str] = Header(default=None) ): """ Execute one action in the environment. Action types: - submit_query: Submit SQL for evaluation (requires 'query' field) - inspect_schema: Get table schema (free action) - inspect_error: Get last error message (free action) - inspect_sample: Get sample rows from table (requires 'table_name') - reset_query: Reset to original broken query (small penalty) """ session_id = x_session_id or "default" if session_id not in _sessions: raise HTTPException(status_code=400, detail="Session not found. Call /reset first.") env = _sessions[session_id] try: observation, reward, done, info = await env.step(request.action) except RuntimeError as e: raise HTTPException(status_code=400, detail=str(e)) except ValueError as e: raise HTTPException(status_code=422, detail=str(e)) return { "observation": observation.model_dump(), "reward": reward, "done": done, "info": info } @app.get("/state") async def state(x_session_id: Optional[str] = Header(default=None)): """Return current full episode state.""" session_id = x_session_id or "default" if session_id not in _sessions: raise HTTPException(status_code=400, detail="No active session. Call /reset first.") env = _sessions[session_id] try: current_state = env.get_state() return current_state.model_dump() except RuntimeError as e: raise HTTPException(status_code=400, detail=str(e))