Spaces:
Running
Running
File size: 5,122 Bytes
f636e22 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 | """Execution checker for text-to-SQL: the single source of truth for pass/fail.
A predicted query is CORRECT iff executing it on the question's SQLite database
returns the same result as the gold query (Spider-style execution accuracy).
Outcome statuses (feed the self-heal Failure objects):
correct - result matches gold
wrong_result - ran fine, result differs from gold
sql_error - sqlite raised (bad column/table/syntax ...); error_message holds it
timeout - exceeded the time budget (likely a runaway / cartesian join)
empty_sql - nothing usable was extracted from the model output
"""
from __future__ import annotations
import re
import sqlite3
import time
from dataclasses import asdict, dataclass
TIMEOUT_S = 10.0
MAX_ROWS = 100_000 # guard against pathological result sets
@dataclass
class ExecResult:
status: str # correct | wrong_result | sql_error | timeout | empty_sql
error_message: str = ""
pred_rows: int | None = None
gold_rows: int | None = None
@property
def correct(self) -> bool:
return self.status == "correct"
def to_dict(self) -> dict:
return asdict(self)
_FENCE = re.compile(r"```(?:sql)?\s*(.*?)```", re.DOTALL | re.IGNORECASE)
_THINK = re.compile(r"<think>.*?</think>", re.DOTALL | re.IGNORECASE)
def extract_sql(text: str) -> str:
"""Pull a SQL statement out of raw model output.
Handles <think> blocks, ```sql fences, and a trailing explanation. Returns ""
if nothing that looks like SQL is found.
"""
if not text:
return ""
text = _THINK.sub("", text)
m = _FENCE.search(text)
if m:
text = m.group(1)
text = text.strip()
# Start at the first SELECT/WITH; drop anything before it.
m = re.search(r"\b(select|with)\b", text, re.IGNORECASE)
if not m:
return ""
text = text[m.start():]
# Keep only the first statement.
text = text.split(";")[0].strip()
return text
def _connect(db_path: str) -> sqlite3.Connection:
# Read-only: a model-written query must never modify the benchmark databases.
uri = "file:" + db_path.replace("\\", "/") + "?mode=ro"
conn = sqlite3.connect(uri, uri=True)
conn.text_factory = lambda b: b.decode("utf-8", errors="ignore")
return conn
def _run(conn: sqlite3.Connection, sql: str, timeout_s: float):
deadline = time.monotonic() + timeout_s
def _abort() -> int:
return 1 if time.monotonic() > deadline else 0
conn.set_progress_handler(_abort, 10_000)
try:
cur = conn.execute(sql)
rows = cur.fetchmany(MAX_ROWS)
finally:
conn.set_progress_handler(None, 0)
return rows
def _norm_cell(v):
if isinstance(v, float):
return round(v, 6)
return v
def _norm_rows(rows) -> list[tuple]:
return [tuple(_norm_cell(c) for c in r) for r in rows]
def _results_match(pred, gold, ordered: bool) -> bool:
pred, gold = _norm_rows(pred), _norm_rows(gold)
if ordered:
return pred == gold
key = lambda r: tuple(repr(c) for c in r) # noqa: E731 (mixed types sort safely)
return sorted(pred, key=key) == sorted(gold, key=key)
def check(db_path: str, pred_sql: str, gold_sql: str, timeout_s: float = TIMEOUT_S) -> ExecResult:
"""Execute pred and gold on db_path and compare. Never raises."""
# Always extract: strips fences/<think>/trailing prose and keeps only the FIRST
# statement (a model that appends "; <explanation>" must not be marked wrong).
pred_sql = extract_sql(pred_sql)
if not pred_sql:
return ExecResult("empty_sql", "no SQL found in model output")
try:
conn = _connect(db_path)
except Exception as e: # noqa: BLE001
return ExecResult("sql_error", f"cannot open database: {e}")
try:
try:
gold = _run(conn, gold_sql, timeout_s * 3)
except Exception as e: # noqa: BLE001
return ExecResult("sql_error", f"GOLD query failed (bad benchmark row): {e}")
try:
pred = _run(conn, pred_sql, timeout_s)
except sqlite3.OperationalError as e:
msg = str(e)
if "interrupted" in msg.lower():
return ExecResult("timeout", f"query exceeded {timeout_s}s", None, len(gold))
return ExecResult("sql_error", msg, None, len(gold))
except Exception as e: # noqa: BLE001
return ExecResult("sql_error", f"{type(e).__name__}: {e}", None, len(gold))
ordered = bool(re.search(r"\border\s+by\b", gold_sql, re.IGNORECASE))
ok = _results_match(pred, gold, ordered)
return ExecResult("correct" if ok else "wrong_result", "", len(pred), len(gold))
finally:
conn.close()
def schema_text(db_path: str) -> str:
"""CREATE TABLE statements for the prompt (the schema the agent is shown)."""
conn = _connect(db_path)
try:
rows = conn.execute(
"SELECT sql FROM sqlite_master WHERE type='table' AND sql IS NOT NULL ORDER BY name"
).fetchall()
finally:
conn.close()
return "\n".join(r[0].strip() + ";" for r in rows)
|