Upload main.py

server/main.py

@@ -0,0 +1,217 @@
+import asyncio
+
+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()
+
+
+@app.get("/", response_class=HTMLResponse)
+async def root():
+    """Root endpoint for HF Space App section - provides testing interface."""
+    return """
+    <!DOCTYPE html>
+    <html>
+    <head>
+        <title>SQL Quality Environment</title>
+        <style>
+            body { font-family: system-ui, sans-serif; max-width: 800px; margin: 40px auto; padding: 20px; }
+            h1 { color: #2563eb; }
+            .endpoint { background: #f3f4f6; padding: 15px; margin: 10px 0; border-radius: 8px; }
+            .method { font-weight: bold; color: #059669; }
+            code { background: #e5e7eb; padding: 2px 6px; border-radius: 4px; }
+            button { background: #2563eb; color: white; border: none; padding: 10px 20px; 
+                     border-radius: 6px; cursor: pointer; margin: 5px; }
+            button:hover { background: #1d4ed8; }
+            #result { background: #1f2937; color: #10b981; padding: 15px; border-radius: 8px; 
+                     margin-top: 20px; font-family: monospace; white-space: pre-wrap; display: none; }
+        </style>
+    </head>
+    <body>
+        <h1>📊 SQL Quality Audit Agent Environment</h1>
+        <p>OpenEnv-compatible environment for training AI agents on SQL data quality tasks.</p>
+        
+        <h2>Available Endpoints</h2>
+        <div class="endpoint">
+            <span class="method">GET</span> <code>/health</code> - Health check
+        </div>
+        <div class="endpoint">
+            <span class="method">POST</span> <code>/ping</code> - Ping for validator
+        </div>
+        <div class="endpoint">
+            <span class="method">POST</span> <code>/reset</code> - Reset environment (requires task_id)
+        </div>
+        <div class="endpoint">
+            <span class="method">POST</span> <code>/step</code> - Execute action
+        </div>
+        <div class="endpoint">
+            <span class="method">GET</span> <code>/state</code> - Get current state
+        </div>
+        
+        <h2>Quick Test</h2>
+        <button onclick="testHealth()">Test /health</button>
+        <button onclick="testReset()">Test /reset (null_audit)</button>
+        
+        <div id="result"></div>
+        
+        <script>
+            async function testHealth() {
+                const res = await fetch('/health');
+                const data = await res.json();
+                showResult('GET /health\\nStatus: ' + res.status + '\\nResponse: ' + JSON.stringify(data, null, 2));
+            }
+            
+            async function testReset() {
+                const res = await fetch('/reset', {
+                    method: 'POST',
+                    headers: {'Content-Type': 'application/json'},
+                    body: JSON.stringify({task_id: 'null_audit', seed: 42})
+                });
+                const data = await res.json();
+                showResult('POST /reset\\nStatus: ' + res.status + '\\nResponse: ' + JSON.stringify(data, null, 2).slice(0, 500) + '...');
+            }
+            
+            function showResult(text) {
+                const el = document.getElementById('result');
+                el.textContent = text;
+                el.style.display = 'block';
+            }
+        </script>
+    </body>
+    </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()