Upload main.py

server/main.py

@@ -0,0 +1,526 @@
+import asyncio
+import csv
+import io
+import os
+import re
+import sqlite3
+from typing import Any
+
+from fastapi import FastAPI, HTTPException, UploadFile, File
+from fastapi.responses import HTMLResponse
+
+from .environment import SQLQualityEnvironment
+from .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()
+
+# Module-level custom database connection (separate from RL environment)
+_custom_db_conn: sqlite3.Connection | None = None
+_custom_db_schema: list[dict] = []
+_custom_db_filename: str = ""
+
+# NEW: RL task support on uploaded database
+_uploaded_db_for_rl: sqlite3.Connection | None = None
+_uploaded_db_schema: dict = {}
+_uploaded_db_tasks: list[str] = []  # Which tasks are compatible
+
+# 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 monitoring and load balancers.
+
+    Returns:
+        dict: Status response indicating the server is running
+    """
+    # Check LLM availability
+    llm_available = False
+    try:
+        import os
+        hf_token = os.getenv("HF_TOKEN", "")
+        api_base_url = os.getenv("API_BASE_URL", "")
+        model_name = os.getenv("MODEL_NAME", "")
+        
+        if hf_token and api_base_url and model_name:
+            llm_available = True
+    except Exception:
+        pass
+    
+    return HealthResponse(status="ok", llm_available=llm_available)
+
+
+@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
+    """
+    global _uploaded_db_for_rl, _uploaded_db_tasks
+    
+    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:
+            # NEW: Check if we have an uploaded DB compatible with this task
+            if _uploaded_db_for_rl is not None and request.task_id in _uploaded_db_tasks:
+                # Use uploaded database for this task
+                return _env.reset_with_connection(
+                    request.task_id, 
+                    _uploaded_db_for_rl,
+                    _uploaded_db_schema
+                )
+            else:
+                # Use synthetic data (original behavior)
+                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))
+
+
+# ─── CUSTOM DATABASE ENDPOINTS ─────────────────────────────────────────
+
+def infer_sqlite_type(values: list[str]) -> str:
+    """Infer SQLite column type from a list of string values."""
+    non_empty = [v for v in values if v and v.strip()]
+    if not non_empty:
+        return 'TEXT'
+    try:
+        for v in non_empty:
+            int(v)
+        return 'INTEGER'
+    except ValueError:
+        pass
+    try:
+        for v in non_empty:
+            float(v)
+        return 'REAL'
+    except ValueError:
+        pass
+    return 'TEXT'
+
+
+def get_db_schema(conn: sqlite3.Connection) -> list[dict]:
+    """Extract schema information from a SQLite connection."""
+    schema = []
+    cursor = conn.execute("SELECT name FROM sqlite_master WHERE type='table'")
+    tables = cursor.fetchall()
+    
+    for (table_name,) in tables:
+        # Get row count
+        count_cursor = conn.execute(f'SELECT COUNT(*) FROM "{table_name}"')
+        row_count = count_cursor.fetchone()[0]
+        
+        # Get columns
+        pragma_cursor = conn.execute(f'PRAGMA table_info("{table_name}")')
+        columns = []
+        for row in pragma_cursor.fetchall():
+            col_name = row[1]
+            col_type = row[2]
+            columns.append(f"{col_name} ({col_type})")
+        
+        schema.append({
+            "table_name": table_name,
+            "row_count": row_count,
+            "columns": columns
+        })
+    
+    return schema
+
+
+def detect_tasks_from_schema(schema: list[dict]) -> list[str]:
+    """Detect which RL tasks are compatible with the uploaded database schema.
+    
+    Args:
+        schema: List of table info dicts with table_name and columns
+        
+    Returns:
+        List of task IDs that can run on this database
+    """
+    compatible_tasks = []
+    table_names = {t['table_name'].lower() for t in schema}
+    
+    # null_audit: requires 'users' table
+    if 'users' in table_names:
+        compatible_tasks.append('null_audit')
+    
+    # type_repair: requires 'transactions' table
+    if 'transactions' in table_names:
+        compatible_tasks.append('type_repair')
+    
+    # dedup_bug: requires 'order_items' table
+    if 'order_items' in table_names:
+        compatible_tasks.append('dedup_bug')
+    
+    # fk_integrity: requires 'invoices' and 'customers' tables
+    if 'invoices' in table_names and 'customers' in table_names:
+        compatible_tasks.append('fk_integrity')
+    
+    # schema_drift: requires 'events_v1' and 'events_v2' tables
+    if 'events_v1' in table_names and 'events_v2' in table_names:
+        compatible_tasks.append('schema_drift')
+    
+    return compatible_tasks
+
+
+@app.post("/upload-db")
+async def upload_db(file: UploadFile = File(...), use_for_tasks: bool = True) -> dict:
+    """Upload a .sqlite, .db, or .csv file for custom exploration and RL tasks.
+    
+    Args:
+        file: The uploaded file (.sqlite, .db, or .csv)
+        use_for_tasks: If True, enables RL tasks on this database
+        
+    Returns:
+        dict: Status, schema info, and compatible RL tasks
+        
+    Raises:
+        HTTPException 400: If file extension is invalid
+        HTTPException 413: If file is too large (>15MB)
+        HTTPException 500: If processing fails
+    """
+    global _custom_db_conn, _custom_db_schema, _custom_db_filename
+    global _uploaded_db_for_rl, _uploaded_db_schema, _uploaded_db_tasks
+    
+    # Validate file extension
+    filename = file.filename or ""
+    ext = os.path.splitext(filename)[1].lower()
+    valid_extensions = {'.sqlite', '.db', '.csv'}
+    
+    if ext not in valid_extensions:
+        raise HTTPException(
+            status_code=400,
+            detail=f"Invalid file extension: {ext}. Must be one of: {valid_extensions}"
+        )
+    
+    # Read file content
+    content = await file.read()
+    
+    # Validate file size (15MB limit)
+    max_size = 15 * 1024 * 1024
+    if len(content) > max_size:
+        raise HTTPException(
+            status_code=413,
+            detail=f"File too large: {len(content)} bytes. Maximum is {max_size} bytes (15MB)"
+        )
+    
+    try:
+        # Close existing connection if any
+        if _custom_db_conn is not None:
+            try:
+                _custom_db_conn.close()
+            except Exception:
+                pass
+            _custom_db_conn = None
+        
+        if ext == '.csv':
+            # Parse CSV and create in-memory SQLite database
+            conn = sqlite3.connect(":memory:", check_same_thread=False)
+            
+            # Read CSV content
+            csv_text = content.decode('utf-8')
+            csv_file = io.StringIO(csv_text)
+            reader = csv.DictReader(csv_file)
+            
+            if not reader.fieldnames:
+                raise HTTPException(status_code=400, detail="CSV file has no headers")
+            
+            # Infer types from first 100 rows
+            rows = list(reader)
+            if not rows:
+                raise HTTPException(status_code=400, detail="CSV file is empty")
+            
+            # Create table
+            table_name = os.path.splitext(os.path.basename(filename))[0]
+            # Sanitize table name
+            table_name = re.sub(r'[^a-zA-Z0-9_]', '_', table_name)
+            if not table_name[0].isalpha() and table_name[0] != '_':
+                table_name = '_' + table_name
+            
+            columns = []
+            for col_name in reader.fieldnames:
+                # Sanitize column name
+                safe_col = re.sub(r'[^a-zA-Z0-9_]', '_', col_name)
+                if not safe_col[0].isalpha() and safe_col[0] != '_':
+                    safe_col = '_' + safe_col
+                
+                values = [row.get(col_name, '') for row in rows[:100]]
+                col_type = infer_sqlite_type(values)
+                columns.append(f'"{safe_col}" {col_type}')
+            
+            cols_joined = ", ".join(columns)
+            create_sql = f'CREATE TABLE "{table_name}" ({cols_joined})'
+            conn.execute(create_sql)
+            
+            # Insert rows
+            placeholders = ', '.join(['?' for _ in reader.fieldnames])
+            insert_sql = f'INSERT INTO "{table_name}" VALUES ({placeholders})'
+            
+            for row in rows:
+                values = [row.get(col_name, None) for col_name in reader.fieldnames]
+                conn.execute(insert_sql, values)
+            
+            conn.commit()
+            _custom_db_conn = conn
+            
+        else:
+            # .sqlite or .db file - write to temp and open
+            import tempfile
+            tmp = tempfile.NamedTemporaryFile(suffix=ext, delete=False)
+            tmp.write(content)
+            tmp.close()
+            
+            conn = sqlite3.connect(tmp.name, check_same_thread=False, uri=True)
+            conn.execute("PRAGMA query_only = OFF")  # Allow modifications for RL tasks
+            _custom_db_conn = conn
+        
+        # Get schema
+        _custom_db_schema = get_db_schema(_custom_db_conn)
+        _custom_db_filename = filename
+        
+        # NEW: Detect compatible RL tasks
+        _uploaded_db_tasks = detect_tasks_from_schema(_custom_db_schema)
+        
+        # NEW: Store for RL task use (copy connection reference)
+        if use_for_tasks and _uploaded_db_tasks:
+            _uploaded_db_for_rl = _custom_db_conn
+            _uploaded_db_schema = {t['table_name']: t for t in _custom_db_schema}
+        
+        return {
+            "status": "ok",
+            "filename": filename,
+            "tables": _custom_db_schema,
+            "compatible_tasks": _uploaded_db_tasks,
+            "use_for_tasks": use_for_tasks and len(_uploaded_db_tasks) > 0
+        }
+        
+    except HTTPException:
+        raise
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Failed to process file: {str(e)}")
+
+
+@app.post("/query-custom")
+async def query_custom(body: dict) -> dict:
+    """Execute a SELECT query against the uploaded custom database.
+    
+    Args:
+        body: JSON body with {"sql": "SELECT ..."}
+        
+    Returns:
+        dict: Query results with rows, columns, and row_count
+        
+    Raises:
+        HTTPException 422: If no database uploaded or SQL is invalid
+    """
+    global _custom_db_conn
+    
+    if _custom_db_conn is None:
+        raise HTTPException(
+            status_code=422,
+            detail="No database uploaded yet. Use /upload-db first."
+        )
+    
+    sql = body.get("sql", "").strip()
+    if not sql:
+        raise HTTPException(status_code=422, detail="No SQL provided")
+    
+    # Validate SQL starts with SELECT (same rule as QUERY action)
+    sql_upper = sql.upper()
+    if not (sql_upper.startswith("SELECT") or sql_upper.startswith("WITH") or 
+            sql_upper.startswith("EXPLAIN") or sql_upper.startswith("PRAGMA")):
+        return {
+            "error": "QUERY action is read-only. This action only accepts SELECT, WITH, EXPLAIN, or PRAGMA.",
+            "message": "SQL Error — check your query. Only SELECT queries are allowed in custom mode."
+        }
+    
+    try:
+        cursor = _custom_db_conn.execute(sql)
+        
+        if cursor.description:
+            columns = [d[0] for d in cursor.description]
+            rows = cursor.fetchall()
+            
+            # Convert rows to dicts
+            row_dicts = []
+            for row in rows:
+                row_dict = {}
+                for i, col in enumerate(columns):
+                    row_dict[col] = row[i]
+                row_dicts.append(row_dict)
+            
+            return {
+                "rows": row_dicts,
+                "columns": columns,
+                "row_count": len(row_dicts)
+            }
+        else:
+            return {
+                "rows": [],
+                "columns": [],
+                "row_count": 0,
+                "message": "Query executed successfully but returned no rows"
+            }
+            
+    except sqlite3.Error as e:
+        return {
+            "error": str(e),
+            "message": "SQL Error — check your query. Check table and column names in the schema panel."
+        }
+
+
+@app.get("/custom-db-status")
+async def custom_db_status() -> dict:
+    """Check if a custom database is loaded and return its schema.
+    
+    Returns:
+        dict: {loaded: true/false, tables: [...]} 
+    """
+    if _custom_db_conn is None:
+        return {"loaded": False}
+    
+    return {
+        "loaded": True,
+        "filename": _custom_db_filename,
+        "tables": _custom_db_schema
+    }
+
+
+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()