Upload main.py

server/main.py

@@ -9,8 +9,9 @@ import tempfile
 from typing import Any
 from difflib import get_close_matches
 
-from fastapi import FastAPI, HTTPException, UploadFile, File, Form
+from fastapi import FastAPI, HTTPException, UploadFile, File, Form, Request
 from fastapi.responses import HTMLResponse, JSONResponse
+from pydantic import ValidationError
 
 try:
     from dotenv import load_dotenv
@@ -62,6 +63,69 @@ _uploaded_db_tasks: list[str] = []  # Which tasks are compatible
 # HTTP-level request lock to serialize concurrent requests
 _request_lock = asyncio.Lock()
 
+TASK_CATALOG = [
+    {
+        "id": "null_audit",
+        "difficulty": "easy",
+        "description": "Identify the column with null rate exceeding 50 percent",
+        "max_attempts": 1,
+        "grader": True,
+        "has_grader": True,
+        "scored": True,
+        "grader_fn": "server.graders:grade_null_audit",
+    },
+    {
+        "id": "type_repair",
+        "difficulty": "medium",
+        "description": "Fix numeric type mismatch in transaction records",
+        "max_attempts": 3,
+        "grader": True,
+        "has_grader": True,
+        "scored": True,
+        "grader_fn": "server.graders:grade_type_repair",
+    },
+    {
+        "id": "dedup_bug",
+        "difficulty": "hard",
+        "description": "Find and repair duplication bug inflating join volumes",
+        "max_attempts": 5,
+        "grader": True,
+        "has_grader": True,
+        "scored": True,
+        "grader_fn": "server.graders:grade_dedup_bug",
+    },
+    {
+        "id": "fk_integrity",
+        "difficulty": "hard",
+        "description": "Remove orphan child records violating referential integrity",
+        "max_attempts": 4,
+        "grader": True,
+        "has_grader": True,
+        "scored": True,
+        "grader_fn": "server.graders:grade_fk_integrity",
+    },
+    {
+        "id": "schema_drift",
+        "difficulty": "medium",
+        "description": "Normalize drifted schemas into canonical analytics artifact",
+        "max_attempts": 4,
+        "grader": True,
+        "has_grader": True,
+        "scored": True,
+        "grader_fn": "server.graders:grade_schema_drift",
+    },
+]
+TASK_IDS = tuple(task["id"] for task in TASK_CATALOG)
+TASK_CATALOG_BY_ID = {task["id"]: task for task in TASK_CATALOG}
+DEFAULT_TASK_ID = TASK_IDS[0]
+DOCUMENTATION_URL = "https://huggingface.co/spaces/Deadman922/sql-quality-env"
+
+JSONRPC_PARSE_ERROR = -32700
+JSONRPC_INVALID_REQUEST = -32600
+JSONRPC_METHOD_NOT_FOUND = -32601
+JSONRPC_INVALID_PARAMS = -32602
+JSONRPC_INTERNAL_ERROR = -32603
+
 
 def open_sqlite_db(db_path: str) -> sqlite3.Connection:
     """Open a SQLite database with row access tuned for API usage."""
