Delete server/main.py

server/main.py

@@ -1,526 +0,0 @@
-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()