server/app.py

"""FastAPI server for the SQL Query Environment.

Exposes HTTP endpoints:
  POST /reset   - Start a new episode
  POST /step    - Submit a SQL query
  GET  /state   - Get current state
  GET  /health  - Health check
"""

import os
import sys
from typing import Optional

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel

# Add parent + server to path for imports
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from models import SQLAction, SQLObservation
from sql_environment import SQLQueryEnvironment


# ── Request/Response models ──

class ResetRequest(BaseModel):
    task_id: Optional[str] = "task_1"


class StepRequest(BaseModel):
    task_id: str = "task_1"
    sql_query: str = ""


# ── Environment instance ──
env = SQLQueryEnvironment()


app = FastAPI(
    title="SQL Query Environment",
    description=(
        "An OpenEnv environment where AI agents learn to write SQL queries. "
        "Features 3 tasks with increasing difficulty and deterministic grading."
    ),
    version="1.0.0",
)


@app.get("/health")
async def health():
    """Health check endpoint."""
    return {"status": "healthy"}


@app.post("/reset")
async def reset(request: ResetRequest = ResetRequest()):
    """Reset the environment and start a new episode."""
    task_id = request.task_id or "task_1"
    obs = env.reset(task_id=task_id)
    return {
        "observation": obs.model_dump(),
        "reward": 0.0,
        "done": False,
        "info": {"episode_started": True, "task_id": task_id},
    }


@app.post("/step")
async def step(request: StepRequest):
    """Submit a SQL query and get the result with grading."""
    if not request.sql_query:
        raise HTTPException(status_code=400, detail="sql_query is required")

    action = SQLAction(task_id=request.task_id, sql_query=request.sql_query)
    obs = env.step(action)

    return {
        "observation": obs.model_dump(),
        "reward": obs.reward,
        "done": obs.done,
        "info": {
            "step_count": obs.step_count,
            "feedback": obs.feedback,
        },
    }


@app.get("/state")
async def get_state():
    """Get current episode state."""
    s = env.state
    return {
        "episode_id": s.episode_id,
        "step_count": s.step_count,
    }


# ── Run with uvicorn ──
if __name__ == "__main__":
    import uvicorn

    port = int(os.environ.get("PORT", 7860))
    uvicorn.run(app, host="0.0.0.0", port=port)