app.py

"""
server/app.py — FastAPI Server
================================
OpenEnv-compliant endpoints + two unique endpoints:
  POST /execute    — run your optimized query against real DuckDB data,
                     see actual speedup + result correctness instantly
  GET  /leaderboard — see best scores + speedups across all tasks
"""

import json
import os
import sys
from contextlib import asynccontextmanager
from pathlib import Path

from fastapi import FastAPI, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import HTMLResponse

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from env import SQLOptimEnv
from executor import get_executor
from graders import grade
from leaderboard import get_board
from models import (
    Action,
    EnvironmentState,
    ExecutionResult,
    Observation,
    StepResult,
)
from tasks import TASKS, get_task_list


# ── Lifespan: pre-warm DuckDB on startup ─────────────────────────────────

@asynccontextmanager
async def lifespan(app: FastAPI):
    # Build all 4 synthetic tables before first request
    get_executor()
    yield


app = FastAPI(
    title="SQL Query Optimization Environment",
    description=(
        "OpenEnv-compliant RL environment where AI agents learn to diagnose "
        "and optimize SQL queries. Uniquely, optimized queries are EXECUTED "
        "against real DuckDB data — reward is based on actual speedup + "
        "result correctness, not keyword heuristics."
    ),
    version="2.0.0",
    lifespan=lifespan,
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

env = SQLOptimEnv()


# ── Serve interactive demo page ─────────────────────────────────────────
_DEMO_HTML = Path(__file__).parent / "demo.html"

@app.get("/demo", response_class=HTMLResponse, include_in_schema=False)
def demo_page():
    """Interactive SQL optimizer demo — paste SQL, hit Execute, see real DuckDB timing."""
    if _DEMO_HTML.exists():
        return HTMLResponse(content=_DEMO_HTML.read_text(encoding="utf-8"))
    return HTMLResponse(content="<h1>demo.html not found</h1>", status_code=404)


# ── Standard OpenEnv endpoints ────────────────────────────────────────────

@app.get("/")
def root():
    ex = get_executor()
    return {
        "status":      "ok",
        "environment": "sql-optim-env",
        "version":     "2.0.0",
        "unique_feature": "Execution-grounded rewards via DuckDB",
        "table_stats": ex.table_stats,
        "tasks":       [t["task_id"] for t in get_task_list()],
    }


@app.post("/reset", response_model=Observation)
async def reset(request: Request):
    """Start a new episode. Body: {"task_id": "..."}  (optional)."""
    try:
        body = await request.body()
        task_id = "task_1_basic_antipatterns"
        if body:
            try:
                data = json.loads(body)
                task_id = data.get("task_id", task_id) or task_id
            except Exception:
                pass
        return env.reset(task_id=task_id)
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc))


@app.post("/step", response_model=StepResult)
def step(action: Action):
    """Submit an optimization action; get real execution feedback."""
    try:
        return env.step(action)
    except RuntimeError as exc:
        raise HTTPException(status_code=400, detail=str(exc))


@app.get("/state", response_model=EnvironmentState)
def state():
    return env.state()


@app.get("/tasks")
def tasks():
    return {"tasks": get_task_list()}


@app.post("/grader")
def grader(action: Action):
    """Grade an action against the current task without advancing the episode."""
    if env._task_data is None:
        raise HTTPException(status_code=400, detail="No active episode. Call /reset first.")
    return grade(env._task_data, action)


@app.post("/baseline")
def baseline():
    """Run the baseline inference script and return output."""
    import subprocess
    try:
        result = subprocess.run(
            ["python", "inference.py"],
            capture_output=True,
            text=True,
            timeout=300,
            cwd=os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
        )
        return {
            "stdout":     result.stdout,
            "stderr":     result.stderr,
            "returncode": result.returncode,
        }
    except Exception as exc:
        raise HTTPException(status_code=500, detail=f"Baseline failed: {exc}")


# ── Unique endpoints (no other team has these) ────────────────────────────

@app.post("/execute", response_model=ExecutionResult)
async def execute(request: Request):
    """
    🚀 UNIQUE ENDPOINT — Execute your optimized query against real DuckDB data.

    Body:
      {
        "task_id": "task_1_basic_antipatterns",
        "optimized_query": "SELECT id, customer_id ... WHERE customer_id = 5000 ..."
      }

    Returns actual execution timing, speedup ratio, result correctness,
    and an EXPLAIN plan — no other OpenEnv environment does this.
    """
    body = await request.body()
    if not body:
        raise HTTPException(status_code=400, detail="Body required: {task_id, optimized_query}")
    try:
        data = json.loads(body)
    except Exception:
        raise HTTPException(status_code=400, detail="Invalid JSON body")

    task_id = data.get("task_id", "task_1_basic_antipatterns")
    optimized_query = (data.get("optimized_query") or "").strip()

    if task_id not in TASKS:
        raise HTTPException(status_code=400, detail=f"Unknown task_id: {task_id}")
    if not optimized_query:
        raise HTTPException(status_code=400, detail="optimized_query is required")

    original_query = TASKS[task_id]["sql_query"]
    ex = get_executor()

    try:
        result = ex.compare(original_query, optimized_query)
        explain = ex.explain(optimized_query)
        return ExecutionResult(
            original_ms=result["original_ms"],
            optimized_ms=result["optimized_ms"],
            speedup=result["speedup"],
            results_match=result["results_match"],
            original_rows=result["original_rows"],
            optimized_rows=result["optimized_rows"],
            original_error=result.get("original_error"),
            optimized_error=result.get("optimized_error"),
            verdict=result["verdict"],
            explain_plan=explain,
        )
    except Exception as exc:
        raise HTTPException(status_code=500, detail=str(exc))


@app.get("/leaderboard")
def leaderboard():
    """
    🏆 UNIQUE ENDPOINT — Real-time leaderboard of best execution scores.

    Shows per-task: best score, best speedup achieved, total attempts,
    how many optimized queries produced correct results.
    """
    return {
        "leaderboard": get_board(),
        "description": (
            "Scores are based on real DuckDB execution: "
            "speedup ratio (35%) + result correctness (20%) + issue detection (25%) + other (20%)"
        ),
    }


# ── Entry point ───────────────────────────────────────────────────────────

def main():
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)


if __name__ == "__main__":
    main()