Spaces:
Sleeping
Sleeping
import asyncio
import csv
import io
import json
import os
import re
import shutil
import sqlite3
import tempfile
from typing import Any
from difflib import get_close_matches
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
except Exception: # pragma: no cover - optional dependency guard
load_dotenv = None
from .environment import (
SQLQualityEnvironment,
analyze_uploaded_database,
first_sql_keyword,
is_readonly_statement,
quote_identifier,
rewrite_known_quoted_identifiers,
strip_sql_comments,
split_sql_statements,
)
from .models import (
HealthResponse,
ResetRequest,
SQLAction,
SQLObservation,
SQLState,
StepResponse,
)
if load_dotenv:
load_dotenv()
from fastapi import APIRouter
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",
)
# Create /v1/ prefixed router for OpenEnv standard compatibility
v1_router = APIRouter(prefix="/v1")
# Module-level singleton β ONE instance shared across all requests
_env = SQLQualityEnvironment()
# Module-level custom database state (separate from RL environment)
_custom_db_path: str | None = None
_custom_db_schema: list[dict] = []
_custom_db_filename: str = ""
_custom_db_analysis: dict[str, Any] = {}
# NEW: RL task support on uploaded database
_uploaded_db_path: str | 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()
TASK_CATALOG = [
{
"id": "null_audit",
"task_id": "null_audit",
"name": "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",
"task_id": "type_repair",
"name": "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",
"task_id": "dedup_bug",
"name": "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",
"task_id": "fk_integrity",
"name": "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",
"task_id": "schema_drift",
"name": "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."""
conn = sqlite3.connect(db_path, check_same_thread=False)
conn.row_factory = sqlite3.Row
return conn
def replace_temp_db(existing_path: str | None, suffix: str = ".sqlite") -> str:
"""Create a fresh temp database file and remove the prior temp file when possible."""
if existing_path and os.path.exists(existing_path):
try:
os.unlink(existing_path)
except OSError:
pass
handle = tempfile.NamedTemporaryFile(suffix=suffix, delete=False)
handle.close()
return handle.name
def schema_table_map(schema: list[dict]) -> dict[str, dict]:
"""Return schema rows indexed by table name."""
return {table["table_name"]: table for table in schema}
def _first_non_empty(*values: str) -> str:
"""Return the first non-empty stripped string."""
for value in values:
candidate = str(value or "").strip()
if candidate:
return candidate
return ""
def _resolve_llm_runtime() -> dict[str, Any]:
"""Resolve OpenAI-compatible runtime config from environment variables."""
groq_key = _first_non_empty(os.getenv("GROQ_API_KEY", ""))
if groq_key:
return {
"configured": True,
"provider": "groq",
"api_key": groq_key,
"base_url": _first_non_empty(
os.getenv("GROQ_BASE_URL", ""),
os.getenv("API_BASE_URL", ""),
"https://api.groq.com/openai/v1",
),
"model": _first_non_empty(
os.getenv("GROQ_MODEL", ""),
os.getenv("MODEL_NAME", ""),
"llama-3.3-70b-versatile",
),
"missing": [],
}
hf_token = _first_non_empty(os.getenv("HF_TOKEN", ""))
openai_key = _first_non_empty(os.getenv("OPENAI_API_KEY", ""), os.getenv("API_KEY", ""))
provider = "huggingface" if hf_token else ("openai" if openai_key else "none")
api_key = hf_token or openai_key
if provider == "openai":
base_url = _first_non_empty(
os.getenv("OPENAI_BASE_URL", ""),
os.getenv("API_BASE_URL", ""),
"https://api.openai.com/v1",
)
else:
base_url = _first_non_empty(os.getenv("API_BASE_URL", ""), os.getenv("OPENAI_BASE_URL", ""))
model_name = _first_non_empty(os.getenv("MODEL_NAME", ""), os.getenv("OPENAI_MODEL", ""))
missing: list[str] = []
if not api_key:
missing.append("HF_TOKEN or OPENAI_API_KEY")
if not base_url:
missing.append("API_BASE_URL")
if not model_name:
missing.append("MODEL_NAME")
return {
"configured": bool(api_key and base_url and model_name),
"provider": provider,
"api_key": api_key,
"base_url": base_url,
"model": model_name,
"missing": missing,
}
def _format_llm_missing_message(missing: list[str]) -> str:
"""Build a user-friendly missing-config message."""
if not missing:
return "LLM configuration incomplete."
return f"Missing LLM configuration: {', '.join(missing)}"
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>"""
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"]
if not LLM_AVAILABLE:
missing_vars = ", ".join(runtime["missing"])
print(f"[WARNING] LLM env vars not set: {missing_vars}. Server starting in mock mode.", flush=True)
@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")
async def health() -> dict:
"""Health check endpoint for monitoring and load balancers.
Returns:
dict: Health status with server info
"""
runtime = _resolve_llm_runtime()
mock_mode = not bool(runtime["configured"])
llm_available = bool(runtime["configured"])
if mock_mode:
message = _format_llm_missing_message(runtime["missing"])
else:
message = runtime["provider"]
return {
"status": "healthy",
"llm": "mock_mode" if mock_mode else "connected",
"llm_available": llm_available,
"mock_mode": mock_mode,
"provider": runtime["provider"],
"message": message,
}
@app.get("/metadata")
async def get_metadata() -> dict:
"""Get environment metadata for openenv-core.
Returns:
dict: Environment metadata including all task IDs with graders
"""
tasks = _task_catalog_records()
return {
"name": "sql-quality-env",
"version": "0.1.0",
"description": "A multi-step SQL data quality environment for training and evaluating data engineering agents",
"author": "Tharak",
"documentation": DOCUMENTATION_URL,
"documentation_url": DOCUMENTATION_URL,
"tasks": tasks,
"task_count": len(tasks),
"graded_count": sum(1 for t in tasks if t.get("has_grader")),
}
@app.get("/tasks")
async def list_tasks():
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")
async def status() -> dict:
"""Status endpoint with LLM heartbeat check for Hugging Face Spaces.
Performs an actual LLM connection test by checking environment variables
and optionally making a lightweight API call.
Returns:
dict: Status response with llm_connected boolean
"""
runtime = _resolve_llm_runtime()
mock_mode = not bool(runtime["configured"])
llm_connected = bool(runtime["configured"])
llm_error = None
if llm_connected:
try:
import httpx
headers = {"Authorization": f"Bearer {runtime['api_key']}"}
probe_url = f"{runtime['base_url'].rstrip('/')}/models"
async with httpx.AsyncClient() as client:
response = await client.get(probe_url, headers=headers, timeout=5)
if response.status_code == 200:
llm_connected = True
elif response.status_code in {401, 403}:
llm_connected = False
llm_error = f"LLM auth failed (HTTP {response.status_code})"
else:
# Some providers do not expose /models; keep configured-as-connected.
llm_connected = True
llm_error = f"LLM heartbeat warning (HTTP {response.status_code})"
except Exception as e:
# Keep configured-as-connected to avoid false negatives from transient network errors.
llm_connected = True
llm_error = f"LLM heartbeat skipped: {e}"
else:
llm_error = _format_llm_missing_message(runtime["missing"])
return {
"status": "online",
"llm": "mock_mode" if mock_mode else ("connected" if llm_connected else "disconnected"),
"llm_connected": llm_connected,
"llm_error": llm_error,
"mock_mode": mock_mode,
"provider": runtime["provider"],
}
@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("/initialize")
async def initialize(request: ResetRequest) -> dict:
"""Initialize episode and return full state including schema for UI.
This endpoint combines reset() with schema extraction for the frontend.
Args:
request: ResetRequest containing task_id and optional seed
Returns:
dict: Full episode state with schema_info for Table Schema panel
"""
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)
async def reset(request: ResetRequest = None) -> 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
"""
return await _reset_observation(request)
@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
"""
return await _execute_step(action)
@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
"""
return await _current_state()
# βββ V1 PREFIXED ENDPOINTS FOR OPENENV STANDARD βββββββββββββββββββββββββ
@v1_router.get("/health")
async def v1_health() -> dict:
"""Health check endpoint at /v1/health for OpenEnv compatibility."""
return await health()
@v1_router.get("/tasks")
async def v1_list_tasks():
"""List all graded tasks at /v1/tasks for OpenEnv compatibility."""
return _build_tasks_payload()
@v1_router.get("/metadata")
async def v1_get_metadata() -> dict:
"""Get environment metadata at /v1/metadata for OpenEnv compatibility."""
return await get_metadata()
@v1_router.get("/schema")
async def v1_get_schema() -> dict:
"""Return action, observation, and state schemas at /v1/schema."""
return _build_schema_payload()
@v1_router.post("/reset", response_model=SQLObservation)
async def v1_reset(request: ResetRequest) -> SQLObservation:
"""Reset environment at /v1/reset for OpenEnv compatibility."""
return await reset(request)
@v1_router.post("/step", response_model=StepResponse)
async def v1_step(action: SQLAction) -> StepResponse:
"""Execute action at /v1/step for OpenEnv compatibility."""
return await step(action)
@v1_router.get("/state", response_model=SQLState)
async def v1_state() -> SQLState:
"""Get current state at /v1/state for OpenEnv compatibility."""
return await state()
# Include the v1 router in the app
app.include_router(v1_router)
# βββ CUSTOM DATABASE ENDPOINTS βββββββββββββββββββββββββββββββββββββββββ
def infer_sqlite_type(values: list[str]) -> str:
"""Infer SQLite column type from a list of string values with smart casting."""
non_empty = [v for v in values if v and v.strip()]
if not non_empty:
return 'TEXT'
strict_numeric_pattern = re.compile(r'^[+-]?\d+(?:\.\d+)?$')
# Try INTEGER first (with cleaning)
int_count = 0
for v in non_empty:
try:
raw = str(v).strip()
if strict_numeric_pattern.match(raw):
int(raw)
int_count += 1
except (ValueError, TypeError):
pass
# If 80% can be integers, use INTEGER
if int_count / len(non_empty) >= 0.8:
return 'INTEGER'
# Try REAL/float
float_count = 0
for v in non_empty:
try:
raw = str(v).strip()
if strict_numeric_pattern.match(raw):
float(raw)
float_count += 1
except (ValueError, TypeError):
pass
if float_count / len(non_empty) >= 0.8:
return 'REAL'
return 'TEXT'
def sanitize_identifier(name: str) -> str:
"""Sanitize table/column names to prevent SQL reserved word collisions."""
# Replace invalid characters
safe = re.sub(r'[^a-zA-Z0-9_]', '_', name)
# Ensure starts with letter or underscore
if safe and not safe[0].isalpha() and safe[0] != '_':
safe = '_' + safe
return safe
def quote_identifier(name: str) -> str:
"""Quote SQL identifier to handle reserved words and special characters."""
safe = sanitize_identifier(name)
return f'"{safe}"'
def perform_health_audit(conn: sqlite3.Connection, table_name: str) -> dict:
"""Perform automatic health audit on uploaded table.
Returns health metrics including:
- Total row count
- Column-level null/empty statistics
- Data type distribution
- Dirty data indicators (NaN, -1, empty strings)
"""
audit = {
"table_name": table_name,
"total_rows": 0,
"columns": {},
"health_score": 100,
"issues": []
}
try:
# Get row count
cursor = conn.execute(f'SELECT COUNT(*) FROM {quote_identifier(table_name)}')
audit["total_rows"] = cursor.fetchone()[0]
if audit["total_rows"] == 0:
audit["health_score"] = 0
audit["issues"].append("Table is empty")
return audit
# Get column info
pragma_cursor = conn.execute(f'PRAGMA table_info({quote_identifier(table_name)})')
columns = [(row[1], row[2]) for row in pragma_cursor.fetchall()] # (name, type)
total_issues = 0
for col_name, col_type in columns:
col_audit = {
"type": col_type,
"null_count": 0,
"empty_string_count": 0,
"nan_count": 0,
"negative_one_count": 0,
"dirty_data_count": 0,
"health_percentage": 100.0
}
# Count NULLs
try:
cursor = conn.execute(
f'SELECT COUNT(*) FROM {quote_identifier(table_name)} WHERE {quote_identifier(col_name)} IS NULL'
)
col_audit["null_count"] = cursor.fetchone()[0]
except:
pass
# Count empty strings
try:
cursor = conn.execute(
f'''SELECT COUNT(*) FROM {quote_identifier(table_name)}
WHERE {quote_identifier(col_name)} = '' OR {quote_identifier(col_name)} = ' '
OR TRIM(CAST({quote_identifier(col_name)} AS TEXT)) = '''
)
col_audit["empty_string_count"] = cursor.fetchone()[0]
except:
pass
# Count 'NaN', 'N/A', 'null' strings
try:
cursor = conn.execute(
f'''SELECT COUNT(*) FROM {quote_identifier(table_name)}
WHERE UPPER(CAST({quote_identifier(col_name)} AS TEXT)) IN ('NAN', 'N/A', 'NULL', 'NONE', '')'''
)
col_audit["nan_count"] = cursor.fetchone()[0]
except:
pass
# Count -1 values (common placeholder for missing data)
try:
cursor = conn.execute(
f'''SELECT COUNT(*) FROM {quote_identifier(table_name)}
WHERE CAST({quote_identifier(col_name)} AS TEXT) = '-1'
OR {quote_identifier(col_name)} = -1'''
)
col_audit["negative_one_count"] = cursor.fetchone()[0]
except:
pass
# Calculate total dirty data
col_audit["dirty_data_count"] = (
col_audit["null_count"] +
col_audit["empty_string_count"] +
col_audit["nan_count"] +
col_audit["negative_one_count"]
)
# Calculate health percentage
if audit["total_rows"] > 0:
clean_rows = audit["total_rows"] - col_audit["dirty_data_count"]
col_audit["health_percentage"] = round((clean_rows / audit["total_rows"]) * 100, 1)
audit["columns"][col_name] = col_audit
total_issues += col_audit["dirty_data_count"]
# Calculate overall health score
if audit["total_rows"] > 0 and len(columns) > 0:
total_cells = audit["total_rows"] * len(columns)
audit["health_score"] = round(((total_cells - total_issues) / total_cells) * 100, 1)
# Identify worst columns
sorted_cols = sorted(
audit["columns"].items(),
key=lambda x: x[1]["health_percentage"]
)
if sorted_cols and sorted_cols[0][1]["health_percentage"] < 50:
audit["issues"].append(
f"Column '{sorted_cols[0][0]}' has {sorted_cols[0][1]['health_percentage']:.1f}% clean data"
)
except Exception as e:
audit["issues"].append(f"Audit error: {str(e)}")
audit["health_score"] = 0
return audit
def build_safe_null_check(column: str) -> str:
"""Build a comprehensive NULL/dirty data check for a column.
Checks for: NULL, '', ' ', 'NaN', 'N/A', 'NULL', 'NONE', -1
"""
col = quote_identifier(column)
checks = [
f"{col} IS NULL",
f"TRIM(CAST({col} AS TEXT)) = ''",
f"UPPER(CAST({col} AS TEXT)) IN ('NAN', 'N/A', 'NULL', 'NONE')",
f"CAST({col} AS TEXT) = '-1'"
]
return " OR ".join(checks)
def suggest_query_alternatives(conn: sqlite3.Connection, sql: str, error_msg: str) -> list[str]:
"""Suggest alternative queries when a query fails or returns 0 rows."""
suggestions = []
sql_clean = strip_sql_comments(sql)
sql_upper = sql_clean.upper()
schema = get_db_schema(conn)
table_names = [table["table_name"] for table in schema]
column_names = []
for table in schema:
for column in table.get("columns", []):
column_names.append(column.split(" (")[0])
missing_table = re.search(r"no such table: ([^\s]+)", error_msg, re.IGNORECASE)
if missing_table:
bad_table = missing_table.group(1).strip('"')
close = get_close_matches(bad_table, table_names, n=3, cutoff=0.4)
if close:
suggestions.extend([f'Try table "{name}" instead.' for name in close])
missing_column = re.search(r"no such column: ([^\s]+)", error_msg, re.IGNORECASE)
if missing_column:
bad_column = missing_column.group(1).strip('"')
close = get_close_matches(bad_column, column_names, n=3, cutoff=0.4)
if close:
suggestions.extend([f'Try column "{name}" instead.' for name in close])
if "0 rows" in error_msg.lower() or "no rows" in error_msg.lower():
match = re.search(r'FROM\s+("?[\w_]+"\b|[\w_]+)', sql_clean, re.IGNORECASE)
if match:
table = match.group(1).strip('"')
suggestions.append(f'SELECT COUNT(*) FROM "{table}"')
suggestions.append(f'SELECT * FROM "{table}" LIMIT 5')
if "IS NULL" in sql_upper and ("0 rows" in error_msg.lower() or "no rows" in error_msg.lower()):
col_match = re.search(r'WHERE\s+("?[\w_]+"\b|[\w_]+)\s+IS\s+NULL', sql_clean, re.IGNORECASE)
from_match = re.search(r'FROM\s+("?[\w_]+"\b|[\w_]+)', sql_clean, re.IGNORECASE)
if col_match and from_match:
col = col_match.group(1).strip('"')
table = from_match.group(1).strip('"')
suggestions.append(
f"SELECT COUNT(*) FROM {quote_identifier(table)} WHERE "
f"{quote_identifier(col)} IS NULL OR TRIM(CAST({quote_identifier(col)} AS TEXT)) = '' "
f"OR LOWER(TRIM(CAST({quote_identifier(col)} AS TEXT))) IN ('n/a','na','nan','null','none','missing','unknown','-1')"
)
if "syntax error" in error_msg.lower():
suggestions.append("Try running EXPLAIN QUERY PLAN before the query to validate complex SQL.")
return suggestions if suggestions else ["Try: SELECT * FROM table LIMIT 5 to explore data"]
def extract_sql_table_names(sql: str) -> list[str]:
"""Return table names referenced in FROM/JOIN clauses."""
sql_clean = strip_sql_comments(sql)
return [
match.strip('"')
for match in re.findall(r"\b(?:FROM|JOIN)\s+(\"?[\w_]+\"?)", sql_clean, re.IGNORECASE)
]
def build_query_quality_warnings(sql: str, analysis: dict[str, Any] | None) -> list[str]:
"""Generate source-aware warnings for analytically risky SQL."""
if not analysis:
return []
sql_clean = strip_sql_comments(sql)
sql_upper = sql_clean.upper()
warnings: list[str] = []
table_profiles = analysis.get("table_profiles", {})
for table_name, table_profile in table_profiles.items():
for column_name, profile in table_profile.get("column_profiles", {}).items():
quoted_column = f'"{column_name}"'
bare_column = re.escape(column_name)
quoted_pattern = re.escape(quoted_column)
aggregate_patterns = [
rf"\bAVG\s*\(\s*(?:{quoted_pattern}|\b{bare_column}\b)\s*\)",
rf"\bSUM\s*\(\s*(?:{quoted_pattern}|\b{bare_column}\b)\s*\)",
rf"\bMIN\s*\(\s*(?:{quoted_pattern}|\b{bare_column}\b)\s*\)",
rf"\bMAX\s*\(\s*(?:{quoted_pattern}|\b{bare_column}\b)\s*\)",
]
aggregate_hit = any(re.search(pattern, sql_clean, re.IGNORECASE) for pattern in aggregate_patterns)
if profile.get("header_like_count", 0) > 0 and (
re.search(quoted_pattern, sql_clean, re.IGNORECASE)
or re.search(rf"\b{bare_column}\b", sql_clean, re.IGNORECASE)
):
warnings.append(
f'Column "{column_name}" contains {profile["header_like_count"]} header-like value(s); the CSV may include repeated header rows or malformed values.'
)
declared_type = str(profile.get("declared_type", "") or "").upper()
suspicious_name = any(token in column_name.lower() for token in ("salary", "revenue", "price", "amount", "cost", "income"))
numeric_ratio = float(profile.get("numeric_ratio", 0.0) or 0.0)
if aggregate_hit and (numeric_ratio >= 0.7 or suspicious_name or declared_type in {"", "TEXT"}):
if "CAST(" not in sql_upper and "REPLACE(" not in sql_upper:
warnings.append(
f'Column "{column_name}" looks numeric-like text. Direct aggregates may coerce dirty values to 0; clean/cast the values first.'
)
return warnings[:4]
def build_dataset_task_suggestions(sql: str, analysis: dict[str, Any] | None) -> list[str]:
"""Prefer backend task candidates over result-grid heuristics."""
if not analysis:
return []
sql_clean = strip_sql_comments(sql)
referenced_tables = set(extract_sql_table_names(sql_clean))
candidates = []
for task_id, candidate in analysis.get("task_candidates", {}).items():
if not candidate.get("supported"):
continue
related_tables = set(candidate.get("related_tables") or ([candidate.get("table")] if candidate.get("table") else []))
relevance = 1.0 if not referenced_tables else (1.0 if related_tables & referenced_tables else 0.4)
candidates.append(
(
relevance * float(candidate.get("confidence", 0.0) or 0.0),
task_id,
candidate,
)
)
suggestions = []
for _, task_id, candidate in sorted(candidates, key=lambda item: item[0], reverse=True)[:3]:
hint = candidate.get("context_hint") or candidate.get("description") or task_id
suggestions.append(f'{task_id}: {hint}')
return suggestions
def build_source_data_insights(sql: str, analysis: dict[str, Any] | None) -> list[str]:
"""Summarize the most relevant source-data quality signals."""
if not analysis:
return []
insights: list[str] = []
warnings = build_query_quality_warnings(sql, analysis)
insights.extend(warnings)
task_candidates = analysis.get("task_candidates", {})
for task_id in ("type_repair", "null_audit", "dedup_bug", "fk_integrity", "schema_drift"):
candidate = task_candidates.get(task_id, {})
if not candidate.get("supported"):
continue
insights.append(candidate.get("context_hint") or candidate.get("description") or task_id)
if len(insights) >= 4:
break
return insights[:6]
def is_aggregate_like_column_name(column_name: str) -> bool:
"""Heuristic for result columns that are derived aggregates or ranks."""
tokens = [token for token in re.split(r"[^a-z0-9]+", (column_name or "").lower()) if token]
aggregate_tokens = {"avg", "sum", "count", "total", "min", "max", "rank", "highest", "lowest"}
return bool(tokens and any(token in aggregate_tokens for token in tokens))
def is_key_like_column_name(column_name: str) -> bool:
"""Heuristic for key-like result columns where duplicate values are suspicious."""
tokens = [token for token in re.split(r"[^a-z0-9]+", (column_name or "").lower()) if token]
key_tokens = {"id", "email", "phone", "ssn", "uuid", "employee", "patient"}
return bool(tokens and any(token in key_tokens for token in tokens))
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
import logging
import json
from datetime import datetime
from pathlib import Path
# Configure confidence logging
confidence_logger = logging.getLogger('confidence_scoring')
confidence_logger.setLevel(logging.INFO)
# Create logs directory if it doesn't exist
log_dir = Path('logs')
log_dir.mkdir(exist_ok=True)
# File handler for confidence scoring
confidence_handler = logging.FileHandler(log_dir / 'confidence_scoring.log')
confidence_handler.setLevel(logging.INFO)
# Formatter for detailed logging
confidence_formatter = logging.Formatter('%(asctime)s - %(message)s')
confidence_handler.setFormatter(confidence_formatter)
confidence_logger.addHandler(confidence_handler)
import uuid
from typing import Dict
# Session-based table alias mapping for manual task override
table_aliases: Dict[str, Dict[str, str]] = {} # {session_id: {actual_table: alias_as}}
# Session management
active_sessions: Dict[str, dict] = {} # {session_id: {created_at, last_used}}
def log_confidence_calculation(table_name: str, task_id: str, confidence: float,
reasoning: str, fingerprint: dict, session_id: str = None) -> dict:
"""Log confidence score calculations for judge review and UI display."""
log_entry = {
"timestamp": datetime.now().isoformat(),
"session_id": session_id or "default",
"table_name": table_name,
"task_id": task_id,
"confidence_score": round(confidence, 3),
"reasoning": reasoning,
"fingerprint_analysis": {
"row_count": fingerprint.get('row_count', 0),
"column_count": fingerprint.get('column_count', 0),
"null_counts": fingerprint.get('null_counts', {}),
"semantic_hints": fingerprint.get('semantic_hints', []),
"data_types": fingerprint.get('data_types', {}),
"total_nulls": sum(fingerprint.get('null_counts', {}).values()),
"null_percentage": round(
sum(fingerprint.get('null_counts', {}).values()) / max(1, fingerprint.get('row_count', 1)) * 100, 2
)
},
"calculation_breakdown": {
"base_confidence": confidence,
"adjustments": [],
"final_score": round(confidence, 3)
}
}
# Log to file for judge review
confidence_logger.info(f"CONFIDENCE_CALCULATION: {json.dumps(log_entry, indent=2)}")
return log_entry
def sanitize_data_value(value: str) -> any:
"""Sanitize dirty data values during upload."""
if value is None:
return None
value_str = str(value).strip()
# Convert common dirty data patterns to NULL
if (value_str == '' or
value_str == ' ' or
value_str.upper() in ['NAN', 'N/A', 'NULL', 'NONE', 'NA', '#N/A', '-1']):
return None
return value_str
def analyze_table_fingerprint(conn: sqlite3.Connection, table_name: str) -> dict:
"""Analyze table structure and content for semantic task mapping."""
fingerprint = {
"table_name": table_name,
"column_count": 0,
"row_count": 0,
"columns": [],
"null_counts": {},
"data_types": {},
"semantic_hints": []
}
try:
# Get basic info
cursor = conn.execute(f'SELECT COUNT(*) FROM {quote_identifier(table_name)}')
fingerprint["row_count"] = cursor.fetchone()[0]
# Get column info
pragma_cursor = conn.execute(f'PRAGMA table_info({quote_identifier(table_name)})')
columns = [(row[1], row[2]) for row in pragma_cursor.fetchall()]
fingerprint["column_count"] = len(columns)
for col_name, col_type in columns:
fingerprint["columns"].append({"name": col_name, "type": col_type})
# Count nulls
try:
null_cursor = conn.execute(
f'SELECT COUNT(*) FROM {quote_identifier(table_name)} WHERE {quote_identifier(col_name)} IS NULL'
)
null_count = null_cursor.fetchone()[0]
fingerprint["null_counts"][col_name] = null_count
except:
fingerprint["null_counts"][col_name] = 0
# Sample data for type analysis
try:
sample_cursor = conn.execute(
f'SELECT {quote_identifier(col_name)} FROM {quote_identifier(table_name)} WHERE {quote_identifier(col_name)} IS NOT NULL LIMIT 10'
)
samples = [row[0] for row in sample_cursor.fetchall()]
fingerprint["data_types"][col_name] = infer_sqlite_type([str(s) for s in samples])
except:
fingerprint["data_types"][col_name] = 'TEXT'
# Generate semantic hints
col_names_lower = [col["name"].lower() for col in fingerprint["columns"]]
# Check for user-like patterns
user_patterns = ['id', 'name', 'email', 'age', 'city', 'country', 'address', 'phone']
if any(pattern in col_names_lower for pattern in user_patterns):
fingerprint["semantic_hints"].append("user_data")
# Check for transaction-like patterns
trans_patterns = ['amount', 'price', 'total', 'cost', 'quantity', 'order', 'transaction', 'date']
if any(pattern in col_names_lower for pattern in trans_patterns):
fingerprint["semantic_hints"].append("transaction_data")
# Check for job/posting patterns
job_patterns = ['job', 'title', 'company', 'salary', 'rating', 'description', 'skills']
if any(pattern in col_names_lower for pattern in job_patterns):
fingerprint["semantic_hints"].append("job_market_data")
# Check for event patterns
event_patterns = ['event', 'timestamp', 'type', 'category', 'action']
if any(pattern in col_names_lower for pattern in event_patterns):
fingerprint["semantic_hints"].append("event_data")
except Exception as e:
fingerprint["error"] = str(e)
return fingerprint
def detect_tasks_semantically(schema: list[dict], fingerprints: dict[str, dict], session_id: str = None) -> list[dict]:
"""Detect compatible tasks using semantic analysis instead of hard-coded names."""
task_suggestions = []
for table_info in schema:
table_name = table_info['table_name']
fingerprint = fingerprints.get(table_name, {})
semantic_hints = fingerprint.get('semantic_hints', [])
null_counts = fingerprint.get('null_counts', {})
total_nulls = sum(null_counts.values())
# Suggest null audit if there are significant nulls
if total_nulls > 0:
confidence = min(0.9, total_nulls / fingerprint.get('row_count', 1))
reasoning = f"Found {total_nulls} null values across {len(null_counts)} columns"
# Log the calculation for judge visibility
log_entry = log_confidence_calculation(
table_name, 'null_audit', confidence, reasoning, fingerprint, session_id
)
task_suggestions.append({
"task_id": "null_audit",
"table_name": table_name,
"confidence": confidence,
"reason": reasoning,
"requires_alias": table_name.lower() != 'users',
"debug_info": log_entry # AI Reasoning for judges
})
# Suggest type repair if we see suspicious type patterns
data_types = fingerprint.get('data_types', {})
suspicious_numeric_cols = [col for col in ['amount', 'price', 'rating', 'age', 'quantity']
if col in data_types and data_types.get(col, '') == 'TEXT']
if suspicious_numeric_cols:
confidence = 0.7 + (len(suspicious_numeric_cols) * 0.05) # Higher confidence for more issues
reasoning = f"Detected {len(suspicious_numeric_cols)} numeric columns stored as text: {', '.join(suspicious_numeric_cols)}"
# Log the calculation
log_entry = log_confidence_calculation(
table_name, 'type_repair', confidence, reasoning, fingerprint, session_id
)
task_suggestions.append({
"task_id": "type_repair",
"table_name": table_name,
"confidence": min(0.9, confidence),
"reason": reasoning,
"requires_alias": table_name.lower() != 'transactions',
"debug_info": log_entry
})
# Suggest dedup if we have potential duplicate indicators
columns = [col.split(' (')[0].lower() for col in table_info.get('columns', [])]
if 'id' not in columns:
confidence = 0.6
reasoning = "No unique ID column detected - potential duplicates"
# Log the calculation
log_entry = log_confidence_calculation(
table_name, 'dedup_bug', confidence, reasoning, fingerprint, session_id
)
task_suggestions.append({
"task_id": "dedup_bug",
"table_name": table_name,
"confidence": confidence,
"reason": reasoning,
"requires_alias": table_name.lower() != 'order_items',
"debug_info": log_entry
})
# Job market specific task
if 'job_market_data' in semantic_hints:
confidence = 0.8
reasoning = "Detected job market data patterns (job titles, companies, ratings)"
# Log the calculation
log_entry = log_confidence_calculation(
table_name, 'job_market_audit', confidence, reasoning, fingerprint, session_id
)
task_suggestions.append({
"task_id": "job_market_audit",
"table_name": table_name,
"confidence": confidence,
"reason": reasoning,
"requires_alias": True,
"suggested_alias": "users",
"debug_info": log_entry
})
# Sort by confidence
task_suggestions.sort(key=lambda x: x['confidence'], reverse=True)
return task_suggestions
def detect_tasks_from_schema(schema: list[dict]) -> list[str]:
"""Legacy hard-coded mapping - kept for compatibility."""
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
# Add base64 upload endpoint for agent automation
@app.post("/upload-db-base64")
async def upload_db_base64(
filename: str = Form(...),
content_base64: str = Form(...),
use_for_tasks: bool = Form(True),
session_id: str = Form(None)
) -> dict:
"""Upload a database file using base64 encoding for agent automation.
This endpoint allows programmatic file upload without using the OS file picker,
enabling test automation agents to upload files directly.
Args:
filename: Original filename (e.g., "jobs_raw.csv")
content_base64: Base64-encoded file content
use_for_tasks: If True, enables RL tasks on this database
session_id: Optional session identifier for alias isolation
Returns:
dict: Status, schema info, and compatible RL tasks
"""
import base64 as b64
try:
# Decode base64 content
content = b64.b64decode(content_base64)
# Create a mock UploadFile
class MockUploadFile:
def __init__(self, filename, content):
self.filename = filename
self._content = content
mock_file = MockUploadFile(filename, content)
# Reuse the existing upload logic by calling the internal handler
return await _process_upload(mock_file, use_for_tasks, session_id)
except Exception as e:
raise HTTPException(status_code=400, detail=f"Failed to decode base64 content: {str(e)}")
# Internal upload processor (refactored from upload_db)
async def _process_upload(
file,
use_for_tasks: bool = True,
session_id: str = None
) -> dict:
"""Internal upload processing logic shared between upload endpoints."""
global _custom_db_path, _custom_db_schema, _custom_db_filename, _custom_db_analysis
global _uploaded_db_path, _uploaded_db_schema, _uploaded_db_tasks
global table_aliases, active_sessions
# Auto-clear aliases for this session on new upload
if session_id and session_id in table_aliases:
table_aliases[session_id].clear()
active_sessions[session_id] = {
"created_at": active_sessions.get(session_id, {}).get("created_at", datetime.now().isoformat()),
"last_used": datetime.now().isoformat()
}
# 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
if hasattr(file, 'read'):
content = await file.read()
else:
content = file._content
# 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:
# Hard reset any active episode so a new upload cannot inherit stale task state.
old_conn = getattr(_env, "_original_conn", None)
if old_conn is not None:
try:
old_conn.close()
except Exception:
pass
_env._original_conn = None
_env._active = False
_env._done = False
_env._episode_id = None
_env._current_task_id = None
_env._trajectory = []
_env._task_context = {}
_uploaded_db_path = None
_uploaded_db_schema = {}
_uploaded_db_tasks = []
db_path = replace_temp_db(_custom_db_path)
conn = None
if ext == '.csv':
conn = open_sqlite_db(db_path)
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 with data sanitization
placeholders = ', '.join(['?' for _ in reader.fieldnames])
insert_sql = f'INSERT INTO "{table_name}" VALUES ({placeholders})'
for row in rows:
# Sanitize each value during insertion
values = [sanitize_data_value(row.get(col_name, None)) for col_name in reader.fieldnames]
conn.execute(insert_sql, values)
conn.commit()
else:
with open(db_path, "wb") as handle:
handle.write(content)
conn = open_sqlite_db(db_path)
_custom_db_path = db_path
_custom_db_schema = get_db_schema(conn)
_custom_db_filename = filename
_custom_db_analysis = analyze_uploaded_database(conn, schema_table_map(_custom_db_schema))
health_audits = {}
fingerprints = {}
for table_info in _custom_db_schema:
table_name = table_info['table_name']
health_audits[table_name] = perform_health_audit(conn, table_name)
fingerprints[table_name] = analyze_table_fingerprint(conn, table_name)
semantic_tasks = []
for task_id, candidate in _custom_db_analysis.get("task_candidates", {}).items():
semantic_tasks.append({
"task_id": task_id,
"table_name": candidate.get("table"),
"confidence": candidate.get("confidence", 0.0),
"reason": candidate.get("context_hint") or candidate.get("reason", ""),
"supported": candidate.get("supported", False),
"verification_sql": candidate.get("verification_sql"),
})
_uploaded_db_tasks = [
task_id
for task_id, candidate in _custom_db_analysis.get("task_candidates", {}).items()
if candidate.get("supported")
]
if use_for_tasks and _uploaded_db_tasks:
_uploaded_db_path = _custom_db_path
_uploaded_db_schema = schema_table_map(_custom_db_schema)
else:
_uploaded_db_path = None
_uploaded_db_schema = {}
conn.close()
conn = None
return {
"status": "ok",
"filename": filename,
"tables": _custom_db_schema,
"compatible_tasks": _uploaded_db_tasks,
"semantic_tasks": semantic_tasks,
"use_for_tasks": use_for_tasks and len(_uploaded_db_tasks) > 0,
"health_audits": health_audits,
"fingerprints": fingerprints,
"task_contexts": _custom_db_analysis.get("task_candidates", {}),
"session_id": session_id,
"detected_table": _custom_db_schema[0]['table_name'] if _custom_db_schema else None
}
except HTTPException:
raise
except Exception as e:
import traceback
error_detail = f"Failed to process file: {str(e)}\n\nTraceback:\n{traceback.format_exc()}"
raise HTTPException(status_code=500, detail=error_detail)
finally:
if conn is not None:
try:
conn.close()
except Exception:
pass
@app.post("/upload-db")
async def upload_db(
file: UploadFile = File(...),
use_for_tasks: bool = True,
session_id: str = None
) -> 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
session_id: Optional session identifier for alias isolation
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
"""
# Use the shared processing logic
return await _process_upload(file, use_for_tasks, session_id)
@app.post("/set-table-alias")
async def set_table_alias(body: dict, session_id: str = None) -> dict:
"""Set a table alias for manual task override.
Args:
body: JSON body with {"actual_table": "cleaned_ds_jobs2", "alias_as": "users"}
session_id: Optional session identifier for multi-user support
Returns:
dict: Status and current aliases
"""
global table_aliases, active_sessions
# Generate or use provided session ID
if not session_id:
session_id = str(uuid.uuid4())
# Initialize session if not exists
if session_id not in table_aliases:
table_aliases[session_id] = {}
actual_table = body.get("actual_table", "").strip()
alias_as = body.get("alias_as", "").strip()
if not actual_table or not alias_as:
raise HTTPException(status_code=422, detail="Both 'actual_table' and 'alias_as' required")
table_aliases[session_id][actual_table] = alias_as
# Update session activity
active_sessions[session_id] = {
"created_at": datetime.now().isoformat(),
"last_used": datetime.now().isoformat()
}
return {
"status": "ok",
"session_id": session_id,
"message": f"Table '{actual_table}' will be treated as '{alias_as}' for RL tasks",
"aliases": table_aliases[session_id].copy()
}
@app.get("/table-aliases")
async def get_table_aliases(session_id: str = None) -> dict:
"""Get current table aliases for session."""
if not session_id or session_id not in table_aliases:
return {"aliases": {}, "session_id": session_id}
return {
"session_id": session_id,
"aliases": table_aliases[session_id].copy()
}
@app.delete("/table-aliases")
async def clear_table_aliases(session_id: str = None) -> dict:
"""Clear all table aliases for session or all sessions."""
global table_aliases, active_sessions
if session_id and session_id in table_aliases:
table_aliases[session_id].clear()
if session_id in active_sessions:
del active_sessions[session_id]
return {
"status": "ok",
"session_id": session_id,
"message": f"Aliases cleared for session {session_id}"
}
else:
# Clear all sessions
table_aliases.clear()
active_sessions.clear()
return {
"status": "ok",
"message": "All table aliases cleared for all sessions"
}
@app.post("/new-session")
async def create_session() -> dict:
"""Create a new session for isolated alias management."""
session_id = str(uuid.uuid4())
table_aliases[session_id] = {}
active_sessions[session_id] = {
"created_at": datetime.now().isoformat(),
"last_used": datetime.now().isoformat()
}
return {
"session_id": session_id,
"status": "ok",
"message": "New session created"
}
def sanitize_sql_identifiers(sql: str) -> str:
"""Pass-through function kept for compatibility instead of aggressive quoting."""
return sql
def apply_table_aliases(sql: str, session_id: str = None) -> str:
"""Apply table aliases to SQL queries intelligently without mangling complex queries."""
global table_aliases
import re
# Get aliases for session or default to empty
session_aliases = {}
if session_id and session_id in table_aliases:
session_aliases = table_aliases[session_id]
# First sanitize all identifiers (now a pass-through)
sanitized_sql = sanitize_sql_identifiers(sql)
# Then apply aliases
for actual, alias in session_aliases.items():
# Replace unquoted FROM table_name
sanitized_sql = re.sub(
r'\bFROM\s+' + re.escape(actual) + r'\b',
f'FROM "{alias}"',
sanitized_sql,
flags=re.IGNORECASE
)
# Replace quoted FROM "table_name"
sanitized_sql = re.sub(
r'\bFROM\s+"' + re.escape(actual) + r'"\b',
f'FROM "{alias}"',
sanitized_sql,
flags=re.IGNORECASE
)
# Replace back-ticked FROM `table_name`
sanitized_sql = re.sub(
r'\bFROM\s+`' + re.escape(actual) + r'`\b',
f'FROM "{alias}"',
sanitized_sql,
flags=re.IGNORECASE
)
# Replace bracketed FROM [table_name]
sanitized_sql = re.sub(
r'\bFROM\s+\[' + re.escape(actual) + r'\]\b',
f'FROM "{alias}"',
sanitized_sql,
flags=re.IGNORECASE
)
# Replace JOIN table_name
sanitized_sql = re.sub(
r'\bJOIN\s+' + re.escape(actual) + r'\b',
f'JOIN "{alias}"',
sanitized_sql,
flags=re.IGNORECASE
)
sanitized_sql = re.sub(
r'\bJOIN\s+"' + re.escape(actual) + r'"\b',
f'JOIN "{alias}"',
sanitized_sql,
flags=re.IGNORECASE
)
# Replace table.column with alias.column (unquoted)
sanitized_sql = re.sub(
r'\b' + re.escape(actual) + r'\.',
f'"{alias}".',
sanitized_sql,
flags=re.IGNORECASE
)
# Replace "table".column with alias.column (quoted)
sanitized_sql = re.sub(
r'"' + re.escape(actual) + r'"\.',
f'"{alias}".',
sanitized_sql,
flags=re.IGNORECASE
)
return sanitized_sql
@app.post("/query-custom")
async def query_custom(body: dict, session_id: str = None) -> dict:
"""Execute a SELECT query against the uploaded custom database.
Args:
body: JSON body with {"sql": "SELECT ..."}
session_id: Optional session identifier for alias isolation
Returns:
dict: Query results with rows, columns, and row_count, plus suggestions if 0 rows
Raises:
HTTPException 422: If no database uploaded or SQL is invalid
"""
global _custom_db_path
if _custom_db_path 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")
rewritten_sql, rewrites, unresolved = rewrite_known_quoted_identifiers(sql, _custom_db_schema)
if unresolved:
return {
"error": "Unknown quoted identifier detected.",
"message": (
"SQL Error β a double-quoted name is not in the uploaded schema. "
"SQLite may treat it as a string literal and produce misleading results."
),
"suggestions": [
f'Use "{target}" instead of "{source}".'
for source, target in unresolved[:3]
] + ["Use single quotes for string literals and schema-panel names for identifiers."],
"active_tables": [table["table_name"] for table in _custom_db_schema],
}
sql = rewritten_sql
if not is_readonly_statement(sql):
return {
"error": "QUERY action is read-only. This action only accepts SELECT, WITH, EXPLAIN, or PRAGMA.",
"message": "SQL Error β check your query. Only read-only SELECT/CTE/EXPLAIN/PRAGMA queries are allowed in custom mode.",
"suggestions": ["Try: SELECT * FROM table LIMIT 5", "Check schema panel for table names"]
}
statements = split_sql_statements(sql)
if any(not is_readonly_statement(stmt) for stmt in statements):
return {
"error": "Only read-only statements are allowed in Zone 1.",
"message": "SQL Error β mixed read/write statements are not allowed.",
"suggestions": ["Run a single SELECT/CTE/EXPLAIN/PRAGMA statement at a time."]
}
conn = open_sqlite_db(_custom_db_path)
try:
warnings = build_query_quality_warnings(sql, _custom_db_analysis)
if rewrites:
warnings = [
"Auto-mapped quoted identifiers: "
+ ", ".join(f'"{src}" -> "{dst}"' for src, dst in rewrites[:3])
] + warnings
task_suggestions = build_dataset_task_suggestions(sql, _custom_db_analysis)
cursor = None
for stmt in statements:
cursor = conn.execute(stmt)
if cursor is None:
return {
"rows": [],
"columns": [],
"row_count": 0,
"message": "No executable SQL statement was found.",
"suggestions": ["Try: SELECT * FROM table LIMIT 5 to explore data"]
}
if cursor.description:
columns = [d[0] for d in cursor.description]
rows = cursor.fetchall()
row_dicts = [dict(row) for row in rows]
suggestions = None
if len(row_dicts) == 0:
suggestions = suggest_query_alternatives(conn, sql, "0 rows")
return {
"rows": row_dicts,
"columns": columns,
"row_count": len(row_dicts),
"suggestions": suggestions,
"warnings": warnings,
"task_suggestions": task_suggestions,
"source_insights": build_source_data_insights(sql, _custom_db_analysis),
"active_tables": [table["table_name"] for table in _custom_db_schema],
}
return {
"rows": [],
"columns": [],
"row_count": 0,
"message": "Query executed successfully but returned no rows",
"suggestions": ["Try: SELECT * FROM table LIMIT 5 to explore data"],
"warnings": warnings,
"task_suggestions": task_suggestions,
"source_insights": build_source_data_insights(sql, _custom_db_analysis),
"active_tables": [table["table_name"] for table in _custom_db_schema],
}
except sqlite3.Error as e:
error_msg = str(e)
suggestions = suggest_query_alternatives(conn, sql, error_msg)
return {
"error": error_msg,
"message": "SQL Error β check your query. Check table and column names in the schema panel.",
"suggestions": suggestions,
"warnings": build_query_quality_warnings(sql, _custom_db_analysis),
"task_suggestions": build_dataset_task_suggestions(sql, _custom_db_analysis),
"active_tables": [table["table_name"] for table in _custom_db_schema],
}
finally:
conn.close()
@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_path is None:
return {"loaded": False}
return {
"loaded": True,
"filename": _custom_db_filename,
"tables": _custom_db_schema,
"compatible_tasks": _uploaded_db_tasks,
"task_contexts": _custom_db_analysis.get("task_candidates", {}),
}
# βββ AI INSIGHTS ENDPOINT (Groq / OpenAI-compatible) ββββββββββββββββββββββ
def _get_ai_client():
"""Get an OpenAI-compatible client for AI insights.
"""
runtime = _resolve_llm_runtime()
if runtime["configured"]:
try:
from openai import OpenAI
return OpenAI(
api_key=runtime["api_key"],
base_url=runtime["base_url"],
), runtime["model"]
except Exception:
pass
return None, None
def _generate_local_insights(sql: str, rows: list, columns: list, row_count: int) -> list[str]:
"""Generate source-aware analysis without LLM (fast fallback)."""
insights = []
if not rows:
insights.append(f"Query returned 0 rows. The SQL may be too restrictive or target wrong table names.")
insights.extend(build_source_data_insights(sql, _custom_db_analysis))
return insights[:6]
insights.append(f"Query returned {row_count} rows across {len(columns)} columns")
insights.extend(build_query_quality_warnings(sql, _custom_db_analysis))
# Null analysis
null_stats = {}
for col in columns:
null_count = sum(1 for r in rows if r.get(col) is None or r.get(col) == '')
if null_count > 0:
pct = round(null_count / len(rows) * 100, 1)
null_stats[col] = {'count': null_count, 'pct': pct}
if null_stats:
worst = max(null_stats.items(), key=lambda x: x[1]['pct'])
insights.append(f"Highest null rate: column '{worst[0]}' at {worst[1]['pct']}% ({worst[1]['count']} nulls)")
if len(null_stats) > 1:
insights.append(f"{len(null_stats)} columns have null/empty values β consider a null_audit task")
else:
insights.append("No NULL values detected in the returned rows, but source-table issues may still exist outside this result set")
# Numeric analysis
for col in columns:
values = [r.get(col) for r in rows if r.get(col) is not None]
numeric = []
for v in values:
try:
numeric.append(float(str(v).replace('$','').replace(',','').replace('%','')))
except (ValueError, TypeError):
pass
if len(numeric) > len(values) * 0.7 and len(numeric) > 3:
avg = sum(numeric) / len(numeric)
std = (sum((x - avg)**2 for x in numeric) / len(numeric)) ** 0.5
outliers = sum(1 for x in numeric if abs(x - avg) > 2 * std)
insights.append(f"Column '{col}': avg={avg:.2f}, std={std:.2f}, {outliers} outlier(s) (>2Ο)")
# Duplicate detection
for col in columns:
if is_aggregate_like_column_name(col) or not is_key_like_column_name(col):
continue
values = [str(r.get(col, '')) for r in rows if r.get(col) is not None]
unique = set(values)
if len(unique) < len(values) * 0.7 and len(values) > 5:
dup_count = len(values) - len(unique)
insights.append(f"Column '{col}' has {dup_count} duplicate values ({len(unique)}/{len(values)} unique) β check for dedup issues")
for extra in build_dataset_task_suggestions(sql, _custom_db_analysis):
if len(insights) >= 8:
break
insights.append(extra)
return insights[:8]
@app.post("/ai-insights")
async def ai_insights(body: dict) -> dict:
"""Generate AI-powered insights from query results.
Uses Groq API (or HF Inference) for intelligent analysis.
Falls back to local heuristic analysis if no API key configured.
Args:
body: JSON with {sql, rows, row_count, columns}
Returns:
dict: {insights: [str, ...]}
"""
sql = body.get("sql", "")
rows = body.get("rows", [])
row_count = body.get("row_count", len(rows))
columns = body.get("columns", [])
# Try AI-powered insights first
client, model = _get_ai_client()
if client and model:
try:
# Build a compact data summary for the LLM
data_summary = f"SQL: {sql[:300]}\n"
data_summary += f"Result: {row_count} rows, {len(columns)} columns\n"
data_summary += f"Columns: {', '.join(columns[:15])}\n"
if rows:
# Include a few sample rows
sample = rows[:5]
data_summary += f"Sample data (first {len(sample)} rows):\n"
for r in sample:
row_str = ', '.join(f"{k}={v}" for k, v in list(r.items())[:8])
data_summary += f" {row_str}\n"
# Calculate null stats
null_info = []
for col in columns:
null_count = sum(1 for r in rows if r.get(col) is None or r.get(col) == '')
if null_count > 0:
null_info.append(f"{col}: {null_count}/{len(rows)} nulls ({round(null_count/len(rows)*100)}%)")
if null_info:
data_summary += f"Null columns: {'; '.join(null_info)}\n"
import asyncio
response = await asyncio.to_thread(
lambda: client.chat.completions.create(
model=model,
messages=[
{"role": "system", "content": (
"You are a SQL data quality expert. Analyze the query results and provide "
"3-5 actionable insights about data quality issues. Be specific and concise. "
"Focus on: null rates, type mismatches, duplicates, outliers, referential issues. "
"Each insight should be 1 sentence. Use emoji prefixes: π β οΈ β
π π π π§"
)},
{"role": "user", "content": data_summary}
],
temperature=0.3,
max_tokens=400,
stream=False
)
)
text = response.choices[0].message.content or ""
# Split into individual insights
insights = [line.strip() for line in text.strip().split('\n') if line.strip() and len(line.strip()) > 5]
if insights:
return {"insights": insights[:6], "source": "ai"}
except Exception as e:
# Log but don't fail β fall back to local
print(f"AI insights error: {e}")
# Fallback: local heuristic insights
insights = _generate_local_insights(sql, rows, columns, row_count)
return {"insights": insights, "source": "local"}
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()
import asyncio
import csv
import io
import json
import os
import re
import shutil
import sqlite3
import tempfile
from typing import Any
from difflib import get_close_matches
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
except Exception: # pragma: no cover - optional dependency guard
load_dotenv = None
from .environment import (
SQLQualityEnvironment,
analyze_uploaded_database,
first_sql_keyword,
is_readonly_statement,
quote_identifier,
rewrite_known_quoted_identifiers,
strip_sql_comments,
split_sql_statements,
)
from .models import (
HealthResponse,
ResetRequest,
SQLAction,
SQLObservation,
SQLState,
StepResponse,
)
if load_dotenv:
load_dotenv()
from fastapi import APIRouter
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",
)
# Create /v1/ prefixed router for OpenEnv standard compatibility
v1_router = APIRouter(prefix="/v1")
# Module-level singleton β ONE instance shared across all requests
_env = SQLQualityEnvironment()
# Module-level custom database state (separate from RL environment)
_custom_db_path: str | None = None
_custom_db_schema: list[dict] = []
_custom_db_filename: str = ""
_custom_db_analysis: dict[str, Any] = {}
# NEW: RL task support on uploaded database
_uploaded_db_path: str | 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()
TASK_CATALOG = [
{
"id": "null_audit",
"task_id": "null_audit",
"name": "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",
"task_id": "type_repair",
"name": "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",
"task_id": "dedup_bug",
"name": "dedup_bug",