# Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. """ SQLSherlock-Env — Production Inference Script. STDOUT FORMAT (mandatory — judges parse this exactly): [START] task= env=sqlsherlock_env model= [STEP] step= action= reward=<0.00> done= error= [END] success= steps= rewards= Environment variables: API_BASE_URL LLM endpoint (default: https://api.openai.com/v1) MODEL_NAME Model id (default: gpt-4.1-mini) HF_TOKEN API key (required — no default) SPACE_URL Server URL (default: http://localhost:7860) """ import json import os import re import sys import time from typing import Any, Optional from openai import OpenAI # --------------------------------------------------------------------------- # Configuration # --------------------------------------------------------------------------- DEMO_DATASET = "phihung/titanic" INFERENCE_MAX_ROWS = 500 ENV_NAME = "sqlsherlock_env" API_BASE_URL = os.getenv("API_BASE_URL", "https://api.openai.com/v1") MODEL_NAME = os.getenv("MODEL_NAME", "gpt-4.1-mini") HF_TOKEN = os.getenv("HF_TOKEN") SPACE_URL = os.getenv("SPACE_URL", "http://localhost:7860") # DEMO_INTENT selects which intent's 3 difficulty tasks to run. # Set via environment variable — defaults to "visualization". DEMO_INTENT = os.getenv("DEMO_INTENT", "visualization") if HF_TOKEN is None: raise ValueError("HF_TOKEN environment variable is required") # 9 tasks: 3 intents × 3 difficulty levels (easy / medium / hard) # Task IDs follow the pattern: {intent_prefix}_{difficulty} # viz_easy, viz_medium, viz_hard # ml_easy, ml_medium, ml_hard # bq_easy, bq_medium, bq_hard STEP_BUDGETS: dict[str, int] = { "viz_easy": 30, "viz_medium": 40, "viz_hard": 50, "ml_easy": 30, "ml_medium": 40, "ml_hard": 50, "bq_easy": 30, "bq_medium": 40, "bq_hard": 50, } # Map each intent to its 3 difficulty tasks (easy → medium → hard) INTENT_TASK_MAP: dict[str, list[tuple[str, str]]] = { "visualization": [("viz_easy", "easy"), ("viz_medium", "medium"), ("viz_hard", "hard")], "ml_training": [("ml_easy", "easy"), ("ml_medium", "medium"), ("ml_hard", "hard")], "business_query": [("bq_easy", "easy"), ("bq_medium", "medium"), ("bq_hard", "hard")], } # Select tasks for the demo intent; fall back to visualization if unknown TASKS: list[tuple[str, str, str]] = [ (task_id, difficulty, DEMO_INTENT) for task_id, difficulty in INTENT_TASK_MAP.get(DEMO_INTENT, INTENT_TASK_MAP["visualization"]) ] # --------------------------------------------------------------------------- # Mandatory log helpers # --------------------------------------------------------------------------- def log_start(task: str, model: str) -> None: print(f"[START] task={task} env={ENV_NAME} model={model}", flush=True) def log_step(step: int, action: str, reward: float, done: bool, error: Optional[str] = None) -> None: action_str = action.replace("\n", " ").replace("\r", " ").strip()[:120] error_str = ( error.replace("\n", " ").replace("\r", " ").strip()[:120] if error else "null" ) print( f"[STEP] step={step} action={action_str} " f"reward={reward:.2f} done={str(done).lower()} " f"error={error_str}", flush=True, ) def log_end(success: bool, steps: int, rewards: list[float]) -> None: # Guard: rewards must never be empty — use "0.00" if no steps executed rewards_str = ",".join(f"{r:.2f}" for r in rewards) if rewards else "0.00" print( f"[END] success={str(success).lower()} steps={steps} " f"rewards={rewards_str}", flush=True, ) def _parse_score(feedback: str) -> Optional[float]: m = re.search(r"[Gg]rader\s+score\s*=?\s*(\d+\.\d+)", feedback) if m: try: # Clamp to (0.01, 0.99) — score must be strictly between 0 and 1 return max(0.01, min(0.99, float(m.group(1)))) except (ValueError, TypeError): pass return None def _label(d: dict) -> str: a = d.get("action_type", "?") if a == "fix_cell": return f"fix_cell(row={d.get('row_id')},col={d.get('column')},val={d.get('value')})" if a == "profile_column": return f"profile_column({d.get('column')})" if a == "run_sql": return f"run_sql({(d.get('sql') or '')[:40]})" if a == "delete_row": return f"delete_row(row={d.get('row_id')})" if a == "join_tables": return f"join_tables({d.get('table')},{d.get('table2')},key={d.get('key')})" if a == "classify_intent": return f"classify_intent({d.get('value')})" if a == "select_tables": return f"select_tables({d.get('tables')})" return a # --------------------------------------------------------------------------- # LLM call helper # --------------------------------------------------------------------------- SYSTEM_PROMPT = """You are an expert data scientist performing a complete data quality audit. You have inspected and profiled the data. Now decide the NEXT action. Available actions (respond with ONE JSON object): - {"action_type":"profile_column","table":"T","column":"C"} - {"action_type":"run_sql","sql":"SELECT id FROM \\"T\\" WHERE \\"C\\" IS NULL LIMIT 10"} - {"action_type":"fix_column","table":"T","column":"C","value":29.5,"reason":"bulk fix nulls, median=29.5"} - {"action_type":"fix_cell","table":"T","row_id":6,"column":"C","value":29.5,"reason":"null, median=29.5"} - {"action_type":"delete_row","table":"T","row_id":8,"reason":"duplicate Name"} - {"action_type":"validate"} - {"action_type":"submit"} PREFER fix_column over fix_cell when fixing nulls — it fixes ALL nulls in one step. RULES: - Use EXACT row ids from SQL results. NEVER guess ids. - Numeric nulls → column MEDIAN. String nulls → column MODE. - Skip columns with null_rate >= 0.70 (structural nulls). - OUTLIER detection uses IQR (Q1 - 3×IQR, Q3 + 3×IQR). Fix detected outliers to the column median. - TRAP: if a value was flagged by the system as suspicious but is NOT in profile_column outlier list, DO NOT TOUCH it. Respond with ONLY one JSON object. No markdown, no text.""" def _call_llm(client: OpenAI, messages: list[dict]) -> Optional[dict]: """Call LLM and parse JSON action. Returns None on failure.""" try: resp = client.chat.completions.create( model=MODEL_NAME, messages=messages, max_tokens=300, temperature=0.0, ) raw = (resp.choices[0].message.content or "").strip() raw = re.sub(r"^```[a-z]*\n?", "", raw) raw = re.sub(r"\n?```\s*$", "", raw) raw = raw.strip() if not raw.startswith("{"): start = raw.find("{") end = raw.rfind("}") if start >= 0 and end > start: raw = raw[start:end + 1] return json.loads(raw) except Exception: return None # --------------------------------------------------------------------------- # Task runner # --------------------------------------------------------------------------- def run_task(task_id: str, intent: str = "") -> float: """Run one task episode. Guarantees [START] … [END] on stdout regardless of errors.""" pkg_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "sqlsherlock_env") if pkg_dir not in sys.path: sys.path.insert(0, pkg_dir) from client import SQLSherlockEnv from models import SQLSherlockAction budget = STEP_BUDGETS[task_id] rewards: list[float] = [] step_num = 0 # initialised here so except/finally can always read it score = 0.0 success = False env = None # initialised inside try so finally can check log_start(task=task_id, model=MODEL_NAME) try: llm = OpenAI(base_url=API_BASE_URL, api_key=HF_TOKEN) env = SQLSherlockEnv(base_url=SPACE_URL) obs = env.reset( dataset=DEMO_DATASET, task_id=task_id, max_rows=INFERENCE_MAX_ROWS, intent=intent or None, ) table = list(obs.tables_summary.keys())[0] if obs.tables_summary else (DEMO_DATASET.split("/")[-1]) columns = obs.tables_summary.get(table, {}).get("columns", []) visible_cols = [c for c in columns if c not in ("id", "_source_format")] done = False col_profiles: dict[str, dict] = {} llm_messages = [{"role": "system", "content": SYSTEM_PROMPT}] def _do_step(action_dict: dict) -> tuple[float, bool]: nonlocal step_num, done, obs step_num += 1 # Always allow submit through — it must execute to produce a valid score is_submit = action_dict.get("action_type") == "submit" if (step_num > budget and not is_submit) or done: return 0.0, True filtered = {k: v for k, v in action_dict.items() if v is not None} action = SQLSherlockAction(**filtered) reward = 0.0 error_msg = None try: obs, reward, done, _ = env.step(action) reward = float(reward or 0.0) except Exception as exc: error_msg = str(exc)[:80] rewards.append(reward) log_step(step_num, _label(action_dict), reward, done, error_msg) return reward, done # Phase 1: Inspect _do_step({"action_type": "inspect", "table": table}) # Phase 2: Profile + bulk fix interleaved for col in visible_cols: if done or step_num >= budget - 2: break _do_step({"action_type": "profile_column", "table": table, "column": col}) if not obs.query_result: continue profile = obs.query_result[0] col_profiles[col] = profile null_count = profile.get("null_count", 0) null_rate = profile.get("null_rate", 0.0) dtype = profile.get("dtype", "unknown") median_val = profile.get("median") mode_val = profile.get("mode") mean_val = profile.get("mean") if null_count == 0: continue if dtype in ("int", "float"): fill_value = median_val or mean_val or 0 elif null_rate >= 0.70: fill_value = "Unknown" else: fill_value = mode_val or "Unknown" strategy = "median" if dtype in ("int", "float") else "mode" _do_step({ "action_type": "fix_column", "table": table, "column": col, "value": fill_value, "reason": f"bulk fix {null_count} nulls in {col}, {strategy}={fill_value}", }) # Phase 3: LLM-assisted advanced cleaning if not done and step_num < budget - 3: fixed_summary = f"Profiled {len(col_profiles)} columns. Bulk-fixed nulls." remaining_budget = budget - step_num - 2 llm_messages.append({"role": "user", "content": ( f"Table: \"{table}\", Columns: {visible_cols}\n" f"Done so far: {fixed_summary}\n" f"Remaining budget: {remaining_budget} actions.\n" "Find remaining issues: type errors, negatives, duplicates, whitespace. " "Respond with one JSON action, or {\"action_type\":\"validate\"} if done." )}) for _ in range(min(remaining_budget, 6)): if done or step_num >= budget - 2: break action_dict = _call_llm(llm, llm_messages) if action_dict is None or action_dict.get("action_type") in ("validate", "submit"): break _, d = _do_step(action_dict) if d: break feedback = (obs.last_feedback or "")[:300] if obs.query_result: ids = [r2.get("id") for r2 in obs.query_result if r2.get("id") is not None] if ids: feedback += f"\nRow IDs: {ids[:15]}" llm_messages.append({"role": "assistant", "content": json.dumps(action_dict)}) llm_messages.append({"role": "user", "content": feedback + "\nNext action?"}) # Phase 4: Validate → Submit if not done and step_num < budget: _do_step({"action_type": "validate"}) if not done: _do_step({"action_type": "submit"}) if obs.last_feedback: parsed = _parse_score(obs.last_feedback) if parsed is not None: score = parsed # already clamped to (0.01, 0.99) # Fallback when grader score wasn't parseable — use scaled cumulative reward if score < 0.01 and rewards: positive = sum(r for r in rewards if r > 0) score = max(0.01, min(0.99, positive / max(budget * 0.15, 0.01))) success = score >= 0.50 except Exception as exc: # Catch-all — log a terminal error step so judges see what failed log_step( max(step_num + 1, 1), "error", 0.0, True, f"{type(exc).__name__}: {str(exc)[:60]}", ) finally: # [END] is ALWAYS emitted here — even after early returns / exceptions if env is not None: try: env.close() except Exception: pass log_end(success=success, steps=step_num, rewards=rewards) return score # --------------------------------------------------------------------------- # Main # --------------------------------------------------------------------------- def main() -> None: wall_start = time.time() all_scores: list[float] = [] print( f"\n=== SQLSherlock-Env intent={DEMO_INTENT} model={MODEL_NAME} ===\n", file=sys.stderr, ) for task_id, _, intent in TASKS: score = run_task(task_id, intent=intent) all_scores.append(score) time.sleep(1) avg = sum(all_scores) / len(all_scores) if all_scores else 0.0 total = time.time() - wall_start print( f"\n=== Results intent={DEMO_INTENT} avg={avg:.3f} " f"runtime={total:.1f}s ===", file=sys.stderr, ) for (tid, diff, _), sc in zip(TASKS, all_scores): bar = "\u2588" * int(sc * 20) + "\u2591" * (20 - int(sc * 20)) print(f" [{diff:<6}] {tid:<20} [{bar}] {sc:.3f}", file=sys.stderr) if __name__ == "__main__": main()