Upload main.py

server/main.py

@@ -0,0 +1,178 @@
+import asyncio
+import os
+
+from fastapi import FastAPI, HTTPException
+from fastapi.responses import HTMLResponse
+
+from server.environment import SQLQualityEnvironment
+from server.models import (
+    HealthResponse,
+    ResetRequest,
+    SQLAction,
+    SQLObservation,
+    SQLState,
+    StepResponse,
+)
+
+app = FastAPI(
+    title="SQL Quality Audit Agent Environment",
+    description="OpenEnv-compatible environment for training AI agents on SQL data quality tasks",
+    version="0.1.0",
+)
+
+# Module-level singleton — ONE instance shared across all requests
+_env = SQLQualityEnvironment()
+
+# HTTP-level request lock to serialize concurrent requests
+_request_lock = asyncio.Lock()
+
+
+def load_static_html() -> str:
+    """Load the static HTML interface for the App section."""
+    # Look for static/index.html in various locations
+    possible_paths = [
+        os.path.join(os.path.dirname(__file__), "..", "static", "index.html"),
+        os.path.join(os.path.dirname(__file__), "static", "index.html"),
+        "/app/static/index.html",
+        "static/index.html",
+    ]
+    for path in possible_paths:
+        if os.path.exists(path):
+            with open(path, "r", encoding="utf-8") as f:
+                return f.read()
+    # Fallback simple HTML if file not found
+    return """<!DOCTYPE html>
+<html>
+<head><title>SQL Quality Environment</title></head>
+<body>
+    <h1>SQL Quality Audit Agent Environment</h1>
+    <p>Environment is running. Use API endpoints directly:</p>
+    <ul>
+        <li>GET /health - Health check</li>
+        <li>POST /reset - Reset task</li>
+        <li>POST /step - Execute action</li>
+        <li>GET /state - Get current state</li>
+    </ul>
+</body>
+</html>"""
+
+
+@app.get("/", response_class=HTMLResponse)
+async def root():
+    """Root endpoint for HF Space App section - provides interactive testing interface."""
+    return load_static_html()
+
+
+@app.get("/health", response_model=HealthResponse)
+async def health() -> HealthResponse:
+    """Health check endpoint for HF Space ping validation.
+
+    Returns:
+        dict: Status response indicating the server is running
+    """
+    return HealthResponse(status="ok")
+
+
+@app.post("/ping", response_model=HealthResponse)
+async def ping() -> HealthResponse:
+    """Ping endpoint for submission validator — accepts empty body.
+
+    Returns:
+        dict: Status response indicating the server is responsive
+    """
+    return HealthResponse(status="ok")
+
+
+@app.post("/reset", response_model=SQLObservation)
+async def reset(request: ResetRequest) -> SQLObservation:
+    """Reset environment and return initial observation for specified task.
+
+    Args:
+        request: ResetRequest containing task_id and optional deterministic seed
+
+    Returns:
+        SQLObservation: Initial state of database for requested task
+
+    Raises:
+        HTTPException 422: If task_id is not a valid task identifier
+        HTTPException 500: If environment initialization fails
+    """
+    async with _request_lock:
+        VALID_TASKS = {
+            "null_audit",
+            "type_repair",
+            "dedup_bug",
+            "fk_integrity",
+            "schema_drift",
+        }
+        if request.task_id not in VALID_TASKS:
+            raise HTTPException(
+                status_code=422,
+                detail=f"Invalid task_id '{request.task_id}'. Must be one of: {VALID_TASKS}",
+            )
+        try:
+            return _env.reset(request.task_id, seed=request.seed)
+        except Exception as e:
+            raise HTTPException(status_code=500, detail=str(e))
+
+
+@app.post("/step", response_model=StepResponse)
+async def step(action: SQLAction) -> StepResponse:
+    """Submit an action and receive observation, reward, done flag, and info.
+
+    Args:
+        action: SQLAction containing task_id and fix_query or answer
+
+    Returns:
+        StepResponse: Updated observation, reward, done status, and info
+
+    Raises:
+        HTTPException 422: If action is invalid or episode is already done
+        HTTPException 500: If environment step execution fails
+    """
+    async with _request_lock:
+        try:
+            return _env.step(action)
+        except HTTPException:
+            raise
+        except ValueError as e:
+            raise HTTPException(status_code=422, detail=str(e))
+        except Exception as e:
+            raise HTTPException(status_code=500, detail=str(e))
+
+
+@app.get("/state", response_model=SQLState)
+async def state() -> SQLState:
+    """Get current episode state.
+
+    Returns:
+        SQLState: Current episode information including task_id, episode_id, attempt_number, cumulative_reward, and done status
+
+    Raises:
+        HTTPException 422: If no active episode exists
+    """
+    try:
+        return _env.state()
+    except ValueError as e:
+        raise HTTPException(status_code=422, detail=str(e))
+
+
+def run_server() -> None:
+    """
+    Entry point for 'uv run server' command.
+
+    Uses PORT environment variable if set, otherwise defaults to 7860
+    (matching the Dockerfile EXPOSE and HF Spaces convention).
+    Override locally: PORT=8000 uv run server
+    """
+    import os
+
+    import uvicorn
+
+    port = int(os.getenv("PORT", "7860"))
+    uvicorn.run("server.main:app", host="0.0.0.0", port=port, reload=False)
+
+
+def main() -> None:
+    """Primary CLI entrypoint expected by OpenEnv validators."""
+    run_server()