@@ -188,6 +252,240 @@ def load_static_html() -> str:
 </html>"""
 
 
+def _task_catalog_records() -> list[dict[str, Any]]:
+    """Return shallow copies of task records for response payloads."""
+    return [dict(task) for task in TASK_CATALOG]
+
+
+def _valid_task_ids_display() -> list[str]:
+    """Return task IDs in canonical order for errors and schemas."""
+    return list(TASK_IDS)
+
+
+def _validate_task_id_or_raise(task_id: str | None) -> str:
+    """Resolve and validate task IDs against the canonical catalog."""
+    resolved_task_id = task_id or DEFAULT_TASK_ID
+    if resolved_task_id not in TASK_CATALOG_BY_ID:
+        raise HTTPException(
+            status_code=422,
+            detail=(
+                f"Invalid task_id '{resolved_task_id}'. "
+                f"Must be one of: {_valid_task_ids_display()}"
+            ),
+        )
+    return resolved_task_id
+
+
+def _patch_task_id_field_schema(field_schema: dict[str, Any]) -> None:
+    """Inject machine-readable task enums into a single schema field."""
+    if "anyOf" in field_schema:
+        for variant in field_schema.get("anyOf", []):
+            if isinstance(variant, dict) and variant.get("type") == "string":
+                variant["enum"] = _valid_task_ids_display()
+    elif field_schema.get("type") == "string":
+        field_schema["enum"] = _valid_task_ids_display()
+
+
+def _patch_task_id_enums(schema: Any) -> None:
+    """Recursively patch schema objects so task_id is an explicit enum."""
+    if isinstance(schema, dict):
+        properties = schema.get("properties")
+        if isinstance(properties, dict) and isinstance(properties.get("task_id"), dict):
+            _patch_task_id_field_schema(properties["task_id"])
+        for value in schema.values():
+            _patch_task_id_enums(value)
+    elif isinstance(schema, list):
+        for item in schema:
+            _patch_task_id_enums(item)
+
+
+def _patched_model_schema(model_cls: Any) -> dict[str, Any]:
+    """Return a model JSON schema with task enums patched in."""
+    schema = model_cls.model_json_schema()
+    _patch_task_id_enums(schema)
+    return schema
+
+
+def _build_schema_payload() -> dict[str, Any]:
+    """Build OpenEnv-style action/observation/state schema payload."""
+    return {
+        "action": _patched_model_schema(SQLAction),
+        "observation": _patched_model_schema(SQLObservation),
+        "state": _patched_model_schema(SQLState),
+    }
+
+
+def _build_tasks_payload() -> dict[str, Any]:
+    """Build the task listing payload from the canonical catalog."""
+    tasks = _task_catalog_records()
+    return {
+        "tasks": tasks,
+        "total": len(tasks),
+        "graded": sum(1 for task in tasks if task.get("has_grader")),
+    }
+
+
+def _jsonrpc_success(result: Any, request_id: Any = None) -> dict[str, Any]:
+    """Build a JSON-RPC 2.0 success response."""
+    return {
+        "jsonrpc": "2.0",
+        "result": result,
+        "id": request_id,
+    }
+
+
+def _jsonrpc_error(
+    code: int,
+    message: str,
+    request_id: Any = None,
+    data: Any = None,
+) -> dict[str, Any]:
+    """Build a JSON-RPC 2.0 error response."""
+    error: dict[str, Any] = {
+        "code": code,
+        "message": message,
+    }
+    if data is not None:
+        error["data"] = data
+    return {
+        "jsonrpc": "2.0",
+        "error": error,
+        "id": request_id,
+    }
+
+
+def _mcp_tool_specs() -> list[dict[str, Any]]:
+    """Return MCP tool specs for task discovery and episode control."""
+    reset_schema = ResetRequest.model_json_schema()
+    _patch_task_id_enums(reset_schema)
+    action_schema = _patched_model_schema(SQLAction)
+    return [
+        {
+            "name": "list_tasks",
+            "description": "Return the registered task catalog for this environment.",
+            "inputSchema": {
+                "type": "object",
+                "properties": {},
+                "additionalProperties": False,
+            },
+        },
+        {
+            "name": "episode_reset",
+            "description": "Reset the environment for a specific task and return the initial observation.",
+            "inputSchema": reset_schema,
+        },
+        {
+            "name": "episode_step",
+            "description": "Execute one environment step using the SQL action payload.",
+            "inputSchema": action_schema,
+        },
+        {
+            "name": "episode_state",
+            "description": "Return the current state for the active environment episode.",
+            "inputSchema": {
+                "type": "object",
+                "properties": {},
+                "additionalProperties": False,
+            },
+        },
+    ]
+
+
+def _schema_rows_from_observation(obs: SQLObservation) -> list[dict[str, Any]]:
+    """Convert observation schema info to frontend-friendly rows."""
+    schema_rows: list[dict[str, Any]] = []
+    if obs.schema_info:
+        for table_info in obs.schema_info:
+            schema_rows.append(
+                {
+                    "table_name": table_info.table_name,
+                    "columns": table_info.columns,
+                    "row_count": table_info.row_count,
+                }
+            )
+    return schema_rows
+
+
+async def _reset_observation(request: ResetRequest | None = None) -> SQLObservation:
+    """Shared reset helper used by HTTP routes and MCP tools."""
+    global _uploaded_db_path, _uploaded_db_tasks, _custom_db_analysis
+
+    if request is None:
+        request = ResetRequest()
+
+    async with _request_lock:
+        task_id = _validate_task_id_or_raise(request.task_id)
+        conn = None
+        try:
+            if _uploaded_db_path is not None and task_id in _uploaded_db_tasks:
+                conn = open_sqlite_db(_uploaded_db_path)
+                return _env.reset_with_connection(
+                    task_id,
+                    conn,
+                    _uploaded_db_schema,
+                    _custom_db_analysis,
+                )
+            if _uploaded_db_path is not None and _custom_db_analysis.get("task_candidates", {}).get(task_id):
+                reason = _custom_db_analysis["task_candidates"][task_id].get(
+                    "reason",
+                    f"Task '{task_id}' is not supported for the uploaded dataset.",
+                )
+                raise ValueError(reason)
+            return _env.reset(task_id, seed=request.seed)
+        except ValueError as e:
+            raise HTTPException(status_code=422, detail=str(e)) from e
+        except HTTPException:
+            raise
+        except Exception as e:
+            raise HTTPException(status_code=500, detail=str(e)) from e
+        finally:
+            if conn is not None:
+                conn.close()
+
+
+async def _execute_step(action: SQLAction) -> StepResponse:
+    """Shared step helper used by HTTP routes and MCP tools."""
+    async with _request_lock:
+        try:
+            return _env.step(action)
+        except HTTPException:
+            raise
+        except ValueError as e:
+            raise HTTPException(status_code=422, detail=str(e)) from e
+        except Exception as e:
+            raise HTTPException(status_code=500, detail=str(e)) from e
+
+
+async def _current_state() -> SQLState:
+    """Shared state helper used by HTTP routes and MCP tools."""
+    try:
+        return _env.state()
+    except ValueError as e:
+        raise HTTPException(status_code=422, detail=str(e)) from e
+
+
+_original_openapi = app.openapi
+
+
+def _custom_openapi() -> dict[str, Any]:
+    """Patch OpenAPI component schemas so task_id becomes a real enum."""
+    if app.openapi_schema:
+        return app.openapi_schema
+
+    openapi_schema = _original_openapi()
+    if isinstance(openapi_schema, dict):
+        components = openapi_schema.get("components", {}).get("schemas", {})
+        for schema_name in ("ResetRequest", "SQLAction", "SQLObservation", "SQLState"):
+            schema = components.get(schema_name)
+            if isinstance(schema, dict):
+                _patch_task_id_enums(schema)
+        app.openapi_schema = openapi_schema
+    return openapi_schema
+
+
+app.openapi = _custom_openapi
+
+
 # Validate LLM configuration at startup
 runtime = _resolve_llm_runtime()
 LLM_AVAILABLE = runtime["configured"]
@@ -219,7 +517,7 @@ async def health() -> dict:
         message = runtime["provider"]
 
     return {
-        "status": "ok",
+        "status": "healthy",
         "llm": "mock_mode" if mock_mode else "connected",
         "llm_available": llm_available,
         "mock_mode": mock_mode,
@@ -240,45 +538,108 @@ async def get_metadata() -> dict:
         "version": "0.1.0",
         "description": "A multi-step SQL data quality environment for training and evaluating data engineering agents",
         "author": "Tharak",
-        "documentation": "https://huggingface.co/spaces/Deadman922/sql-quality-env",
-        "tasks": [
-            {
-                "id": "null_audit",
-                "has_grader": True
-            },
-            {
-                "id": "type_repair", 
-                "has_grader": True
-            },
-            {
-                "id": "dedup_bug",
-                "has_grader": True
-            },
-            {
-                "id": "fk_integrity",
-                "has_grader": True
-            },
-            {
-                "id": "schema_drift",
-                "has_grader": True
-            }
-        ]
+        "documentation": DOCUMENTATION_URL,
+        "documentation_url": DOCUMENTATION_URL,
+        "tasks": _task_catalog_records(),
     }
 
 
 @app.get("/tasks")
 async def list_tasks():
-    return {
-        "tasks": [
-            {"id": "null_audit", "has_grader": True},
-            {"id": "type_repair", "has_grader": True},
-            {"id": "dedup_bug", "has_grader": True},
-            {"id": "fk_integrity", "has_grader": True},
-            {"id": "schema_drift", "has_grader": True},
-        ],
-        "total": 5,
-        "graded": 5,
-    }
+    return _build_tasks_payload()
+
+
+@app.get("/schema")
+async def get_schema() -> dict:
+    """Return action, observation, and state schemas for OpenEnv clients."""
+    return _build_schema_payload()
+
+
+@app.post("/mcp")
+async def mcp_endpoint(request_raw: Request) -> dict[str, Any]:
+    """Minimal MCP JSON-RPC surface for task discovery and episode control."""
+    try:
+        body = await request_raw.body()
+        request_dict = json.loads(body or b"{}")
+    except json.JSONDecodeError:
+        return _jsonrpc_error(JSONRPC_PARSE_ERROR, "Parse error")
+
+    if not isinstance(request_dict, dict):
+        return _jsonrpc_error(JSONRPC_INVALID_REQUEST, "Invalid Request")
+
+    request_id = request_dict.get("id")
+    if request_dict.get("jsonrpc") != "2.0" or not isinstance(request_dict.get("method"), str):
+        return _jsonrpc_error(JSONRPC_INVALID_REQUEST, "Invalid Request", request_id=request_id)
+
+    method = request_dict["method"]
+    params = request_dict.get("params", {})
+    if params is None:
+        params = {}
+    if not isinstance(params, dict):
+        return _jsonrpc_error(
+            JSONRPC_INVALID_PARAMS,
+            "Invalid params: 'params' must be an object",
+            request_id=request_id,
+        )
+
+    if method == "tools/list":
+        return _jsonrpc_success({"tools": _mcp_tool_specs()}, request_id=request_id)
+
+    if method != "tools/call":
+        return _jsonrpc_error(
+            JSONRPC_METHOD_NOT_FOUND,
+            f"Method not found: {method}",
+            request_id=request_id,
+        )
+
+    tool_name = params.get("name")
+    arguments = params.get("arguments", {})
+    if arguments is None:
+        arguments = {}
+    if not tool_name:
+        return _jsonrpc_error(
+            JSONRPC_INVALID_PARAMS,
+            "Missing 'name' in params",
+            request_id=request_id,
+        )
+    if not isinstance(arguments, dict):
+        return _jsonrpc_error(
+            JSONRPC_INVALID_PARAMS,
+            "Invalid params: 'arguments' must be an object",
+            request_id=request_id,
+        )
+
+    try:
+        if tool_name == "list_tasks":
+            result: Any = _build_tasks_payload()
+        elif tool_name == "episode_reset":
+            reset_request = ResetRequest.model_validate(arguments)
+            result = (await _reset_observation(reset_request)).model_dump()
+        elif tool_name == "episode_step":
+            action = SQLAction.model_validate(arguments)
+            result = (await _execute_step(action)).model_dump()
+        elif tool_name == "episode_state":
+            result = (await _current_state()).model_dump()
+        else:
+            return _jsonrpc_error(
+                JSONRPC_INVALID_PARAMS,
+                f"Tool not found: {tool_name}",
+                request_id=request_id,
+            )
+    except ValidationError as e:
+        return _jsonrpc_error(
+            JSONRPC_INVALID_PARAMS,
+            "Invalid params",
+            request_id=request_id,
+            data=e.errors(),
+        )
+    except HTTPException as e:
+        code = JSONRPC_INVALID_PARAMS if e.status_code == 422 else JSONRPC_INTERNAL_ERROR
+        return _jsonrpc_error(code, str(e.detail), request_id=request_id)
+    except Exception as e:
+        return _jsonrpc_error(JSONRPC_INTERNAL_ERROR, str(e), request_id=request_id)
+
+    return _jsonrpc_success(result, request_id=request_id)
 
 
 @app.get("/status")
@@ -352,67 +713,19 @@ async def initialize(request: ResetRequest) -> dict:
     Returns:
         dict: Full episode state with schema_info for Table Schema panel
     """
