Spaces:
Sleeping
Sleeping
"""SQL Quality Environment - Core environment manager."""
from __future__ import annotations
from datetime import datetime
from difflib import get_close_matches
from itertools import combinations
import random as _random
import re
import sqlite3
import threading
import uuid
from typing import Any, Dict, List, Tuple
from fastapi import HTTPException
from .models import (
ActionType,
SQLAction,
SQLObservation,
SQLReward,
SQLState,
StepResponse,
TrajectoryStep,
)
from .tasks import (
create_dedup_bug_task,
create_fk_integrity_task,
create_null_audit_task,
create_schema_drift_task,
create_type_repair_task,
)
TASK_CREATORS = {
"null_audit": create_null_audit_task,
"type_repair": create_type_repair_task,
"dedup_bug": create_dedup_bug_task,
"fk_integrity": create_fk_integrity_task,
"schema_drift": create_schema_drift_task,
}
NULLISH_TEXT_VALUES = {
"",
"n/a",
"na",
"nan",
"null",
"none",
"missing",
"unknown",
"-1",
}
ALLOWED_READ_KEYWORDS = {"SELECT", "WITH", "EXPLAIN", "PRAGMA", "SHOW"}
IDENTITY_NAME_TOKENS = {
"id",
"title",
"name",
"email",
"company",
"location",
"address",
"original",
"imdb",
}
BUSINESS_NUMERIC_TOKENS = {
"score",
"rating",
"votes",
"income",
"revenue",
"salary",
"amount",
"price",
"cost",
"duration",
"founded",
}
LOW_SIGNAL_DEDUP_TOKENS = {
"genre",
"industry",
"sector",
"country",
"release",
"year",
"type",
"category",
}
def strip_sql_comments(sql: str) -> str:
"""Remove SQL comments while preserving string literals."""
if not sql:
return ""
result = []
i = 0
n = len(sql)
while i < n:
char = sql[i]
# Handle string literals
if char in ("'", '"'):
quote_char = char
result.append(char)
i += 1
while i < n:
char = sql[i]
result.append(char)
if char == quote_char:
# Check for doubled quote escape
if i + 1 < n and sql[i + 1] == quote_char:
result.append(sql[i + 1]) # Add the escaped quote
i += 2
else:
i += 1
break
elif char == "\\" and i + 1 < n:
# Handle backslash escapes
result.append(sql[i + 1])
i += 2
else:
i += 1
elif char == "-" and i + 1 < n and sql[i + 1] == "-":
# Skip -- comments until end of line
i += 2
while i < n and sql[i] != "\n":
i += 1
elif char == "/" and i + 1 < n and sql[i + 1] == "*":
# Skip /* */ comments
i += 2
while i + 1 < n and not (sql[i] == "*" and sql[i + 1] == "/"):
i += 1
i += 2 # Skip the closing */
else:
result.append(char)
i += 1
return "".join(result).strip()
def split_sql_statements(sql: str) -> list[str]:
"""Split SQL into top-level statements without breaking string literals."""
if not sql:
return []
statements: list[str] = []
current = []
in_string = False
string_char = ""
i = 0
n = len(sql)
while i < n:
char = sql[i]
prev = sql[i - 1] if i > 0 else ""
next_char = sql[i + 1] if i < n - 1 else ""
# Always append quote characters to current buffer
if char in ("'", '"'):
current.append(char)
# Check for doubled-quote escape (SQL standard)
if next_char == char:
current.append(next_char) # Add the escaped quote
i += 1 # Skip the second quote
elif not in_string:
in_string = True
string_char = char
elif string_char == char:
in_string = False
string_char = ""
elif char == "\\" and in_string and next_char:
# Handle backslash escapes in strings
current.append(char)
current.append(next_char)
i += 1
elif char == ";" and not in_string:
# Statement delimiter outside of strings
statement = "".join(current).strip()
if statement:
statements.append(statement)
current = []
else:
current.append(char)
i += 1
# Add the final statement if non-empty
tail = "".join(current).strip()
if tail:
statements.append(tail)
return statements
def first_sql_keyword(sql: str) -> str:
"""Return the first SQL keyword after stripping comments and whitespace."""
cleaned = strip_sql_comments(sql)
if not cleaned:
return ""
return cleaned.split()[0].upper()
def is_readonly_statement(sql: str) -> bool:
"""Return True when a statement is allowed in Zone 1 read-only execution."""
# Remove comments and split into individual statements
statements = split_sql_statements(sql)
# Empty or whitespace-only SQL should return False
if not statements:
return False
# Check every statement's first keyword
for statement in statements:
statement = statement.strip()
if statement: # Skip empty statements
first_keyword = first_sql_keyword(statement)
if not first_keyword or first_keyword not in ALLOWED_READ_KEYWORDS:
return False
return True
def quote_identifier(name: str) -> str:
"""Quote SQLite identifiers safely."""
return '"' + str(name).replace('"', '""') + '"'
def collect_schema_identifiers(schema: Any) -> tuple[set[str], set[str], dict[str, str]]:
"""Collect table/column identifiers plus normalized lookup keys from schema-like objects."""
tables: set[str] = set()
columns: set[str] = set()
normalized_lookup: dict[str, str] = {}
for table_info in schema or []:
if isinstance(table_info, dict):
table_name = table_info.get("table_name", "")
if table_name:
tables.add(table_name)
normalized_lookup.setdefault(normalize_schema_label(table_name), table_name)
# Handle both string and dict column formats
for col in table_info.get("columns", []):
if isinstance(col, str):
column_name = col.split(" (")[0]
else:
column_name = col.get("name", str(col))
if column_name:
columns.add(column_name)
normalized_lookup.setdefault(normalize_schema_label(column_name), column_name)
return tables, columns, normalized_lookup
def rewrite_known_quoted_identifiers(sql: str, schema: Any) -> Tuple[str, List[Tuple[str, str]], List[Tuple[str, str]]]:
"""Map quoted original CSV headers to sanitized schema identifiers when the match is exact after normalization."""
sql_clean = strip_sql_comments(sql)
tables, columns, normalized_lookup = collect_schema_identifiers(schema)
alias_tokens = set(re.findall(r'(?i)\bAS\s+"([^"]+)"', sql_clean))
rewrites: list[tuple[str, str]] = []
unresolved: list[tuple[str, str]] = []
rewritten_sql = sql
for token in re.findall(r'"([^"]+)"', sql_clean):
if token in alias_tokens or token in tables or token in columns:
continue
normalized = normalize_schema_label(token)
if not normalized:
continue
replacement = normalized_lookup.get(normalized)
if replacement and replacement != token:
rewrites.append((token, replacement))
else:
close = get_close_matches(normalized, list(normalized_lookup.keys()), n=1, cutoff=0.75)
if close:
unresolved.append((token, normalized_lookup[close[0]]))
for source, target in rewrites:
# Use regex to match double-quoted identifiers only, not inside string literals
pattern = fr'"({re.escape(source)})"'
replacement = f'"{target}"'
rewritten_sql = re.sub(pattern, replacement, rewritten_sql)
# Deduplicate while preserving order.
dedup_rewrites = list(dict.fromkeys(rewrites))
dedup_unresolved = list(dict.fromkeys(unresolved))
return rewritten_sql, dedup_rewrites, dedup_unresolved
def normalize_text(value: Any) -> str:
"""Normalize values for placeholder, boolean, and date heuristics."""
if value is None:
return ""
return str(value).strip()
def normalize_schema_label(value: str) -> str:
"""Collapse punctuation and spacing for schema-aware comparisons."""
import re
return re.sub(r"[^a-z0-9]+", "", (value or "").lower())
def identifier_tokens(name: str) -> list[str]:
"""Split a column name into normalized semantic tokens."""
import re
return [token for token in re.split(r"[^a-z0-9]+", (name or "").lower()) if token]
def is_identifier_like_name(name: str) -> bool:
"""Return True for classic identifier columns such as *_id."""
tokens = identifier_tokens(name)
return bool(name and ((name or "").lower() == "id" or (tokens and tokens[-1] == "id")))
def is_identity_like_name(name: str) -> bool:
"""Return True for columns likely to participate in business keys."""
tokens = set(identifier_tokens(name))
return bool(tokens & IDENTITY_NAME_TOKENS)
def is_low_signal_dedup_name(name: str) -> bool:
"""Return True for categorical/date-ish names that rarely form a good dedup key alone."""
tokens = set(identifier_tokens(name))
return bool(tokens & LOW_SIGNAL_DEDUP_TOKENS)
def is_nullish_value(value: Any) -> bool:
"""Treat NULLs, empty strings, and common placeholders as missing."""
if value is None:
return True
text = normalize_text(value).lower()
return text in NULLISH_TEXT_VALUES
def clean_numeric_text(value: Any) -> str:
"""Remove common numeric noise such as currency and commas."""
text = normalize_text(value)
return re.sub(r"\s+", "", text.replace(",", "").replace("$", "").replace("%", ""))
def parse_numeric_scalar(value: Any) -> float | None:
"""Parse scalar numeric values while rejecting IDs and textual ranges."""
text = clean_numeric_text(value)
if not text:
return None
if re.fullmatch(r"[+-]?\d+(?:\.\d+)?", text):
try:
return float(text)
except ValueError:
return None
suffix_match = re.fullmatch(r"([+-]?\d+(?:\.\d+)?)([KMB])", text, flags=re.IGNORECASE)
if suffix_match:
magnitude = {"K": 1_000, "M": 1_000_000, "B": 1_000_000_000}[suffix_match.group(2).upper()]
try:
return float(suffix_match.group(1)) * magnitude
except ValueError:
return None
return None
def identifier_like_ratio(values: list[Any]) -> float:
"""Estimate how many values look like alphanumeric identifiers rather than measures."""
clean_values = [normalize_text(v) for v in values if not is_nullish_value(v)]
if not clean_values:
return 0.0
matches = 0
for value in clean_values:
compact = re.sub(r"[\s_-]+", "", value)
if re.fullmatch(r"[A-Za-z]{1,8}\d{3,}", compact):
matches += 1
return matches / max(1, len(clean_values))
def numeric_parse_ratio(values: list[Any]) -> float:
"""Estimate how many non-nullish values can be parsed as numbers."""
clean_values = [v for v in values if not is_nullish_value(v)]
if not clean_values:
return 0.0
parsed = 0
for value in clean_values:
if parse_numeric_scalar(value) is not None:
parsed += 1
return parsed / max(1, len(clean_values))
def boolean_parse_ratio(values: list[Any]) -> float:
"""Estimate how many non-nullish values look like booleans."""
clean_values = [normalize_text(v).lower() for v in values if not is_nullish_value(v)]
if not clean_values:
return 0.0
boolean_tokens = {"true", "false", "yes", "no", "y", "n", "0", "1"}
parsed = sum(1 for value in clean_values if value in boolean_tokens)
return parsed / max(1, len(clean_values))
def date_parse_ratio(values: list[Any]) -> float:
"""Estimate how many non-nullish values look like dates or timestamps."""
clean_values = [normalize_text(v) for v in values if not is_nullish_value(v)]
if not clean_values:
return 0.0
parsed = 0
patterns = [
("%Y-%m-%d",),
("%Y-%m-%d %H:%M:%S",),
("%Y-%m-%dT%H:%M:%S",),
("%m/%d/%Y",),
("%d/%m/%Y",),
]
for value in clean_values:
candidate = value.replace("Z", "")
matched = False
for (fmt,) in patterns:
try:
datetime.strptime(candidate, fmt)
parsed += 1
matched = True
break
except ValueError:
continue
if matched:
continue
try:
datetime.fromisoformat(candidate)
parsed += 1
except ValueError:
pass
return parsed / max(1, len(clean_values))
def normalize_base_table_name(name: str) -> str:
"""Normalize versioned table names to support schema-drift detection."""
text = name.lower()
text = re.sub(r"(?:_v\d+|_version\d+|_ver\d+)$", "", text)
text = re.sub(r"(?:\d{4}|\d+)$", "", text)
return text.rstrip("_")
def fetch_sample_values(conn: sqlite3.Connection, table_name: str, column_name: str, limit: int = 200) -> list[Any]:
"""Fetch a bounded sample for heuristic analysis."""
query = (
f"SELECT {quote_identifier(column_name)} "
f"FROM {quote_identifier(table_name)} "
f"LIMIT {limit}"
)
return [row[0] for row in conn.execute(query).fetchall()]
def nullish_count_sql(table_name: str, column_name: str) -> str:
"""Build a deterministic SQL predicate for nullish values."""
col = quote_identifier(column_name)
return (
f"SELECT COUNT(*) FROM {quote_identifier(table_name)} WHERE "
f"{col} IS NULL OR TRIM(CAST({col} AS TEXT)) = '' OR "
f"LOWER(TRIM(CAST({col} AS TEXT))) IN ('n/a','na','nan','null','none','missing','unknown','-1')"
)
def analyze_uploaded_database(conn: sqlite3.Connection, schema_info: dict[str, dict]) -> dict[str, Any]:
"""Profile uploaded data and derive supported task contexts deterministically."""
table_profiles: dict[str, dict[str, Any]] = {}
task_candidates: dict[str, dict[str, Any]] = {}
tables = list(schema_info.keys())
for table_name in tables:
table_schema = schema_info.get(table_name, {})
raw_columns = table_schema.get("columns", [])
columns = [
col.split(" (")[0] if isinstance(col, str) else col.get("name", "")
for col in raw_columns
]
row_count = int(table_schema.get("row_count", 0) or 0)
column_profiles: dict[str, dict[str, Any]] = {}
duplicate_candidates: list[dict[str, Any]] = []
for column_name in columns:
try:
nullish_count = conn.execute(nullish_count_sql(table_name, column_name)).fetchone()[0]
except sqlite3.Error:
nullish_count = 0
try:
header_like_count = conn.execute(
f"SELECT COUNT(*) FROM {quote_identifier(table_name)} "
f"WHERE LOWER(TRIM(CAST({quote_identifier(column_name)} AS TEXT))) = ?",
(column_name.lower(),),
).fetchone()[0]
except sqlite3.Error:
header_like_count = 0
try:
distinct_count = conn.execute(
f"SELECT COUNT(DISTINCT {quote_identifier(column_name)}) "
f"FROM {quote_identifier(table_name)}"
).fetchone()[0]
except sqlite3.Error:
distinct_count = 0
sample_values = fetch_sample_values(conn, table_name, column_name)
numeric_ratio = numeric_parse_ratio(sample_values)
boolean_ratio = boolean_parse_ratio(sample_values)
date_ratio = date_parse_ratio(sample_values)
id_like_ratio = identifier_like_ratio(sample_values)
numeric_samples = [parsed for parsed in (parse_numeric_scalar(v) for v in sample_values) if parsed is not None]
semantic_anomaly = ""
anomaly_ratio = 0.0
name_tokens = set(identifier_tokens(column_name))
if numeric_samples and name_tokens & {"score", "rating"}:
too_large = sum(1 for value in numeric_samples if value > 10)
in_range = sum(1 for value in numeric_samples if 0 <= value <= 10)
if too_large > 0 and in_range >= max(3, int(len(numeric_samples) * 0.6)):
semantic_anomaly = "score_scale"
anomaly_ratio = round(too_large / max(1, len(numeric_samples)), 4)
declared_type = ""
for raw in raw_columns:
if isinstance(raw, str) and raw.startswith(f"{column_name} "):
declared_type = raw.rsplit("(", 1)[-1].rstrip(")")
break
column_profiles[column_name] = {
"declared_type": declared_type.upper(),
"nullish_count": nullish_count,
"nullish_ratio": round(nullish_count / max(1, row_count), 4),
"distinct_count": distinct_count,
"distinct_ratio": round(distinct_count / max(1, row_count), 4),
"numeric_ratio": round(numeric_ratio, 4),
"boolean_ratio": round(boolean_ratio, 4),
"date_ratio": round(date_ratio, 4),
"identifier_like_ratio": round(id_like_ratio, 4),
"semantic_anomaly": semantic_anomaly,
"anomaly_ratio": anomaly_ratio,
"header_like_count": header_like_count,
}
try:
duplicate_groups = conn.execute(
f"SELECT COUNT(*) FROM ("
f"SELECT {quote_identifier(column_name)} "
f"FROM {quote_identifier(table_name)} "
f"WHERE {quote_identifier(column_name)} IS NOT NULL "
f"GROUP BY {quote_identifier(column_name)} "
f"HAVING COUNT(*) > 1)"
).fetchone()[0]
expected_clean = conn.execute(
f"SELECT COUNT(DISTINCT {quote_identifier(column_name)}) "
f"FROM {quote_identifier(table_name)} "
f"WHERE {quote_identifier(column_name)} IS NOT NULL"
).fetchone()[0]
except sqlite3.Error:
duplicate_groups = 0
expected_clean = row_count
is_key_like = is_identifier_like_name(column_name) or column_name.lower() in {
"id",
"order_id",
"invoice_id",
"customer_id",
"user_id",
"event_id",
"job_id",
}
distinct_ratio = column_profiles[column_name]["distinct_ratio"]
if duplicate_groups >= 2 and (
is_key_like or (is_identity_like_name(column_name) and distinct_ratio >= 0.7)
):
duplicate_candidates.append(
{
"table": table_name,
"key_columns": [column_name],
"duplicate_groups": duplicate_groups,
"expected_clean_count": expected_clean,
"score": (
distinct_ratio
+ (0.45 if is_key_like else 0.0)
+ (0.2 if is_identity_like_name(column_name) else 0.0)
- min(0.3, duplicate_groups / max(1, row_count))
),
}
)
if not duplicate_candidates and len(columns) >= 2:
composite_pool = [
column_name
for column_name in columns
if is_identity_like_name(column_name)
or column_profiles[column_name]["distinct_ratio"] >= 0.55
][:6]
for combo_size in (3, 2):
for combo in combinations(composite_pool, combo_size):
if not any(is_identity_like_name(column_name) for column_name in combo):
continue
if all(is_low_signal_dedup_name(column_name) for column_name in combo):
continue
try:
key_expr = ", ".join(quote_identifier(column_name) for column_name in combo)
duplicate_groups = conn.execute(
f"SELECT COUNT(*) FROM ("
f"SELECT {key_expr} "
f"FROM {quote_identifier(table_name)} "
f"GROUP BY {key_expr} "
f"HAVING COUNT(*) > 1)"
).fetchone()[0]
expected_clean = conn.execute(
f"SELECT COUNT(*) FROM ("
f"SELECT DISTINCT {key_expr} "
f"FROM {quote_identifier(table_name)})"
).fetchone()[0]
except sqlite3.Error:
duplicate_groups = 0
expected_clean = row_count
if duplicate_groups <= 1:
continue
combined_distinct_ratio = expected_clean / max(1, row_count)
if combined_distinct_ratio < 0.6:
continue
semantic_bonus = 0.12 * sum(
1 for column_name in combo if is_identity_like_name(column_name)
)
duplicate_candidates.append(
{
"table": table_name,
"key_columns": list(combo),
"duplicate_groups": duplicate_groups,
"expected_clean_count": expected_clean,
"score": combined_distinct_ratio + semantic_bonus,
}
)
if duplicate_candidates:
break
table_profiles[table_name] = {
"table_name": table_name,
"row_count": row_count,
"columns": columns,
"column_profiles": column_profiles,
"duplicate_candidates": sorted(
duplicate_candidates,
key=lambda item: item["score"],
reverse=True,
),
}
# null_audit candidate
null_candidate = None
for table_profile in table_profiles.values():
for column_name, profile in table_profile["column_profiles"].items():
if profile["nullish_ratio"] < 0.05:
continue
candidate = {
"task_id": "null_audit",
"supported": True,
"confidence": min(0.99, profile["nullish_ratio"] + 0.2),
"table": table_profile["table_name"],
"related_tables": [table_profile["table_name"]],
"column": column_name,
"description": f"Identify the highest nullish-rate column in {table_profile['table_name']}.",
"input_mode": "answer",
"input_label": "Column name",
"input_placeholder": column_name,
"action_guidance": "Enter only the single column name with the highest nullish rate.",
"context_hint": (
f"Target table: {table_profile['table_name']}; focus on column {column_name}. "
f"Nullish ratio: {round(profile['nullish_ratio'] * 100, 1)}%."
),
"verification_sql": (
f"SELECT COUNT(*) AS nullish_count FROM {quote_identifier(table_profile['table_name'])} "
f"WHERE {quote_identifier(column_name)} IS NULL "
f"OR TRIM(CAST({quote_identifier(column_name)} AS TEXT)) = ''"
),
"ground_truth": {
"task_type": "null_audit",
"table": table_profile["table_name"],
"column": column_name,
"null_column": column_name,
"null_rate": profile["nullish_ratio"],
"original_row_count": table_profile["row_count"],
"all_columns": table_profile["columns"],
"uploaded_mode": True,
},
}
if null_candidate is None or candidate["ground_truth"]["null_rate"] > null_candidate["ground_truth"]["null_rate"]:
null_candidate = candidate
task_candidates["null_audit"] = null_candidate or {
"task_id": "null_audit",
"supported": False,
"reason": "No column with meaningful nullish values was detected in the uploaded data.",
}
# type_repair candidate
type_candidate = None
for table_profile in table_profiles.values():
for column_name, profile in table_profile["column_profiles"].items():
declared = profile["declared_type"]
family = None
ratio = 0.0
name_tokens = set(identifier_tokens(column_name))
if is_identifier_like_name(column_name) or profile.get("identifier_like_ratio", 0.0) >= 0.5:
continue
anomaly_type = profile.get("semantic_anomaly", "")
if anomaly_type == "score_scale":
family = "score_scale"
ratio = max(0.8, 1.0 - profile.get("anomaly_ratio", 0.0))
elif declared not in {"INTEGER", "REAL", "NUMERIC", "FLOAT", "DOUBLE"} and profile["numeric_ratio"] >= 0.75:
family = "numeric"
ratio = profile["numeric_ratio"]
elif declared not in {"BOOLEAN", "INTEGER"} and profile["boolean_ratio"] >= 0.85:
family = "boolean"
ratio = profile["boolean_ratio"]
elif declared not in {"DATE", "DATETIME", "TIMESTAMP"} and profile["date_ratio"] >= 0.75:
family = "date"
ratio = profile["date_ratio"]
if not family:
continue
if family == "date" and ratio >= 0.98:
# Fully parseable date-like text does not need repair unless another anomaly exists.
continue
score = ratio
if family in {"numeric", "score_scale"}:
score += 0.15
if name_tokens & BUSINESS_NUMERIC_TOKENS:
score += 0.2
if name_tokens & {"score", "rating"} and family == "score_scale":
score += 0.25
if declared in {"", "TEXT"}:
score += 0.05
candidate = {
"task_id": "type_repair",
"supported": True,
"confidence": min(0.99, score),
"ranking_score": score,
"table": table_profile["table_name"],
"related_tables": [table_profile["table_name"]],
"column": column_name,
"description": (
f"Repair numeric anomalies in {table_profile['table_name']}.{column_name}."
if family == "score_scale"
else f"Repair typed values stored inconsistently in {table_profile['table_name']}.{column_name}."
),
"input_mode": "fix_sql",
"input_label": "Fix SQL",
"input_placeholder": f'UPDATE {quote_identifier(table_profile["table_name"])} SET {quote_identifier(column_name)} = ...',
"action_guidance": "Submit SQL that repairs the affected values on a cloned database.",
"context_hint": (
f"Target table: {table_profile['table_name']}; column {column_name} contains score/rating values above 10 that likely need rescaling."
if family == "score_scale"
else f"Target table: {table_profile['table_name']}; column {column_name} looks {family}-like "
f"but is declared as {declared or 'TEXT'}."
),
"verification_sql": (
f"SELECT {quote_identifier(column_name)}, typeof({quote_identifier(column_name)}) "
f"FROM {quote_identifier(table_profile['table_name'])} "
f"WHERE {quote_identifier(column_name)} IS NOT NULL LIMIT 10"
),
"ground_truth": {
"task_type": "type_repair",
"table": table_profile["table_name"],
"column": column_name,
"target_column": column_name,
"issue_type": family,
"baseline_parse_ratio": ratio,
"baseline_anomaly_ratio": profile.get("anomaly_ratio", 0.0),
"declared_type": declared or "TEXT",
"original_row_count": table_profile["row_count"],
"all_columns": table_profile["columns"],
"uploaded_mode": True,
},
}
if type_candidate is None or candidate["ranking_score"] > type_candidate.get("ranking_score", type_candidate["confidence"]):
type_candidate = candidate
task_candidates["type_repair"] = type_candidate or {
"task_id": "type_repair",
"supported": False,
"reason": "No strongly typed-as-text column was detected in the uploaded data.",
}
# dedup_bug candidate
dedup_candidate = None
for table_profile in table_profiles.values():
if not table_profile["duplicate_candidates"]:
continue
best = table_profile["duplicate_candidates"][0]
key_columns = best["key_columns"]
candidate = {
"task_id": "dedup_bug",
"supported": True,
"confidence": 0.85,
"table": table_profile["table_name"],
"related_tables": [table_profile["table_name"]],
"column": key_columns[0],
"key_columns": key_columns,
"description": f"Repair duplicate rows in {table_profile['table_name']}.",
"input_mode": "fix_sql",
"input_label": "Fix SQL",
"input_placeholder": f'DELETE FROM {quote_identifier(table_profile["table_name"])} WHERE rowid NOT IN (...)',
"action_guidance": "Submit SQL that removes duplicate rows while preserving one valid row per duplicate key.",
"context_hint": f"Target table: {table_profile['table_name']}; duplicate key candidate: {', '.join(best.get('key_columns', ['unknown']))}",
"ranking_score": best["score"],
"ground_truth": {
"task_type": "dedup_bug",
"table": table_profile["table_name"],
"key_columns": key_columns,
"column": key_columns[0] if key_columns else None,
"expected_clean_count": table_profile["row_count"] - best.get("duplicate_count", 0),
"original_row_count": table_profile["row_count"],
"uploaded_mode": True,
}
}
if dedup_candidate is None or best["score"] > dedup_candidate.get("ranking_score", 0):
dedup_candidate = candidate
task_candidates["dedup_bug"] = dedup_candidate or {
"task_id": "dedup_bug",
"supported": False,
"reason": "No meaningful duplicate-key pattern was detected in the uploaded data.",
}
# fk_integrity candidate
fk_candidate = None
for child_table, child_profile in table_profiles.items():
for column_name in child_profile["columns"]:
if not column_name.lower().endswith("_id") or column_name.lower() == "id":
continue
stem = column_name[:-3].lower()
for parent_table, parent_profile in table_profiles.items():
if parent_table == child_table:
continue
parent_keys = [col for col in parent_profile["columns"] if col.lower() == "id" or col.lower() == stem or col.lower() == f"{stem}_id"]
if not parent_keys:
continue
parent_key = parent_keys[0]
try:
orphan_count = conn.execute(
f"SELECT COUNT(*) FROM {quote_identifier(child_table)} c "
f"LEFT JOIN {quote_identifier(parent_table)} p "
f"ON c.{quote_identifier(column_name)} = p.{quote_identifier(parent_key)} "
f"WHERE c.{quote_identifier(column_name)} IS NOT NULL "
f"AND p.{quote_identifier(parent_key)} IS NULL"
).fetchone()[0]
valid_count = conn.execute(
f"SELECT COUNT(*) FROM {quote_identifier(child_table)} c "
f"JOIN {quote_identifier(parent_table)} p "
f"ON c.{quote_identifier(column_name)} = p.{quote_identifier(parent_key)}"
).fetchone()[0]
except sqlite3.Error:
orphan_count = 0
valid_count = 0
if orphan_count <= 0 or valid_count <= 0:
continue
candidate = {
"task_id": "fk_integrity",
"supported": True,
"confidence": min(0.95, 0.5 + orphan_count / max(1, child_profile["row_count"])),
"table": child_table,
"related_tables": [child_table, parent_table],
"column": column_name,
"description": f"Resolve orphaned rows in {child_table}.{column_name}.",
"input_mode": "fix_sql",
"input_label": "Fix SQL",
"input_placeholder": f'UPDATE {quote_identifier(child_table)} SET {quote_identifier(column_name)} = ... WHERE ...',
"action_guidance": "Submit SQL that resolves orphan child rows without damaging valid relationships.",
"context_hint": (
f"Child table: {child_table}; parent table: {parent_table}; "
f"foreign key candidate: {column_name} -> {parent_key}."
),
"verification_sql": (
f"SELECT COUNT(*) AS orphan_count FROM {quote_identifier(child_table)} c "
f"LEFT JOIN {quote_identifier(parent_table)} p "
f"ON c.{quote_identifier(column_name)} = p.{quote_identifier(parent_key)} "
f"WHERE c.{quote_identifier(column_name)} IS NOT NULL "
f"AND p.{quote_identifier(parent_key)} IS NULL"
),
"ground_truth": {
"task_type": "fk_integrity",
"table": child_table,
"column": column_name,
"parent_table": parent_table,
"parent_key": parent_key,
"expected_clean_count": valid_count,
"orphan_count": orphan_count,
"original_row_count": child_profile["row_count"],
"all_columns": child_profile["columns"],
"uploaded_mode": True,
},
}
if fk_candidate is None or candidate["ground_truth"]["orphan_count"] > fk_candidate["ground_truth"]["orphan_count"]:
fk_candidate = candidate
task_candidates["fk_integrity"] = fk_candidate or {
"task_id": "fk_integrity",
"supported": False,
"reason": "No real parent-child orphan pattern was detected across the uploaded tables.",
}
# schema_drift candidate
drift_candidate = None
for left_table, right_table in combinations(tables, 2):
left_base = normalize_base_table_name(left_table)
right_base = normalize_base_table_name(right_table)
left_cols = set(table_profiles[left_table]["columns"])
right_cols = set(table_profiles[right_table]["columns"])
overlap = len(left_cols & right_cols)
if left_base != right_base and overlap < 2:
continue
if overlap == 0:
continue
canonical_name = "canonical_events" if "event" in left_base or "event" in right_base else f"canonical_{left_base or right_base}"
expected_columns = sorted(left_cols | right_cols)
candidate = {
"task_id": "schema_drift",
"supported": True,
"confidence": min(0.9, overlap / max(1, min(len(left_cols), len(right_cols))) + 0.2),
"table": canonical_name,
"related_tables": [left_table, right_table],
"column": expected_columns[0] if expected_columns else "",
"description": f"Create a canonical table from {left_table} and {right_table}.",
"input_mode": "fix_sql",
"input_label": "Fix SQL",
"input_placeholder": f'CREATE TABLE {quote_identifier(canonical_name)} AS SELECT ...',
"action_guidance": "Submit SQL that creates a canonical table covering both source tables.",
"context_hint": (
f"Comparable source tables detected: {left_table} and {right_table}. "
f"Canonical output should be {canonical_name}."
),
"verification_sql": f"SELECT COUNT(*) FROM {quote_identifier(canonical_name)}",
"ground_truth": {
"task_type": "schema_drift",
"table": canonical_name,
"canonical_table": canonical_name,
"source_tables": [left_table, right_table],
"expected_columns": expected_columns,
"expected_row_count": table_profiles[left_table]["row_count"] + table_profiles[right_table]["row_count"],
"original_row_count": table_profiles[left_table]["row_count"] + table_profiles[right_table]["row_count"],
"uploaded_mode": True,
},
}
drift_candidate = candidate
break
task_candidates["schema_drift"] = drift_candidate or {
"task_id": "schema_drift",
"supported": False,
"reason": "No comparable versioned tables were detected in the uploaded database.",
}
return {
"table_profiles": table_profiles,
"task_candidates": task_candidates,
}
class SQLQualityEnvironment:
"""Main environment class for SQL Quality tasks."""
def __init__(self) -> None:
"""Initialize environment."""
self._current_task_id: str | None = None
self._episode_id: str | None = None
self._attempt_number: int = 1
self._step_count: int = (
0 # increments on every action (QUERY, PROFILE, FIX, SUBMIT)
)
self._max_steps_per_episode: int = 20 # prevent infinite loops
self._cumulative_reward: float = 0.0
self._done: bool = False
self._ground_truth: Dict[str, Any] = {}
self._cached_observation: SQLObservation | None = None
self._lock = threading.Lock()
# NEW: Multi-step trajectory support
self._trajectory: list = []
self._original_snapshot: dict = {}
self._original_conn: sqlite3.Connection | None = None
self._uploaded_analysis: dict[str, Any] = {}
self._task_context: dict[str, Any] = {}
self._active: bool = False
self._last_feedback: str | None = None
def reset(self, task_id: str, seed: int | None = None) -> SQLObservation:
"""Reset environment for a new episode.
Args:
task_id: One of the registered task IDs
seed: Optional deterministic episode seed
Returns:
SQLObservation for the new episode
Raises:
ValueError: If task_id is not valid
"""
with self._lock:
# ── NEW: close any live connection from the previous episode ──────
if hasattr(self, "_original_conn") and self._original_conn is not None:
try:
self._original_conn.close()
except Exception:
pass # Already closed or never opened — safe to ignore
self._original_conn = None
# Validate task_id
if task_id not in TASK_CREATORS:
raise ValueError(f"Invalid task_id: {task_id}")
# Generate new episode
self._episode_id = str(uuid.uuid4())
self._attempt_number = 1
self._step_count = (
0 # increments on every action (QUERY, PROFILE, FIX, SUBMIT)
)
self._cumulative_reward = 0.0
self._done = False
self._current_task_id = task_id
self._trajectory = []
self._original_snapshot = {}
self._ground_truth = {}
self._task_context = {}
self._cached_observation = None
self._last_feedback = None
# Create task-specific database to extract observation data
conn = None
try:
# Generate random episode seed for diversity
episode_seed = seed if seed is not None else _random.randint(0, 2**31 - 1)
creator = TASK_CREATORS[task_id]
conn, observation, self._ground_truth = creator(episode_seed)
# Store original connection for graders (read-only use)
self._original_conn = conn
# Snapshot the target table for data preservation grading
target_table = self._ground_truth.get("table", "")
if target_table:
from .graders import _snapshot_table
self._original_snapshot = _snapshot_table(conn, target_table)
# Mark episode as active
self._active = True
# Cache observation for reuse in step()
self._cached_observation = observation.model_copy(
update={"episode_id": self._episode_id}
)
# Store schema info for dynamic updates
self._schema_info = observation.schema_info
# Keep connection open for multi-step operations
# Connection will be closed when episode ends or environment is reset
except Exception as e:
if conn is not None:
conn.close()
raise e
return self._cached_observation
def reset_with_connection(
self,
task_id: str,
conn: sqlite3.Connection,
schema_info: dict,
uploaded_analysis: dict[str, Any] | None = None,
) -> SQLObservation:
"""Reset environment using an uploaded database connection.
This allows RL tasks to run on user-uploaded databases instead of
synthetic data.
Args:
task_id: One of the registered task IDs
conn: SQLite connection to the uploaded database
schema_info: Dictionary mapping table names to schema info
Returns:
SQLObservation for the new episode
Raises:
ValueError: If task_id is not valid
"""
with self._lock:
# ── close any live connection from the previous episode ──────
if hasattr(self, "_original_conn") and self._original_conn is not None:
try:
self._original_conn.close()
except Exception:
pass
self._original_conn = None
# Validate task_id
if task_id not in TASK_CREATORS:
raise ValueError(f"Invalid task_id: {task_id}")
# Generate new episode
self._episode_id = str(uuid.uuid4())
self._attempt_number = 1
self._step_count = 0
self._cumulative_reward = 0.0
self._done = False
self._current_task_id = task_id
self._trajectory = []
self._original_snapshot = {}
self._ground_truth = {}
self._task_context = {}
self._cached_observation = None
self._last_feedback = None
try:
# Use the uploaded connection
self._original_conn = conn
self._uploaded_analysis = uploaded_analysis or analyze_uploaded_database(conn, schema_info)
task_candidates = self._uploaded_analysis.get("task_candidates", {})
task_candidate = task_candidates.get(task_id, {})
if not task_candidate.get("supported"):
reason = task_candidate.get(
"reason",
f"Task '{task_id}' is not supported for the uploaded database.",
)
raise ValueError(reason)
self._ground_truth = task_candidate.get("ground_truth", {}).copy()
self._task_context = {
"task_id": task_id,
"target_table": task_candidate.get("table"),
"target_column": task_candidate.get("column"),
"related_tables": task_candidate.get("related_tables", []),
"key_columns": task_candidate.get("key_columns", []),
"verification_sql": task_candidate.get("verification_sql"),
"context_hint": task_candidate.get("context_hint"),
"confidence": task_candidate.get("confidence"),
"input_mode": task_candidate.get("input_mode", "fix_sql"),
"input_label": task_candidate.get("input_label", "Fix SQL"),
"input_placeholder": task_candidate.get("input_placeholder"),
"action_guidance": task_candidate.get("action_guidance"),
"description": task_candidate.get("description", f"Investigate task {task_id} on uploaded data."),
"max_attempts": task_candidate.get("max_attempts", 1),
"investigation_hint": task_candidate.get("context_hint", f"Using uploaded database."),
}
target_table = self._ground_truth.get("table", "")
# Snapshot the target table for data preservation grading
if target_table:
from .graders import _snapshot_table
self._original_snapshot = _snapshot_table(conn, target_table)
# Build observation from uploaded schema
from .models import SchemaInfo, SQLObservation
schema_list = []
sample_rows = {}
related_tables = task_candidate.get("related_tables") or ([target_table] if target_table else [])
for related_table in related_tables:
if related_table not in schema_info:
continue
table_data = schema_info[related_table]
schema_list.append(SchemaInfo(
table_name=related_table,
columns=table_data.get('columns', []),
row_count=table_data.get('row_count', 0)
))
try:
cursor = conn.execute(f'SELECT * FROM "{related_table}" LIMIT 3')
cols = [d[0] for d in cursor.description]
rows = cursor.fetchall()
sample_rows[related_table] = [dict(zip(cols, row)) for row in rows]
except Exception:
sample_rows[related_table] = []
observation = SQLObservation(
task_id=task_id,
task_description=task_candidate.get("description", f"Investigate task {task_id} on uploaded data."),
schema_info=schema_list,
sample_rows=sample_rows,
attempt_number=1,
episode_id=self._episode_id,
trajectory_history=[],
investigation_hint=task_candidate.get("context_hint", f"Using uploaded database: {target_table} table"),
max_attempts=1,
)
# Mark episode as active
self._active = True
self._cached_observation = observation
self._schema_info = schema_list
except Exception as e:
raise e
return self._cached_observation
def task_context(self) -> dict[str, Any]:
"""Expose the current uploaded-data task context for UI helpers."""
return self._task_context.copy()
def _analyze_data_for_ground_truth(
self,
conn: sqlite3.Connection,
task_id: str,
target_table: str
) -> dict:
"""Analyze uploaded database to determine ground truth for scoring.
Args:
conn: SQLite connection
task_id: The task type
target_table: The primary table for this task
Returns:
Dictionary with ground truth information
"""
ground_truth = {"task_type": task_id, "table": target_table}
try:
if task_id == "null_audit" and target_table:
# Find column with highest null rate
cursor = conn.execute(f'PRAGMA table_info("{target_table}")')
columns = [row[1] for row in cursor.fetchall()]
total_rows = conn.execute(
f'SELECT COUNT(*) FROM "{target_table}"'
).fetchone()[0]
if total_rows > 0:
max_null_rate = 0
max_null_col = None
for col in columns:
try:
null_count = conn.execute(
f'SELECT COUNT(*) FROM "{target_table}" WHERE "{col}" IS NULL'
).fetchone()[0]
null_rate = null_count / total_rows
if null_rate > max_null_rate:
max_null_rate = null_rate
max_null_col = col
except Exception:
continue
ground_truth["column"] = max_null_col
ground_truth["null_rate"] = max_null_rate
ground_truth["null_column"] = max_null_col
elif task_id == "dedup_bug" and target_table:
# Find duplicate keys
cursor = conn.execute(f'PRAGMA table_info("{target_table}")')
columns = [row[1] for row in cursor.fetchall()]
ground_truth["all_columns"] = columns
if columns:
# Try to find duplicates on first column
key_col = columns[0]
ground_truth["column"] = key_col
ground_truth["offending_key"] = key_col
# Count total duplicates
cursor = conn.execute(f'''
SELECT COUNT(*) FROM (
SELECT "{key_col}"
FROM "{target_table}"
GROUP BY "{key_col}"
HAVING COUNT(*) > 1
)
''')
dup_groups = cursor.fetchone()[0]
# Count expected rows after dedup (unique keys)
cursor = conn.execute(f'''
SELECT COUNT(DISTINCT "{key_col}") FROM "{target_table}"
''')
expected_clean = cursor.fetchone()[0]
ground_truth["expected_clean_count"] = expected_clean
# Get first duplicate example
cursor = conn.execute(f'''
SELECT "{key_col}", COUNT(*) as cnt
FROM "{target_table}"
GROUP BY "{key_col}"
HAVING cnt > 1
LIMIT 1
''')
result = cursor.fetchone()
if result:
ground_truth["duplicate_key"] = result[0]
ground_truth["duplicate_count"] = result[1]
ground_truth["duplicate_groups"] = dup_groups
elif task_id == "type_repair" and target_table:
# Find columns that might have type issues (TEXT containing numbers)
cursor = conn.execute(f'PRAGMA table_info("{target_table}")')
columns = [row[1] for row in cursor.fetchall()]
ground_truth["all_columns"] = columns
# Find columns with numeric-like text values
for col in columns:
try:
# Check if column contains numeric-looking values stored as TEXT
cursor = conn.execute(f'''
SELECT "{col}" FROM "{target_table}"
WHERE "{col}" IS NOT NULL
LIMIT 100
''')
values = [row[0] for row in cursor.fetchall()]
# Check if majority look like numbers but stored as text
numeric_like = 0
for v in values:
if v and str(v).strip():
try:
float(str(v).replace(',', '').replace('$', ''))
numeric_like += 1
except (ValueError, TypeError):
pass
if numeric_like > len(values) * 0.5:
ground_truth["column"] = col
ground_truth["target_column"] = col
ground_truth["issue_type"] = "TEXT_STORED_NUMERIC"
break
except Exception:
continue
# If no type issue found, default to first column
if "column" not in ground_truth and columns:
ground_truth["column"] = columns[0]
ground_truth["target_column"] = columns[0]
elif task_id == "fk_integrity" and target_table:
# For FK integrity, we need to simulate orphan records
# Since uploaded data may not have parent/child tables,
# we'll create a virtual parent table based on existing data
cursor = conn.execute(f'PRAGMA table_info("{target_table}")')
columns = [row[1] for row in cursor.fetchall()]
ground_truth["all_columns"] = columns
# Look for columns that might be foreign keys (ending in _id)
fk_candidates = [c for c in columns if c.lower().endswith('_id') and c.lower() != 'id']
if fk_candidates:
fk_col = fk_candidates[0]
ground_truth["column"] = fk_col
# Find values that would be orphans (using modulo to simulate missing parents)
cursor = conn.execute(f'''
SELECT DISTINCT "{fk_col}" FROM "{target_table}"
WHERE "{fk_col}" IS NOT NULL
''')
all_ids = [row[0] for row in cursor.fetchall()]
if all_ids:
# Mark some IDs as invalid to simulate orphans
rng = _random.Random(42) # Deterministic for testing
valid_ids = rng.sample(all_ids, max(1, len(all_ids) * 2 // 3))
valid_ids_set = set(valid_ids)
invalid_ids = [v for v in all_ids if v not in valid_ids_set]
ground_truth["parent_table"] = f"{target_table}_parent"
ground_truth["parent_key"] = fk_col.replace('_id', '')
ground_truth["valid_invoice_ids"] = valid_ids
ground_truth["expected_clean_count"] = len([v for v in all_ids if v in valid_ids])
ground_truth["orphan_count"] = len(invalid_ids)
elif task_id == "schema_drift" and target_table:
# Schema drift compares two versions of same table
# For uploaded data, we'll treat it as canonical and create expected schema
cursor = conn.execute(f'PRAGMA table_info("{target_table}")')
columns_info = cursor.fetchall()
columns = [row[1] for row in columns_info]
ground_truth["table"] = target_table
ground_truth["canonical_table"] = target_table
ground_truth["expected_columns"] = columns
# Get row count
cursor = conn.execute(f'SELECT COUNT(*) FROM "{target_table}"')
ground_truth["expected_row_count"] = cursor.fetchone()[0]
# Get distinct IDs if available
if 'id' in [c.lower() for c in columns]:
id_col = next(c for c in columns if c.lower() == 'id')
cursor = conn.execute(f'SELECT COUNT(DISTINCT "{id_col}") FROM "{target_table}"')
ground_truth["expected_distinct_ids"] = cursor.fetchone()[0]
except Exception as e:
# Ground truth analysis failed, but we can still run the task
ground_truth["analysis_error"] = str(e)
return ground_truth
def step(self, action: SQLAction) -> StepResponse:
"""
Execute one agent action and return the resulting observation and reward.
Routes by action_type:
QUERY - execute SELECT on episode conn, return result in observation
PROFILE - compute column statistics, return in observation
FIX - execute on fresh_test_db(), score partially, continue episode
SUBMIT - call compute_episode_score(), end episode
"""
with self._lock:
if not self._active:
raise HTTPException(
status_code=422, detail="No active episode. Call reset() first."
)
if self._done:
raise HTTPException(
status_code=422,
detail="Episode is already done. Call reset() to start a new episode.",
)
# Check max steps guard
if self._step_count >= self._max_steps_per_episode:
raise HTTPException(
status_code=422,
detail="Max investigation steps reached. Call SUBMIT or reset().",
)
action = self._normalize_action(action)
action_type = action.action_type
# --- QUERY action ---
if action_type == ActionType.QUERY:
return self._handle_query(action)
# --- PROFILE action ---
elif action_type == ActionType.PROFILE:
return self._handle_profile(action)
# --- FIX action ---
elif action_type == ActionType.FIX:
return self._handle_fix(action)
# --- SUBMIT action (or legacy single-step) ---
else:
return self._handle_submit(action)
def _normalize_action(self, action: SQLAction) -> SQLAction:
"""Infer v2 action_type for backward-compatible payloads without action_type."""
if "action_type" in action.model_fields_set:
return action
inferred = ActionType.SUBMIT
if action.column_name:
inferred = ActionType.PROFILE
elif action.fix_query:
inferred = ActionType.FIX
elif action.sql:
inferred = ActionType.QUERY if is_readonly_statement(action.sql) else ActionType.FIX
elif action.answer or action.root_cause_table or action.offending_key:
inferred = ActionType.SUBMIT
return action.model_copy(update={"action_type": inferred})
def _handle_query(self, action: SQLAction) -> StepResponse:
"""
Handle a QUERY action.
Executes the SQL on the episode connection (supports multi-statement).
Returns a result preview in the observation. Does NOT end the episode.
Awards small step_reward (0.05) for queries that target the right area.
"""
# Minimal implementation to isolate the issue
sql = (action.sql or "").strip()
# Execute the query directly
try:
cursor = self._original_conn.execute(sql)
if cursor.description:
rows = cursor.fetchmany(5)
cols = [d[0] for d in cursor.description]
result_preview = {
"type": "table",
"columns": cols,
"rows": [dict(zip(cols, row)) for row in rows],
"total_rows": len(rows),
}
else:
result_preview = f"Query executed successfully. Rows affected: {cursor.rowcount}"
except Exception as e:
result_preview = f"SQL Error: {e}"
# Add to trajectory
self._trajectory.append(
TrajectoryStep(
step_number=len(self._trajectory) + 1,
action_type=ActionType.QUERY,
sql_submitted=sql,
result_preview=str(result_preview)[:400],
step_reward=0.01,
)
)
# Build minimal observation
from .models import SQLObservation
obs = SQLObservation(
task_id=getattr(self, '_current_task_id', 'unknown'),
episode_id=getattr(self, '_episode_id', 'unknown'),
task_description="Test task",
schema_info=getattr(self, '_schema_info', []),
sample_rows={},
attempt_number=getattr(self, '_attempt_number', 1),
max_attempts=1,
previous_feedback=None,
trajectory_history=[],
available_actions=["query", "profile", "fix", "submit"],
investigation_hint=None,
)
return StepResponse(
observation=obs,
reward=SQLReward(
score=self._cumulative_reward,
step_reward=0.01,
feedback=str(result_preview)[:400],
dimensions={},
),
done=False,
info={"action_type": "query", "step": len(self._trajectory)},
)
def _handle_profile(self, action: SQLAction) -> StepResponse:
"""
Handle a PROFILE action — statistical summary of a named column.
Computes: null rate, distinct count, data type distribution, min/max.
Does NOT end the episode. Awards step_reward based on column relevance.
"""
col = (action.column_name or "").strip()
result_preview = ""
if not col:
result_preview = "Error: column_name required for PROFILE action."
else:
# Try to find which table has this column
try:
tables = self._original_conn.execute(
"SELECT name FROM sqlite_master WHERE type='table'"
).fetchall()
profile_results = []
for (tbl,) in tables:
try:
cols = [
row[1]
for row in self._original_conn.execute(
f"PRAGMA table_info({tbl})"
).fetchall()
]
if col in cols:
total = self._original_conn.execute(
f"SELECT COUNT(*) FROM {quote_identifier(tbl)}"
).fetchone()[0]
nulls = self._original_conn.execute(
nullish_count_sql(tbl, col)
).fetchone()[0]
distinct = self._original_conn.execute(
f"SELECT COUNT(DISTINCT {quote_identifier(col)}) FROM {quote_identifier(tbl)}"
).fetchone()[0]
profile_results.append(
f"{tbl}.{col}: total={total}, nulls={nulls} ({100*nulls//total if total else 0}%), distinct={distinct}"
)
except sqlite3.Error:
pass
result_preview = (
" | ".join(profile_results)
if profile_results
else f"Column '{col}' not found."
)
except sqlite3.Error as e:
result_preview = f"Profile error: {e}"
# Determine step_reward based on relevance to ground truth
ground_table = self._ground_truth.get("table", "").lower()
ground_col = self._ground_truth.get(
"column", self._ground_truth.get("null_column", "")
).lower()
col_lower = col.lower()
if col_lower == ground_col:
step_reward = 0.08 # On-target: profiling exact relevant column
elif ground_table and ground_table in result_preview.lower():
step_reward = 0.05 # Near-target: profiled something from right table
elif col_lower in result_preview.lower() and len(col_lower) > 2:
step_reward = 0.01 # Off-target: profiled something real but irrelevant
else:
step_reward = 0.01 # Minimal: invalid column or junk probe
self._trajectory.append(
TrajectoryStep(
step_number=len(self._trajectory) + 1,
action_type=ActionType.PROFILE,
sql_submitted=f"PROFILE({col})",
result_preview=str(result_preview)[:200] if result_preview else "",
step_reward=step_reward,
)
)
self._cumulative_reward = min(
1.0, round(self._cumulative_reward + step_reward, 4)
)
self._last_feedback = str(result_preview)[:200] if result_preview else ""
self._step_count += 1
obs = self._build_observation(result_preview=result_preview)
return StepResponse(
observation=obs,
reward=SQLReward(
score=self._cumulative_reward,
step_reward=step_reward,
feedback=str(result_preview)[:200] if result_preview else "",
),
done=False,
info={"action_type": "profile", "column_profiled": col},
)
def _handle_fix(self, action: SQLAction) -> StepResponse:
"""
Handle a FIX action — execute write query on isolated copy for preview scoring.
Executes on a _fresh_test_db() copy. Does NOT mutate episode state.
Returns intermediate feedback on whether fix appears to work.
Does NOT end episode — agent can follow up with verification query or SUBMIT.
Awards intermediate reward based on quick fix assessment.
"""
fix_sql = action.effective_sql or ""
step_reward = 0.0
feedback = ""
rewrite_note = ""
if fix_sql:
if getattr(self, "_schema_info", None):
# fix_sql, rewrites, unresolved = rewrite_known_quoted_identifiers(fix_sql, self._schema_info)
fix_sql, rewrites, unresolved = fix_sql, [], []
if unresolved:
feedback = (
"Fix preview blocked: unknown quoted identifier(s). Closest schema matches: "
+ ", ".join(f'"{src}" -> "{dst}"' for src, dst in unresolved[:3])
)
fix_sql = ""
if rewrites:
rewrite_note = " Auto-mapped: " + ", ".join(
f'"{src}" -> "{dst}"' for src, dst in rewrites[:3]
) + "."
from .graders import grade_fix_correctness
if fix_sql:
quick_score, severity = grade_fix_correctness(
fix_sql,
self._original_conn,
self._original_snapshot,
self._ground_truth,
)
step_reward = quick_score * 0.3 # Partial credit for fix attempt
preview_label = "Answer preview" if self._ground_truth.get("task_type") == "null_audit" else "Fix preview"
feedback = f"{preview_label}: {severity} ({quick_score:.2f}). Submit to finalize score.{rewrite_note}"
else:
feedback = "No SQL provided in fix action."
self._trajectory.append(
TrajectoryStep(
step_number=len(self._trajectory) + 1,
action_type=ActionType.FIX,
sql_submitted=fix_sql,
result_preview=feedback[:200],
step_reward=step_reward,
)
)
self._cumulative_reward = min(
1.0, round(self._cumulative_reward + step_reward, 4)
)
self._last_feedback = feedback
self._step_count += 1
obs = self._build_observation(result_preview=feedback)
return StepResponse(
observation=obs,
reward=SQLReward(
score=self._cumulative_reward,
step_reward=step_reward,
feedback=feedback,
),
done=False,
info={
"action_type": "fix",
"severity_preview": severity if fix_sql else "N/A",
},
)
def _handle_submit(self, action: SQLAction) -> StepResponse:
"""
Handle a SUBMIT action — compute final composite score and end episode.
Calls compute_episode_score() with the full trajectory.
Sets done=True. This is the only action that ends the episode.
"""
from .graders import compute_episode_score
reward = compute_episode_score(
trajectory=self._trajectory,
ground_truth=self._ground_truth,
original_conn=self._original_conn,
original_snapshot=self._original_snapshot,
submit_action=action,
)
self._done = True
self._active = False
self._cumulative_reward = min(
1.0, round(self._cumulative_reward + reward.score, 4)
)
self._last_feedback = reward.feedback
obs = self._build_observation(result_preview="")
# Close connection when episode ends
if self._original_conn:
self._original_conn.close()
self._original_conn = None
return StepResponse(
observation=obs,
reward=reward,
done=True,
info={
"action_type": "submit",
"final_score": reward.score,
"dimensions": reward.dimensions,
}
)
if fix_sql:
# quick_score, severity = grade_fix_correctness(
# fix_sql,
# self._original_conn,
# self._original_snapshot,
# self._ground_truth,
# )
quick_score, severity = 0.5, "UNKNOWN"
step_reward = quick_score * 0.3 # Partial credit for fix attempt
preview_label = "Answer preview" if self._ground_truth.get("task_type") == "null_audit" else "Fix preview"
feedback = f"{preview_label}: {severity} ({quick_score:.2f}). Submit to finalize score.{rewrite_note}"
else:
feedback = "No SQL provided in fix action."
self._trajectory.append(
TrajectoryStep(
step_number=len(self._trajectory) + 1,
action_type=ActionType.FIX,
sql_submitted=fix_sql,
result_preview=feedback[:200],
step_reward=step_reward,
)
)
self._cumulative_reward = min(
1.0, round(self._cumulative_reward + step_reward, 4)
)
self._last_feedback = feedback[:200] if feedback else ""
self._step_count += 1
# Build observation
obs = self._build_observation(result_preview=feedback)
return StepResponse(
observation=obs,
reward=SQLReward(
score=self._cumulative_reward,
step_reward=step_reward,
feedback=feedback[:200] if feedback else "",
dimensions={},
),
done=False,
info={"action_type": "fix", "step": len(self._trajectory)},
)
def _build_observation(self, result_preview: str = "") -> SQLObservation:
"""Build observation with current state and optional result preview."""
# Format trajectory history for frontend
trajectory_history = []
for step in self._trajectory:
trajectory_history.append({
"step_number": step.step_number,
"action_type": step.action_type.value,
"sql_submitted": step.sql_submitted,
"result_preview": step.result_preview,
"step_reward": step.step_reward,
})
# Fallback to _cached_observation for synthetic tasks where _task_context may be empty
task_description = self._task_context.get("description") or (self._cached_observation.task_description if self._cached_observation else "")
max_attempts = self._task_context.get("max_attempts") or (self._cached_observation.max_attempts if self._cached_observation else 1)
investigation_hint = self._task_context.get("investigation_hint") or (self._cached_observation.investigation_hint if self._cached_observation else None)
# Check if current task has a grader (all tasks currently have graders)
# This makes available_actions dynamic based on task grader status
tasks_with_graders = {"null_audit", "type_repair", "dedup_bug", "fk_integrity", "schema_drift"}
has_grader = self._current_task_id in tasks_with_graders
available_actions = ["query", "profile", "fix"]
if has_grader:
available_actions.append("submit")
return SQLObservation(
task_id=self._current_task_id,
episode_id=self._episode_id,
task_description=task_description or "",
schema_info=self._schema_info,
sample_rows=self._cached_observation.sample_rows if self._cached_observation else {},
attempt_number=self._attempt_number,
max_attempts=max_attempts or 1,
previous_feedback=self._last_feedback,
trajectory_history=trajectory_history,
available_actions=available_actions,
investigation_hint=investigation_hint,
)
def state(self) -> SQLState:
"""Get current environment state.
Returns:
SQLState with current episode information
Raises:
ValueError: If no episode has been started
"""
if self._episode_id is None:
raise ValueError("No episode started. Call reset() first.")
return SQLState(
task_id=self._current_task_id,
episode_id=self._episode_id,
attempt_number=self._attempt_number,
cumulative_reward=self._cumulative_reward,
done=self._done,
)
"""SQL Quality Environment - Core environment manager."""
from __future__ import annotations
from datetime import datetime
from difflib import get_close_matches
from itertools import combinations
import random as _random
import re
import sqlite3
import threading
import uuid
from typing import Any, Dict, List, Tuple
from fastapi import HTTPException
from .models import (
ActionType,
SQLAction,
SQLObservation,
SQLReward,
SQLState,
StepResponse,
TrajectoryStep,
)
from .tasks import (
create_dedup_bug_task,
create_fk_integrity_task,
create_null_audit_task,
create_schema_drift_task,
create_type_repair_task,
)
TASK_CREATORS = {
"null_audit": create_null_audit_task,
"type_repair": create_type_repair_task,
"dedup_bug": create_dedup_bug_task,
"fk_integrity": create_fk_integrity_task,
"schema_drift": create_schema_drift_task,
}
NULLISH_TEXT_VALUES = {
"",
"n/a",
"na",
"nan",
"null",
"none",
"missing",
"unknown",
"-1",
}
ALLOWED_READ_KEYWORDS = {"SELECT", "WITH", "EXPLAIN", "PRAGMA", "SHOW"}
IDENTITY_NAME_TOKENS = {
"id",
"title",
"name",
"email",
"company",
"location",
"address",
"original",
"imdb",
}
BUSINESS_NUMERIC_TOKENS = {
"score",
"rating",
"votes",
"income",
"revenue",
"salary",
"amount",
"price",
"cost",
"duration",
"founded",
}
LOW_SIGNAL_DEDUP_TOKENS = {
"genre",
"industry",
"sector",
"country",
"release",
"year",
"type",
"category",
}
def strip_sql_comments(sql: str) -> str:
"""Remove SQL comments while preserving string literals."""
if not sql:
return ""
result = []
i = 0
n = len(sql)
while i < n:
char = sql[i]