from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Request
from fastapi.responses import HTMLResponse, RedirectResponse
from pydantic import BaseModel
from typing import Optional
import json
import asyncio
from server.environment import SQLEnvironment
app = FastAPI(title="SQL Debugger OpenEnv", version="1.0.0")
# One environment instance per HTTP session (stateless endpoints)
# WebSocket sessions get their own instance
env = SQLEnvironment()
# ── Pydantic models ───────────────────────────────────────────────────────────
class ResetRequest(BaseModel):
task_id: Optional[str] = None
difficulty: Optional[str] = None
class StepRequest(BaseModel):
action_type: str = "write_query"
sql: str = ""
explanation: Optional[str] = None
metadata: dict = {}
class GradeRequest(BaseModel):
task_id: str
action: dict
# ── HTTP endpoints ────────────────────────────────────────────────────────────
@app.get("/")
def root(request: Request):
# Smart routing: If a human browser loads the root URL, direct them to the UI!
# If the Hackathon Validator automated bot calls the root URL, serve the required JSON.
accept = request.headers.get("accept", "")
if "text/html" in accept:
return RedirectResponse(url="/web")
return {
"info": "SQL Agent Environment API",
"version": "1.0.0",
"tasks": 6,
"status": "running"
}
@app.get("/health")
def health():
return {"status": "healthy"}
@app.post("/reset")
def reset(req: ResetRequest = ResetRequest()):
return env.reset(task_id=req.task_id, difficulty=req.difficulty)
@app.post("/step")
def step(req: StepRequest):
return env.step(req.dict())
@app.get("/state")
def state():
return env.state
# ── CRITICAL MISSING ENDPOINTS — ADD THESE ───────────────────────────────────
@app.get("/tasks")
def list_tasks():
"""
Validator uses this endpoint to find and enumerate all graders.
Every task MUST have grader: true.
"""
return {
"tasks": [
{
"id": "easy_01",
"name": "Filter high earners",
"difficulty": "easy",
"type": "write_query",
"grader": True,
"description": "Find all employees in Engineering with salary above 90000, return name and salary ordered by salary DESC"
},
{
"id": "easy_02",
"name": "Count by department",
"difficulty": "easy",
"type": "write_query",
"grader": True,
"description": "Count employees per department, return department and count ordered by count DESC"
},
{
"id": "medium_01",
"name": "Fix broken JOIN",
"difficulty": "medium",
"type": "fix_query",
"grader": True,
"description": "Fix the broken JOIN query with missing ON keyword and wrong table alias in WHERE"
},
{
"id": "medium_02",
"name": "Fix wrong GROUP BY",
"difficulty": "medium",
"type": "fix_query",
"grader": True,
"description": "Fix the GROUP BY that incorrectly groups by id instead of department"
},
{
"id": "hard_01",
"name": "Optimize correlated subquery",
"difficulty": "hard",
"type": "optimize_query",
"grader": True,
"description": "Replace correlated subqueries with window functions or CTEs to find top earner per department"
},
{
"id": "hard_02",
"name": "Eliminate N+1 problem",
"difficulty": "hard",
"type": "optimize_query",
"grader": True,
"description": "Rewrite N+1 correlated subquery using LEFT JOIN and GROUP BY"
}
]
}
@app.post("/grade")
def grade(req: GradeRequest):
"""
Validator calls this to get a score for a specific task.
Score MUST be strictly between 0 and 1 (never 0.0, never 1.0).
This is enforced by _clamp() in the environment AND by the
absolute safety clamp below.
"""
# Use a fresh environment instance for grading to avoid state pollution
grade_env = SQLEnvironment()
grade_env.reset(task_id=req.task_id)
result = grade_env.step(req.action)
raw_score = result.get("reward", 0.05)
# ABSOLUTE SAFETY CLAMP — even if environment has a bug,
# this line guarantees the validator never sees 0.0 or 1.0
score = round(max(0.05, min(0.95, float(raw_score))), 4)
return {
"task_id": req.task_id,
"score": score,
"feedback": result["observation"]["feedback"],
"done": result["done"],
"info": result.get("info", {})
}
# ── WebSocket endpoint — required by OpenEnv spec ─────────────────────────────
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
"""
Persistent WebSocket session. Each connection gets its own
SQLEnvironment instance so sessions are isolated.
"""
await websocket.accept()
ws_env = SQLEnvironment()
try:
while True:
data = await websocket.receive_text()
try:
message = json.loads(data)
except json.JSONDecodeError:
await websocket.send_text(json.dumps({
"error": "Invalid JSON"
}))
continue
msg_type = message.get("type", "")
if msg_type == "reset":
result = ws_env.reset(
task_id=message.get("task_id"),
difficulty=message.get("difficulty")
)
await websocket.send_text(json.dumps(result))
elif msg_type == "step":
action = message.get("action", {})
result = ws_env.step(action)
await websocket.send_text(json.dumps(result))
elif msg_type == "state":
await websocket.send_text(json.dumps(ws_env.state))
else:
await websocket.send_text(json.dumps({
"error": f"Unknown message type: {msg_type}"
}))
except WebSocketDisconnect:
pass
# ── Web UI ─────────────────────────────────────────────────────────────────────
@app.get("/web", response_class=HTMLResponse)
def web_ui():
"""
Interactive web playground for testing the environment manually.
"""
return """
SQL data analyst playground
Interactive mode uses a persistent WebSocket session at /ws
(OpenEnv HTTP /step is stateless). Connect, then run reset → SQL steps.
CURRENT TASK
—
—
Click "Start episode" to load the first task.
YOUR SQL
Type only valid SQLite here.
LAST RESULT
No result yet.
Not connected. Click "Start episode" to begin.
"""
def main():
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=7860)
if __name__ == "__main__":
main()