-    global _uploaded_db_path, _uploaded_db_tasks, _custom_db_analysis
-    
-    async with _request_lock:
-        VALID_TASKS = {
-            "null_audit",
-            "type_repair",
-            "dedup_bug",
-            "fk_integrity",
-            "schema_drift",
-        }
-        task_id = request.task_id or "null_audit"
-        if task_id not in VALID_TASKS:
-            raise HTTPException(
-                status_code=422,
-                detail=f"Invalid task_id '{task_id}'. Must be one of: {VALID_TASKS}",
-            )
-        conn = None
-        try:
-            if _uploaded_db_path is not None and task_id in _uploaded_db_tasks:
-                conn = open_sqlite_db(_uploaded_db_path)
-                obs = _env.reset_with_connection(
-                    task_id,
-                    conn,
-                    _uploaded_db_schema,
-                    _custom_db_analysis,
-                )
-            elif _uploaded_db_path is not None and _custom_db_analysis.get("task_candidates", {}).get(task_id):
-                reason = _custom_db_analysis["task_candidates"][task_id].get("reason", f"Task '{task_id}' is not supported for the uploaded dataset.")
-                raise ValueError(reason)
-            else:
-                obs = _env.reset(task_id, seed=request.seed)
-        except ValueError as e:
-            raise HTTPException(status_code=422, detail=str(e))
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=str(e))
-        finally:
-            if conn is not None:
-                conn.close()
-            
-            # Extract schema for frontend
-            schema_data = []
-            if obs.schema_info:
-                for table_info in obs.schema_info:
-                    schema_data.append({
-                        "table_name": table_info.table_name,
-                        "columns": table_info.columns,
-                        "row_count": table_info.row_count
-                    })
-            
-            return {
-                "status": "ok",
-                "episode_id": obs.episode_id,
-                "task_id": obs.task_id,
-                "task_description": obs.task_description,
-                "schema": schema_data,
-                "sample_rows": obs.sample_rows,
-                "attempt_number": obs.attempt_number,
-                "max_attempts": obs.max_attempts,
-                "investigation_hint": obs.investigation_hint,
-                "task_context": _env.task_context(),
-            }
+    obs = await _reset_observation(request)
+    return {
+        "status": "ok",
+        "episode_id": obs.episode_id,
+        "task_id": obs.task_id,
+        "task_description": obs.task_description,
+        "schema": _schema_rows_from_observation(obs),
+        "sample_rows": obs.sample_rows,
+        "attempt_number": obs.attempt_number,
+        "max_attempts": obs.max_attempts,
+        "investigation_hint": obs.investigation_hint,
+        "task_context": _env.task_context(),
+    }
 
 
 @app.post("/reset", response_model=SQLObservation)
