| """ |
| SQL Query Optimizer — OpenEnv Environment Server |
| FastAPI server implementing the OpenEnv spec: |
| POST /reset → initial observation |
| POST /step → observation, reward, done, info |
| GET /state → current state |
| """ |
|
|
| from __future__ import annotations |
|
|
| import re |
| import time |
| import uuid |
| from typing import Any |
|
|
| import sqlglot |
| import sqlglot.errors |
| from fastapi import FastAPI, HTTPException |
| from fastapi.responses import RedirectResponse |
| from fastapi.middleware.cors import CORSMiddleware |
| from pydantic import BaseModel, Field |
|
|
| from .tasks import TASKS, Task |
| from .grader import grade_submission |
|
|
| |
| |
| |
|
|
| class ExecutionStats(BaseModel): |
| estimated_cost: float = Field(description="Query planner estimated cost (lower = better)") |
| scan_type: str = Field(description="Dominant scan type: seq_scan | index_scan | index_only_scan") |
| rows_examined: int = Field(description="Estimated rows examined") |
| has_anti_patterns: list[str] = Field(default_factory=list, description="Detected anti-patterns") |
|
|
|
|
| class Observation(BaseModel): |
| task_id: str |
| schema_ddl: str = Field(description="Table DDL describing the database schema") |
| slow_query: str = Field(description="The SQL query to be optimized") |
| execution_stats: ExecutionStats = Field(description="Stats for the slow (original) query") |
| step_number: int = Field(default=0) |
| last_reward: float | None = Field(default=None) |
| last_feedback: str | None = Field(default=None) |
| history: list[str] = Field(default_factory=list, description="Previous action summaries") |
| done: bool = Field(default=False) |
|
|
|
|
| class Action(BaseModel): |
| optimized_query: str = Field(description="Your rewritten, optimized SQL query") |
| explanation: str = Field(default="", description="Brief explanation of changes made") |
|
|
|
|
| class Reward(BaseModel): |
| total: float = Field(ge=0.0, le=1.0, description="Combined score 0.0–1.0") |
| correctness: float = Field(ge=0.0, le=1.0, description="Does it return the same result set?") |
| performance: float = Field(ge=0.0, le=1.0, description="Estimated cost reduction ratio") |
| style: float = Field(ge=0.0, le=1.0, description="Absence of anti-patterns, readability") |
| feedback: str = Field(description="Human-readable explanation of the score") |
|
|
|
|
| class StepResult(BaseModel): |
| observation: Observation |
| reward: Reward |
| done: bool |
| info: dict[str, Any] = Field(default_factory=dict) |
|
|
|
|
| class StateResponse(BaseModel): |
| session_id: str |
| task_id: str | None |
| step_number: int |
| total_reward: float |
| best_reward: float |
| is_done: bool |
|
|
|
|
| |
| |
| |
|
|
| class Session: |
| def __init__(self, task: Task) -> None: |
| self.session_id = str(uuid.uuid4()) |
| self.task = task |
| self.step_number = 0 |
| self.total_reward = 0.0 |
| self.best_reward = 0.0 |
| self.is_done = False |
| self.history: list[str] = [] |
| self.created_at = time.time() |
|
|
| def to_initial_observation(self) -> Observation: |
| return Observation( |
| task_id=self.task.task_id, |
| schema_ddl=self.task.schema_ddl, |
| slow_query=self.task.slow_query, |
| execution_stats=ExecutionStats(**self.task.execution_stats), |
| step_number=0, |
| history=[], |
| done=False, |
| ) |
|
|
| def to_state(self) -> StateResponse: |
| return StateResponse( |
| session_id=self.session_id, |
| task_id=self.task.task_id, |
| step_number=self.step_number, |
| total_reward=round(self.total_reward, 4), |
| best_reward=round(self.best_reward, 4), |
| is_done=self.is_done, |
| ) |
|
|
|
|
| _sessions: dict[str, Session] = {} |
| _current_session_id: str | None = None |
|
|
| |
| |
| |
|
|
| app = FastAPI( |
| title="SQL Query Optimizer — OpenEnv", |
| description="AI agent environment for learning to optimize SQL queries.", |
| version="1.0.0", |
| ) |
| app.add_middleware( |
| CORSMiddleware, |
| allow_origins=["*"], |
| allow_methods=["*"], |
| allow_headers=["*"], |
| ) |
|
|
|
|
| |
| |
| |
|
|
| @app.post("/reset", response_model=Observation, summary="Reset environment to a new task") |
| def reset(task_id: str | None = None) -> Observation: |
| """ |
| Start a new episode. Optionally specify a task_id; otherwise the first |
| task (easy) is selected. Returns the initial observation. |
| """ |
| global _current_session_id |
|
|
| if task_id is None: |
| task_id = list(TASKS.keys())[0] |
|
|
| if task_id not in TASKS: |
| raise HTTPException(status_code=400, detail=f"Unknown task_id '{task_id}'. " |
| f"Valid: {list(TASKS.keys())}") |
|
|
| task = TASKS[task_id] |
| session = Session(task) |
| _sessions[session.session_id] = session |
| _current_session_id = session.session_id |
|
|
| return session.to_initial_observation() |
|
|
|
|
| @app.post("/step", response_model=StepResult, summary="Submit an optimized query") |
| def step(action: Action) -> StepResult: |
| """ |
| Submit the agent's optimized SQL query. Returns updated observation, |
| reward breakdown, done flag, and debug info. |
| """ |
| global _current_session_id |
|
|
| if _current_session_id is None or _current_session_id not in _sessions: |
| raise HTTPException(status_code=400, detail="No active session. Call /reset first.") |
|
|
| session = _sessions[_current_session_id] |
|
|
| if session.is_done: |
| raise HTTPException(status_code=400, detail="Episode already done. Call /reset to start a new one.") |
|
|
| session.step_number += 1 |
| MAX_STEPS = 5 |
|
|
| |
| reward_obj, grader_info = grade_submission( |
| task=session.task, |
| optimized_query=action.optimized_query, |
| explanation=action.explanation, |
| ) |
|
|
| session.total_reward += reward_obj.total |
| session.best_reward = max(session.best_reward, reward_obj.total) |
|
|
| |
| summary = ( |
| f"Step {session.step_number}: score={reward_obj.total:.2f} " |
| f"[correct={reward_obj.correctness:.2f}, perf={reward_obj.performance:.2f}, " |
| f"style={reward_obj.style:.2f}]" |
| ) |
| session.history.append(summary) |
|
|
| |
| done = reward_obj.total >= 0.95 or session.step_number >= MAX_STEPS |
| session.is_done = done |
|
|
| obs = Observation( |
| task_id=session.task.task_id, |
| schema_ddl=session.task.schema_ddl, |
| slow_query=session.task.slow_query, |
| execution_stats=ExecutionStats(**session.task.execution_stats), |
| step_number=session.step_number, |
| last_reward=reward_obj.total, |
| last_feedback=reward_obj.feedback, |
| history=session.history.copy(), |
| done=done, |
| ) |
|
|
| return StepResult( |
| observation=obs, |
| reward=reward_obj, |
| done=done, |
| info={**grader_info, "session_id": session.session_id}, |
| ) |
|
|
|
|
| @app.get("/state", response_model=StateResponse, summary="Current session state") |
| def state() -> StateResponse: |
| """Return the current session state without advancing the episode.""" |
| if _current_session_id is None or _current_session_id not in _sessions: |
| raise HTTPException(status_code=400, detail="No active session. Call /reset first.") |
| return _sessions[_current_session_id].to_state() |
|
|
|
|
| @app.get("/tasks", summary="List all available tasks") |
| def list_tasks() -> dict: |
| return { |
| "tasks": [ |
| { |
| "task_id": t.task_id, |
| "difficulty": t.difficulty, |
| "description": t.description, |
| } |
| for t in TASKS.values() |
| ] |
| } |
|
|
|
|
| @app.get("/health") |
| def health() -> dict: |
| return {"status": "ok", "version": "1.0.0"} |
|
|
|
|
| @app.get("/", summary="Redirect to API docs", include_in_schema=False) |
| def root(): |
| return RedirectResponse(url="/docs") |
|
|
|
|
| def main(): |
| import uvicorn |
| uvicorn.run("server.app:app", host="0.0.0.0", port=7860) |
|
|
| if __name__ == "__main__": |
| main() |
|
|
|
|
