Spaces:
Running
Running
| """The heal step: classify each failure, generate a corrected SQL query, keep it only | |
| if it ACTUALLY produces the gold result. Output: DPO preference pairs. | |
| Replaces AgentTune's TrainingExampleGenerator, which is tool-call-specific (it drops | |
| every failure whose root cause isn't wrong_tool/loop_collapse, asks the LLM for "the | |
| EXACT JSON tool call", and trains on a placeholder prompt). The classification stage | |
| keeps AgentTune's FailureClassifier contract ({root_cause, confidence, analysis}) but | |
| with a SQL-specific taxonomy. The output rows carry chosen/rejected exactly as | |
| agenttune.decide.closed_loop.contracts.TrainingExample expects, so the retrain step | |
| consumes them through Lexsi's own examples_to_dpo_dataset / run_retrain. | |
| The healer NEVER sees the gold SQL. Gold is used only to verify a proposed fix by | |
| executing both and comparing results (execcheck). Failed fixes are retried with the | |
| new error message, up to --max-attempts. | |
| Runs locally (needs the Spider DBs on disk + the HF router for LLM calls). | |
| Usage: | |
| python heal.py --run baseline-pool600 --split-file ../../data/splits/pool.jsonl \ | |
| --out ../../data/heal [--limit 6] [--workers 6] | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import json | |
| import os | |
| import re | |
| import sys | |
| import threading | |
| import time | |
| import urllib.error | |
| import urllib.request | |
| from collections import Counter | |
| from concurrent.futures import ThreadPoolExecutor | |
| HERE = os.path.dirname(os.path.abspath(__file__)) | |
| sys.path.insert(0, HERE) | |
| from execcheck import check, extract_sql # noqa: E402 | |
| from student import PROMPT # noqa: E402 (the student's prompt - reused verbatim) | |
| ROUTER = "https://router.huggingface.co/v1/chat/completions" | |
| HEALER_MODEL = "CohereLabs/aya-expanse-32b" # on the hackathon's allowed model list | |
| RUNS_REPO = "HideIron/selfheal-sql-runs" | |
| SCHEMA_CAP = 7000 # chars of schema shown to the healer | |
| # SQL-specific root causes (replaces the tool-agent taxonomy in | |
| # agenttune FailureClassifier.ROOT_CAUSES; same output contract). | |
| ROOT_CAUSES = [ | |
| "schema_misread", # non-existent column/table, column on wrong table | |
| "wrong_join", # wrong join path, missing join, wrong join keys | |
| "wrong_aggregation", # count/sum/avg/group-by/having misuse | |
| "missing_or_wrong_filter", # WHERE/HAVING condition missing or wrong | |
| "value_mismatch", # wrong literal (e.g. 'US' vs 'United States') | |
| "wrong_output_columns", # right rows, wrong selected columns/order/distinct | |
| "overcomplicated_query", # needless subquery/join that changes the result | |
| "other", | |
| ] | |
| _counts_lock = threading.Lock() | |
| TOKENS = [0, 0] # prompt, completion | |
| CALLS = [0] | |
| def llm(messages: list[dict], temperature: float, max_tokens: int = 500) -> str: | |
| """One chat call to the HF router, with backoff on 429/5xx. Returns content.""" | |
| body = json.dumps( | |
| {"model": HEALER_MODEL, "messages": messages, "temperature": temperature, | |
| "max_tokens": max_tokens} | |
| ).encode() | |
| last = None | |
| for i in range(5): | |
| req = urllib.request.Request( | |
| ROUTER, data=body, | |
| headers={"Authorization": "Bearer " + os.environ["HF_TOKEN"], | |
| "Content-Type": "application/json"}, | |
| ) | |
| try: | |
| with urllib.request.urlopen(req, timeout=180) as r: | |
| d = json.load(r) | |
| u = d.get("usage", {}) | |
| with _counts_lock: | |
| TOKENS[0] += u.get("prompt_tokens", 0) | |
| TOKENS[1] += u.get("completion_tokens", 0) | |
| CALLS[0] += 1 | |
| return d["choices"][0]["message"]["content"] | |
| except urllib.error.HTTPError as e: | |
| last = f"HTTP {e.code}: {e.read()[:200]!r}" | |
| if e.code in (429, 500, 502, 503, 504): | |
| time.sleep(2 ** i) | |
| continue | |
| raise RuntimeError(last) from e | |
| except Exception as e: # noqa: BLE001 (timeouts, resets) | |
| last = repr(e) | |
| time.sleep(2 ** i) | |
| raise RuntimeError(f"router call failed after retries: {last}") | |
| def _json_block(text: str) -> dict: | |
| m = re.search(r"\{.*\}", text, re.DOTALL) | |
| return json.loads(m.group(0)) if m else {} | |
| def classify(fail: dict, srow: dict) -> dict: | |
| """SQL-specific root-cause classification (AgentTune ClassifiedFailure contract).""" | |
| sys_p = ( | |
| "You are an expert SQL failure analyst. Classify the failure into EXACTLY ONE " | |
| "root cause from this list:\n" + ", ".join(ROOT_CAUSES) + "\n\n" | |
| 'Return ONLY a JSON object: {"root_cause": "<one of the list>", ' | |
| '"confidence": 0.0-1.0, "analysis": "one short sentence"}' | |
| ) | |
| user = ( | |
| f"Schema:\n{srow['schema'][:SCHEMA_CAP]}\n\nQuestion: {srow['question']}\n\n" | |
| f"Model's SQL:\n{fail['raw_output'][:1200]}\n\n" | |
| f"Outcome: {fail['status']}" | |
| + (f" - {fail['error_message']}" if fail.get("error_message") else "") | |
| + ( | |
| f" (query returned {fail['pred_rows']} rows, expected result has {fail['gold_rows']} rows)" | |
| if fail.get("pred_rows") is not None and fail.get("gold_rows") is not None | |
| else "" | |
| ) | |
| ) | |
| try: | |
| p = _json_block(llm([{"role": "system", "content": sys_p}, | |
| {"role": "user", "content": user}], 0.0, 250)) | |
| rc = p.get("root_cause", "other") | |
| return { | |
| "root_cause": rc if rc in ROOT_CAUSES else "other", | |
| "confidence": float(p.get("confidence", 0.5)), | |
| "analysis": str(p.get("analysis", ""))[:400], | |
| "classify_error": None, | |
| } | |
| except Exception as e: # noqa: BLE001 | |
| return {"root_cause": "other", "confidence": 0.0, "analysis": "", | |
| "classify_error": str(e)[:300]} | |
| def generate_fix(fail: dict, srow: dict, cls: dict, max_attempts: int) -> tuple[str | None, list]: | |
| """Ask the healer for corrected SQL; keep only a fix that executes to the gold | |
| result. Retries feed the candidate's own error back. Returns (sql|None, attempts).""" | |
| base = ( | |
| f"A SQLite query is wrong and you must fix it.\n\n" | |
| f"Schema:\n{srow['schema'][:SCHEMA_CAP]}\n\n" | |
| f"Question: {srow['question']}\n\n" | |
| f"Wrong SQL:\n{fail['raw_output'][:1200]}\n\n" | |
| f"Observed problem: {fail['status']}" | |
| + (f" - {fail['error_message']}" if fail.get("error_message") else "") | |
| + (f"\nLikely root cause: {cls['root_cause']} ({cls['analysis']})" if cls.get("analysis") else "") | |
| + "\n\nRules: make the MINIMAL edit that fixes the query. Keep the original's " | |
| "formatting, aliasing style, capitalisation and column order; do not add joins, " | |
| "aliases, extra columns or reformat. Change only what is wrong.\n" | |
| "Return ONLY the corrected SQLite query, nothing else." | |
| ) | |
| msgs = [{"role": "user", "content": base}] | |
| attempts = [] | |
| for k in range(max_attempts): | |
| try: | |
| out = llm(msgs, 0.0 if k == 0 else 0.4, 400) | |
| except Exception as e: # noqa: BLE001 | |
| attempts.append({"attempt": k + 1, "error": f"llm: {str(e)[:200]}"}) | |
| break | |
| cand = extract_sql(out) | |
| if not cand: | |
| attempts.append({"attempt": k + 1, "error": "no SQL in healer output"}) | |
| feedback = "Your reply contained no SQL. Return ONLY the corrected SQLite query." | |
| else: | |
| res = check(srow["db_path"], cand, srow["gold_sql"]) | |
| attempts.append({"attempt": k + 1, "sql": cand, "status": res.status, | |
| "error": res.error_message[:200]}) | |
| if res.correct: | |
| return cand, attempts | |
| feedback = ( | |
| f"Your query:\n{cand}\n\nIt is still wrong: {res.status}" | |
| + (f" - {res.error_message[:300]}" if res.error_message else "") | |
| + ( | |
| f" (it returned {res.pred_rows} rows; the correct result has {res.gold_rows} rows)" | |
| if res.pred_rows is not None and res.gold_rows is not None | |
| else "" | |
| ) | |
| + "\nTry again. Return ONLY the corrected SQLite query." | |
| ) | |
| msgs = msgs + [{"role": "assistant", "content": out}, | |
| {"role": "user", "content": feedback}] | |
| return None, attempts | |
| def heal_one(fail: dict, srow: dict, max_attempts: int) -> dict: | |
| cls = classify(fail, srow) | |
| fixed_sql, attempts = generate_fix(fail, srow, cls, max_attempts) | |
| rejected = extract_sql(fail["raw_output"]) or fail["raw_output"].strip() | |
| return { | |
| "qid": fail["qid"], "db_id": fail["db_id"], "question": srow["question"], | |
| "status_before": fail["status"], "error_before": fail.get("error_message", ""), | |
| **{k: cls[k] for k in ("root_cause", "confidence", "analysis", "classify_error")}, | |
| "healed": fixed_sql is not None, | |
| "attempts_used": len(attempts), "attempts": attempts, | |
| # --- TrainingExample-ready fields (consumed by train_dpo.py) --- | |
| "prompt_user_content": PROMPT.format(schema=srow["schema"], question=srow["question"]), | |
| "chosen_sql": fixed_sql, | |
| "rejected_sql": rejected, | |
| "healer_model": HEALER_MODEL, | |
| } | |
| def main() -> None: | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--run", default="baseline-pool600", help="runs-repo folder with results.jsonl") | |
| ap.add_argument("--split-file", required=True) | |
| ap.add_argument("--out", required=True) | |
| ap.add_argument("--limit", type=int, default=None) | |
| ap.add_argument("--workers", type=int, default=6) | |
| ap.add_argument("--max-attempts", type=int, default=3) | |
| ap.add_argument("--no-push", action="store_true") | |
| ap.add_argument("--tag", default="", help="suffix for the runs-repo push path") | |
| args = ap.parse_args() | |
| from huggingface_hub import HfApi, hf_hub_download | |
| rp = hf_hub_download(RUNS_REPO, f"runs/{args.run}/results.jsonl", repo_type="dataset") | |
| results = [json.loads(l) for l in open(rp, encoding="utf-8")] | |
| splits = {r["qid"]: r for r in map(json.loads, open(args.split_file, encoding="utf-8"))} | |
| fails = [r for r in results if r["status"] != "correct"] | |
| if args.limit: | |
| fails = fails[: args.limit] | |
| print(f"[heal] {len(fails)} failures to heal (of {len(results)} runs)", flush=True) | |
| t0 = time.time() | |
| done = [0] | |
| def work(f): | |
| row = heal_one(f, splits[f["qid"]], args.max_attempts) | |
| with _counts_lock: | |
| done[0] += 1 | |
| if done[0] % 20 == 0: | |
| print(f"[heal] {done[0]}/{len(fails)} ({time.time()-t0:.0f}s, " | |
| f"{CALLS[0]} calls)", flush=True) | |
| return row | |
| with ThreadPoolExecutor(max_workers=args.workers) as ex: | |
| rows = list(ex.map(work, fails)) | |
| os.makedirs(args.out, exist_ok=True) | |
| with open(os.path.join(args.out, "pairs.jsonl"), "w", encoding="utf-8") as f: | |
| for r in rows: | |
| f.write(json.dumps(r, ensure_ascii=False) + "\n") | |
| healed = [r for r in rows if r["healed"]] | |
| summary = { | |
| "run": args.run, "healer_model": HEALER_MODEL, "n_failures": len(fails), | |
| "n_healed": len(healed), "heal_yield": round(len(healed) / max(len(fails), 1), 4), | |
| "by_root_cause": dict(Counter(r["root_cause"] for r in rows)), | |
| "healed_by_root_cause": dict(Counter(r["root_cause"] for r in healed)), | |
| "attempts_histogram": dict(Counter(r["attempts_used"] for r in healed)), | |
| "classify_errors": sum(1 for r in rows if r["classify_error"]), | |
| "llm_calls": CALLS[0], "prompt_tokens": TOKENS[0], "completion_tokens": TOKENS[1], | |
| "elapsed_s": round(time.time() - t0, 1), "max_attempts": args.max_attempts, | |
| } | |
| with open(os.path.join(args.out, "heal_summary.json"), "w", encoding="utf-8") as f: | |
| json.dump(summary, f, indent=2) | |
| print("[heal] SUMMARY", json.dumps(summary), flush=True) | |
| if not args.no_push: | |
| dest = f"heal/{args.run}{args.tag}" | |
| HfApi().upload_folder(repo_id=RUNS_REPO, repo_type="dataset", folder_path=args.out, | |
| path_in_repo=dest, commit_message=f"heal {dest}") | |
| print("[heal] pushed to", RUNS_REPO, dest, flush=True) | |
| if __name__ == "__main__": | |
| main() | |