@@ -429,50 +742,7 @@ async def reset(request: ResetRequest = None) -> SQLObservation:
         HTTPException 422: If task_id is not a valid task identifier
         HTTPException 500: If environment initialization fails
     """
-    global _uploaded_db_path, _uploaded_db_tasks, _custom_db_analysis
-    
-    async with _request_lock:
-        VALID_TASKS = {
-            "null_audit",
-            "type_repair",
-            "dedup_bug",
-            "fk_integrity",
-            "schema_drift",
-        }
-        
-        # Handle None request (validator sends empty body)
-        if request is None:
-            request = ResetRequest()
-        
-        task_id = request.task_id or "null_audit"
-        if task_id not in VALID_TASKS:
-            raise HTTPException(
-                status_code=422,
-                detail=f"Invalid task_id '{task_id}'. Must be one of: {VALID_TASKS}",
-            )
-        conn = None
-        try:
-            if _uploaded_db_path is not None and task_id in _uploaded_db_tasks:
-                conn = open_sqlite_db(_uploaded_db_path)
-                obs = _env.reset_with_connection(
-                    task_id,
-                    conn,
-                    _uploaded_db_schema,
-                    _custom_db_analysis,
-                )
-                return obs
-            elif _uploaded_db_path is not None and _custom_db_analysis.get("task_candidates", {}).get(task_id):
-                reason = _custom_db_analysis["task_candidates"][task_id].get("reason", f"Task '{task_id}' is not supported for the uploaded dataset.")
-                raise ValueError(reason)
-            else:
-                return _env.reset(task_id, seed=request.seed)
-        except ValueError as e:
-            raise HTTPException(status_code=422, detail=str(e))
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=str(e))
-        finally:
-            if conn is not None:
-                conn.close()
+    return await _reset_observation(request)
 
 
 @app.post("/step", response_model=StepResponse)
@@ -489,15 +759,7 @@ async def step(action: SQLAction) -> StepResponse:
         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))
+    return await _execute_step(action)
 
 
 @app.get("/state", response_model=SQLState)
@@ -510,10 +772,7 @@ async def state() -> SQLState:
     Raises:
         HTTPException 422: If no active episode exists
     """
-    try:
-        return _env.state()
-    except ValueError as e:
-        raise HTTPException(status_code=422, detail=str(e))
+    return await _current_state()
 
 
 # ─── CUSTOM DATABASE ENDPOINTS ─────────────────────────────────────────