readme + dash cleanup in source

src/selfheal_sql/audit_eval.py

@@ -1,107 +1,107 @@
-"""AuditKIT agent_eval over a student run: independent outcome verification.
-
-Each logged run row becomes a recorded AgentEpisode (user prompt -> assistant raw
-output). The outcome oracle is a CustomPredicate that RE-EXECUTES the model's SQL
-against the real database and compares with the gold result — AuditKIT's verdict is a
-fresh execution, not a read of our logged status. We then cross-check AuditKIT's
-verdicts against the harness log (they must agree 100%).
-
-Usage: python audit_eval.py --run baseline-test --split-file ../../data/splits/test.jsonl
-"""
-
-from __future__ import annotations
-
-import argparse
-import json
-import os
-import sys
-from collections import Counter
-
-HERE = os.path.dirname(os.path.abspath(__file__))
-sys.path.insert(0, HERE)
-from execcheck import check  # noqa: E402
-from student import PROMPT  # noqa: E402
-
-RUNS_REPO = "HideIron/selfheal-sql-runs"
-
-
-def main() -> None:
-    ap = argparse.ArgumentParser()
-    ap.add_argument("--run", required=True)
-    ap.add_argument("--split-file", required=True)
-    ap.add_argument("--out", default=None)
-    ap.add_argument("--no-push", action="store_true")
-    args = ap.parse_args()
-
-    from huggingface_hub import HfApi, hf_hub_download
-
-    from auditkit.agent_eval import (
-        AgentCase,
-        AgentEvalRunner,
-        AgentEvalSpec,
-        CustomPredicate,
-        episode_from_openai_messages,
-    )
-
-    rp = hf_hub_download(RUNS_REPO, f"runs/{args.run}/results.jsonl", repo_type="dataset")
-    rows = [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"))}
-
-    def exec_predicate(episode, case):
-        s = splits[case.id]
-        res = check(s["db_path"], episode.final_answer or "", s["gold_sql"])
-        return bool(res.correct)
-
-    episodes, cases = [], []
-    for r in rows:
-        s = splits[r["qid"]]
-        episodes.append(
-            episode_from_openai_messages(
-                [
-                    {"role": "user",
-                     "content": PROMPT.format(schema=s["schema"], question=s["question"])},
-                    {"role": "assistant", "content": r["raw_output"]},
-                ],
-                case_id=r["qid"],
-                final_answer=r["raw_output"],
-            )
-        )
-        cases.append(
-            AgentCase(
-                id=r["qid"], task=s["question"], category="text_to_sql",
-                outcome=CustomPredicate(exec_predicate, version="execcheck-v1",
-                                        name="spider_execution_accuracy"),
-                metadata={"db_id": s["db_id"]},
-            )
-        )
-
-    spec = AgentEvalSpec(cases=cases, mode="recorded", episodes=episodes,
-                         scorers=[], trials=1)
-    result = AgentEvalRunner().run(spec)
-
-    verdicts = {row.case_id: row.outcome["verdict"] for row in result.rows}
-    logged = {r["qid"]: ("success" if r["status"] == "correct" else "failure") for r in rows}
-    agree = sum(verdicts[q] == logged[q] for q in verdicts)
-    out = {
-        "run": args.run,
-        "auditkit_summary": result.summary(),
-        "verdict_counts": dict(Counter(verdicts.values())),
-        "n": len(verdicts),
-        "success_rate": round(sum(v == "success" for v in verdicts.values()) / len(verdicts), 4),
-        "crosscheck_agree_with_harness": f"{agree}/{len(verdicts)}",
-        "oracle": "CustomPredicate(spider_execution_accuracy, version=execcheck-v1) — fresh re-execution",
-    }
-    out_dir = args.out or os.path.join(HERE, "..", "..", "data", "audit", args.run)
-    os.makedirs(out_dir, exist_ok=True)
-    with open(os.path.join(out_dir, "auditkit_report.json"), "w", encoding="utf-8") as f:
-        json.dump(out, f, indent=2, default=str)
-    print(json.dumps(out, indent=2, default=str))
-
-    if not args.no_push:
-        HfApi().upload_folder(repo_id=RUNS_REPO, repo_type="dataset", folder_path=out_dir,
-                              path_in_repo=f"audit/{args.run}",
-                              commit_message=f"auditkit report {args.run}")
-
-
-if __name__ == "__main__":
-    main()
+"""AuditKIT agent_eval over a student run: independent outcome verification.
+
+Each logged run row becomes a recorded AgentEpisode (user prompt -> assistant raw
+output). The outcome oracle is a CustomPredicate that RE-EXECUTES the model's SQL
+against the real database and compares with the gold result - AuditKIT's verdict is a
+fresh execution, not a read of our logged status. We then cross-check AuditKIT's
+verdicts against the harness log (they must agree 100%).
+
+Usage: python audit_eval.py --run baseline-test --split-file ../../data/splits/test.jsonl
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import sys
+from collections import Counter
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+sys.path.insert(0, HERE)
+from execcheck import check  # noqa: E402
+from student import PROMPT  # noqa: E402
+
+RUNS_REPO = "HideIron/selfheal-sql-runs"
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--run", required=True)
+    ap.add_argument("--split-file", required=True)
+    ap.add_argument("--out", default=None)
+    ap.add_argument("--no-push", action="store_true")
+    args = ap.parse_args()
+
+    from huggingface_hub import HfApi, hf_hub_download
+
+    from auditkit.agent_eval import (
+        AgentCase,
+        AgentEvalRunner,
+        AgentEvalSpec,
+        CustomPredicate,
+        episode_from_openai_messages,
+    )
+
+    rp = hf_hub_download(RUNS_REPO, f"runs/{args.run}/results.jsonl", repo_type="dataset")
+    rows = [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"))}
+
+    def exec_predicate(episode, case):
+        s = splits[case.id]
+        res = check(s["db_path"], episode.final_answer or "", s["gold_sql"])
+        return bool(res.correct)
+
+    episodes, cases = [], []
+    for r in rows:
+        s = splits[r["qid"]]
+        episodes.append(
+            episode_from_openai_messages(
+                [
+                    {"role": "user",
+                     "content": PROMPT.format(schema=s["schema"], question=s["question"])},
+                    {"role": "assistant", "content": r["raw_output"]},
+                ],
+                case_id=r["qid"],
+                final_answer=r["raw_output"],
+            )
+        )
+        cases.append(
+            AgentCase(
+                id=r["qid"], task=s["question"], category="text_to_sql",
+                outcome=CustomPredicate(exec_predicate, version="execcheck-v1",
+                                        name="spider_execution_accuracy"),
+                metadata={"db_id": s["db_id"]},
+            )
+        )
+
+    spec = AgentEvalSpec(cases=cases, mode="recorded", episodes=episodes,
+                         scorers=[], trials=1)
+    result = AgentEvalRunner().run(spec)
+
+    verdicts = {row.case_id: row.outcome["verdict"] for row in result.rows}
+    logged = {r["qid"]: ("success" if r["status"] == "correct" else "failure") for r in rows}
+    agree = sum(verdicts[q] == logged[q] for q in verdicts)
+    out = {
+        "run": args.run,
+        "auditkit_summary": result.summary(),
+        "verdict_counts": dict(Counter(verdicts.values())),
+        "n": len(verdicts),
+        "success_rate": round(sum(v == "success" for v in verdicts.values()) / len(verdicts), 4),
+        "crosscheck_agree_with_harness": f"{agree}/{len(verdicts)}",
+        "oracle": "CustomPredicate(spider_execution_accuracy, version=execcheck-v1) - fresh re-execution",
+    }
+    out_dir = args.out or os.path.join(HERE, "..", "..", "data", "audit", args.run)
+    os.makedirs(out_dir, exist_ok=True)
+    with open(os.path.join(out_dir, "auditkit_report.json"), "w", encoding="utf-8") as f:
+        json.dump(out, f, indent=2, default=str)
+    print(json.dumps(out, indent=2, default=str))
+
+    if not args.no_push:
+        HfApi().upload_folder(repo_id=RUNS_REPO, repo_type="dataset", folder_path=out_dir,
+                              path_in_repo=f"audit/{args.run}",
+                              commit_message=f"auditkit report {args.run}")
+
+
+if __name__ == "__main__":
+    main()

src/selfheal_sql/gate_decision.py

@@ -1,115 +1,115 @@
-"""Deployment gate: decide whether the retrained adapter ships, using AgentTune's
-DeploymentGate verbatim (vendored at third_party/agenttune_gate, commit 1dfd054).
-
-The gate is model-agnostic by design: `model_runner_fn(input_text) -> verdict` is
-supplied by the caller and compared against `expected_verdict` per case ("keeps the
-gate model-agnostic and testable without a GPU" — their docstring). Our runners look
-up each gate question's execution-checker outcome from the student harness's logged
-runs, so the verdict IS Spider execution accuracy.
-
-Usage:
-  python gate_decision.py --old baseline-gate --new healed-gate --out ../../data/gate
-"""
-
-from __future__ import annotations
-
-import argparse
-import importlib.util
-import json
-import os
-import sys
-import types
-from dataclasses import asdict
-
-HERE = os.path.dirname(os.path.abspath(__file__))
-ROOT = os.path.abspath(os.path.join(HERE, "..", ".."))
-RUNS_REPO = "HideIron/selfheal-sql-runs"
-
-
-def load_gate_module():
-    """Import the vendored deployment_gate.py without installing agenttune
-    (its only non-stdlib import is the PASS_VERDICTS constant, stubbed here 1:1)."""
-    pkg_names = ["agenttune", "agenttune.decide"]
-    for n in pkg_names:
-        if n not in sys.modules:
-            m = types.ModuleType(n)
-            m.__path__ = []  # mark as package
-            sys.modules[n] = m
-    st = types.ModuleType("agenttune.decide.state")
-    st.PASS_VERDICTS = frozenset({"PASS", "APPROVE", "COMPLETE"})  # verbatim
-    sys.modules["agenttune.decide.state"] = st
-    path = os.path.join(ROOT, "third_party", "agenttune_gate", "deployment_gate.py")
-    spec = importlib.util.spec_from_file_location("deployment_gate", path)
-    mod = importlib.util.module_from_spec(spec)
-    sys.modules["deployment_gate"] = mod  # dataclasses resolves cls.__module__ here
-    spec.loader.exec_module(mod)
-    return mod
-
-
-def results_lookup(run_name: str) -> dict[str, str]:
-    from huggingface_hub import hf_hub_download
-
-    p = hf_hub_download(RUNS_REPO, f"runs/{run_name}/results.jsonl", repo_type="dataset")
-    out = {}
-    for line in open(p, encoding="utf-8"):
-        r = json.loads(line)
-        out[r["qid"]] = "CORRECT" if r["status"] == "correct" else r["status"].upper()
-    return out
-
-
-def main() -> None:
-    ap = argparse.ArgumentParser()
-    ap.add_argument("--old", default="baseline-gate")
-    ap.add_argument("--new", default="healed-gate")
-    ap.add_argument("--out", default=os.path.join(ROOT, "data", "gate"))
-    ap.add_argument("--task-regression-tol", type=float, default=0.0)
-    ap.add_argument("--no-push", action="store_true")
-    args = ap.parse_args()
-
-    gate_mod = load_gate_module()
-    old_v, new_v = results_lookup(args.old), results_lookup(args.new)
-    qids = sorted(old_v)
-    assert set(qids) == set(new_v), "old/new runs cover different questions"
-
-    test_set = [{"pipeline_id": q, "input_text": q, "expected_verdict": "CORRECT"}
-                for q in qids]
-
-    gate = gate_mod.DeploymentGate()
-    decision = gate.evaluate_decision(
-        test_set=test_set,
-        old_model_fn=lambda qid: old_v[qid],
-        new_model_fn=lambda qid: new_v[qid],
-        task_regression_tol=args.task_regression_tol,
-    )
-    comparison = gate.compare_models(
-        test_set=test_set,
-        old_model_fn=lambda qid: old_v[qid],
-        new_model_fn=lambda qid: new_v[qid],
-    )
-
-    out = {
-        "gate_decision": asdict(decision),
-        "ab_comparison": asdict(comparison),
-        "old_run": args.old, "new_run": args.new,
-        "per_question_flips": {
-            "fixed": [q for q in qids if old_v[q] != "CORRECT" and new_v[q] == "CORRECT"],
-            "broken": [q for q in qids if old_v[q] == "CORRECT" and new_v[q] != "CORRECT"],
-        },
-        "gate_source": "AgentTune deployment_gate.py @ 1dfd054 (vendored verbatim)",
-    }
-    os.makedirs(args.out, exist_ok=True)
-    with open(os.path.join(args.out, "gate_decision.json"), "w", encoding="utf-8") as f:
-        json.dump(out, f, indent=2)
-    print(json.dumps({k: out[k] for k in ("gate_decision", "ab_comparison")}, indent=2))
-    print("fixed:", len(out["per_question_flips"]["fixed"]),
-          "broken:", len(out["per_question_flips"]["broken"]))
-
-    if not args.no_push:
-        from huggingface_hub import HfApi
-
-        HfApi().upload_folder(repo_id=RUNS_REPO, repo_type="dataset", folder_path=args.out,
-                              path_in_repo="gate", commit_message="gate decision")
-
-
-if __name__ == "__main__":
-    main()
+"""Deployment gate: decide whether the retrained adapter ships, using AgentTune's
+DeploymentGate verbatim (vendored at third_party/agenttune_gate, commit 1dfd054).
+
+The gate is model-agnostic by design: `model_runner_fn(input_text) -> verdict` is
+supplied by the caller and compared against `expected_verdict` per case ("keeps the
+gate model-agnostic and testable without a GPU" - their docstring). Our runners look
+up each gate question's execution-checker outcome from the student harness's logged
+runs, so the verdict IS Spider execution accuracy.
+
+Usage:
+  python gate_decision.py --old baseline-gate --new healed-gate --out ../../data/gate
+"""
+
+from __future__ import annotations
+
+import argparse
+import importlib.util
+import json
+import os
+import sys
+import types
+from dataclasses import asdict
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+ROOT = os.path.abspath(os.path.join(HERE, "..", ".."))
+RUNS_REPO = "HideIron/selfheal-sql-runs"
+
+
+def load_gate_module():
+    """Import the vendored deployment_gate.py without installing agenttune
+    (its only non-stdlib import is the PASS_VERDICTS constant, stubbed here 1:1)."""
+    pkg_names = ["agenttune", "agenttune.decide"]
+    for n in pkg_names:
+        if n not in sys.modules:
+            m = types.ModuleType(n)
+            m.__path__ = []  # mark as package
+            sys.modules[n] = m
+    st = types.ModuleType("agenttune.decide.state")
+    st.PASS_VERDICTS = frozenset({"PASS", "APPROVE", "COMPLETE"})  # verbatim
+    sys.modules["agenttune.decide.state"] = st
+    path = os.path.join(ROOT, "third_party", "agenttune_gate", "deployment_gate.py")
+    spec = importlib.util.spec_from_file_location("deployment_gate", path)
+    mod = importlib.util.module_from_spec(spec)
+    sys.modules["deployment_gate"] = mod  # dataclasses resolves cls.__module__ here
+    spec.loader.exec_module(mod)
+    return mod
+
+
+def results_lookup(run_name: str) -> dict[str, str]:
+    from huggingface_hub import hf_hub_download
+
+    p = hf_hub_download(RUNS_REPO, f"runs/{run_name}/results.jsonl", repo_type="dataset")
+    out = {}
+    for line in open(p, encoding="utf-8"):
+        r = json.loads(line)
+        out[r["qid"]] = "CORRECT" if r["status"] == "correct" else r["status"].upper()
+    return out
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--old", default="baseline-gate")
+    ap.add_argument("--new", default="healed-gate")
+    ap.add_argument("--out", default=os.path.join(ROOT, "data", "gate"))
+    ap.add_argument("--task-regression-tol", type=float, default=0.0)
+    ap.add_argument("--no-push", action="store_true")
+    args = ap.parse_args()
+
+    gate_mod = load_gate_module()
+    old_v, new_v = results_lookup(args.old), results_lookup(args.new)
+    qids = sorted(old_v)
+    assert set(qids) == set(new_v), "old/new runs cover different questions"
+
+    test_set = [{"pipeline_id": q, "input_text": q, "expected_verdict": "CORRECT"}
+                for q in qids]
+
+    gate = gate_mod.DeploymentGate()
+    decision = gate.evaluate_decision(
+        test_set=test_set,
+        old_model_fn=lambda qid: old_v[qid],
+        new_model_fn=lambda qid: new_v[qid],
+        task_regression_tol=args.task_regression_tol,
+    )
+    comparison = gate.compare_models(
+        test_set=test_set,
+        old_model_fn=lambda qid: old_v[qid],
+        new_model_fn=lambda qid: new_v[qid],
+    )
+
+    out = {
+        "gate_decision": asdict(decision),
+        "ab_comparison": asdict(comparison),
+        "old_run": args.old, "new_run": args.new,
+        "per_question_flips": {
+            "fixed": [q for q in qids if old_v[q] != "CORRECT" and new_v[q] == "CORRECT"],
+            "broken": [q for q in qids if old_v[q] == "CORRECT" and new_v[q] != "CORRECT"],
+        },
+        "gate_source": "AgentTune deployment_gate.py @ 1dfd054 (vendored verbatim)",
+    }
+    os.makedirs(args.out, exist_ok=True)
+    with open(os.path.join(args.out, "gate_decision.json"), "w", encoding="utf-8") as f:
+        json.dump(out, f, indent=2)
+    print(json.dumps({k: out[k] for k in ("gate_decision", "ab_comparison")}, indent=2))
+    print("fixed:", len(out["per_question_flips"]["fixed"]),
+          "broken:", len(out["per_question_flips"]["broken"]))
+
+    if not args.no_push:
+        from huggingface_hub import HfApi
+
+        HfApi().upload_folder(repo_id=RUNS_REPO, repo_type="dataset", folder_path=args.out,
+                              path_in_repo="gate", commit_message="gate decision")
+
+
+if __name__ == "__main__":
+    main()

src/selfheal_sql/heal.py

@@ -1,270 +1,270 @@
-"""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()
+"""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()

src/selfheal_sql/train_dpo.py

@@ -1,211 +1,211 @@
-"""DPO retrain via AgentTune's closed-loop retrain path (run inside an HF GPU Job).
-
-Consumes the healer's verified pairs, builds agenttune TrainingExample objects, and
-trains an adapter-only LoRA through agenttune's own
-  examples_to_dpo_dataset -> build_retrainer (create_agentic_trainer) -> train()
-exactly as the closed loop's run_retrain does (we call build_retrainer so we can save
-the adapter explicitly afterwards).
-
-Train/eval consistency: the student evaluates with the Cohere chat template
-(student.build_prompt). examples_to_dpo_dataset flattens prompt messages as
-"role: content", BUT with role="" it emits content verbatim — so we pre-apply the
-chat template here and hand it the final string. Training sees exactly the text the
-model sees at inference.
-
-Launched by launch_job.py, which injects --data/--out/--push-repo/--push-path.
-"""
-
-from __future__ import annotations
-
-import argparse
-import json
-import os
-import sys
-
-
-def main() -> None:
-    ap = argparse.ArgumentParser()
-    ap.add_argument("--data", default=None)  # injected by launcher; unused
-    ap.add_argument("--out", required=True)
-    ap.add_argument("--push-repo", default=None)
-    ap.add_argument("--push-path", default=None)
-    ap.add_argument("--pairs-repo", default="HideIron/selfheal-sql-runs")
-    ap.add_argument("--pairs-path", default="heal/baseline-pool600/pairs.jsonl")
-    ap.add_argument("--model", default="CohereLabs/tiny-aya-global")
-    ap.add_argument("--adapter-repo", default="HideIron/tiny-aya-global-sql-selfheal")
-    ap.add_argument("--epochs", type=float, default=3.0)
-    ap.add_argument("--lr", type=float, default=1e-5)
-    ap.add_argument("--beta", type=float, default=0.1)
-    ap.add_argument("--lora-r", type=int, default=16)
-    ap.add_argument("--batch-size", type=int, default=2)
-    ap.add_argument("--grad-accum", type=int, default=8)
-    ap.add_argument("--max-steps", type=int, default=None, help="smoke-test cap")
-    ap.add_argument("--max-prompt-len", type=int, default=1408)
-    ap.add_argument("--rpo-alpha", type=float, default=0.0)
-    ap.add_argument("--algorithm", default="dpo", choices=["dpo", "bco"])
-    ap.add_argument("--replay-run", default=None,
-                    help="BCO only: student run whose CORRECT answers are added as "
-                         "positive (label=True) replay rows against forgetting")
-    ap.add_argument("--no-push-adapter", action="store_true")
-    args = ap.parse_args()
-
-    from huggingface_hub import HfApi, hf_hub_download
-    from transformers import AutoTokenizer
-
-    from agenttune.decide.closed_loop.contracts import TrainingExample
-    from agenttune.decide.closed_loop.retrain_config import RetrainConfig, build_retrainer
-
-    pp = hf_hub_download(args.pairs_repo, args.pairs_path, repo_type="dataset")
-    pairs = [json.loads(l) for l in open(pp, encoding="utf-8")]
-    healed = [p for p in pairs if p.get("healed") and p.get("chosen_sql")]
-    print(f"[dpo] {len(healed)} verified pairs of {len(pairs)} failures", flush=True)
-
-    tok = AutoTokenizer.from_pretrained(args.model)
-
-    def templ(user_content: str) -> str:
-        return tok.apply_chat_template(
-            [{"role": "user", "content": user_content}],
-            add_generation_prompt=True, tokenize=False,
-        )
-
-    examples = [
-        TrainingExample(
-            trajectory_id=p["qid"],
-            original_failure_type=p["status_before"],
-            root_cause=p["root_cause"],
-            # role "" -> _messages_to_text emits content verbatim (no "user: " prefix)
-            prompt=[{"role": "", "content": templ(p["prompt_user_content"])}],
-            chosen=[{"role": "", "content": p["chosen_sql"]}],
-            rejected=[{"role": "", "content": p["rejected_sql"]}],
-        )
-        for p in healed
-    ]
-
-    if args.replay_run:
-        assert args.algorithm == "bco", "success replay needs BCO's independent rows"
-        import re as _re
-
-        def _extract(t):  # minimal extract_sql (fences + first statement)
-            t = _re.sub(r"<think>.*?</think>", "", t, flags=_re.S)
-            m = _re.search(r"```(?:sql)?\s*(.*?)```", t, _re.S | _re.I)
-            if m:
-                t = m.group(1)
-            m = _re.search(r"\b(select|with)\b", t, _re.I)
-            return t[m.start():].split(";")[0].strip() if m else t.strip()
-
-        sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-        from student import PROMPT  # the exact student prompt builder
-
-        split_rows = {}
-        sp = os.path.join(args.data or "/data", "splits", "pool.jsonl")
-        for s in map(json.loads, open(sp, encoding="utf-8")):
-            split_rows[s["qid"]] = s
-        rp2 = hf_hub_download(args.pairs_repo, f"runs/{args.replay_run}/results.jsonl",
-                              repo_type="dataset")
-        n0 = len(examples)
-        for r in map(json.loads, open(rp2, encoding="utf-8")):
-            if r["status"] == "correct" and r["qid"] in split_rows:
-                s = split_rows[r["qid"]]
-                examples.append(TrainingExample(
-                    trajectory_id=r["qid"] + "-replay",
-                    original_failure_type="success_replay", root_cause="none",
-                    prompt=[{"role": "", "content": templ(
-                        PROMPT.format(schema=s["schema"], question=s["question"]))}],
-                    chosen=[{"role": "", "content": _extract(r["raw_output"])}],
-                ))
-        print(f"[dpo] +{len(examples)-n0} success-replay rows (label=True)", flush=True)
-
-    cfg = RetrainConfig(
-        model=args.model,
-        algorithm=args.algorithm,
-        output_dir=os.path.join(args.out, "adapter"),
-        lora_r=args.lora_r,
-        lora_alpha=args.lora_r * 2,
-        lora_dropout=0.05,
-        lora_target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
-                             "gate_proj", "up_proj", "down_proj"],
-        num_train_epochs=args.epochs,
-        per_device_train_batch_size=args.batch_size,
-        gradient_accumulation_steps=args.grad_accum,
-        learning_rate=args.lr,
-        beta=args.beta,
-        max_steps=args.max_steps,
-        extra={
-            "bf16": True,
-            "logging_steps": 2,
-            "save_strategy": "no",
-            "report_to": "none",
-            "gradient_checkpointing": True,
-        },
-    )
-    if args.algorithm == "dpo":
-        # The BCO wrapper instantiates the model itself and rejects model_init_kwargs.
-        cfg.extra["model_init_kwargs"] = {"dtype": "bfloat16"}
-
-    # Cohere vocab is ~256k: full-sequence fp32 logits are ~4 GiB per copy and OOM a
-    # 24 GB card. use_logits_to_keep restricts logits to completion tokens (~60 here).
-    # If this TRL lacks it, shrink sequence caps hard instead.
-    import inspect
-
-    import trl
-
-    cfg_keys = set(inspect.signature(trl.DPOConfig.__init__).parameters)
-    if "use_logits_to_keep" in cfg_keys:
-        cfg.extra.update({"use_logits_to_keep": True,
-                          "max_prompt_length": 2944, "max_length": 3072})
-        print("[dpo] use_logits_to_keep=True (trl", trl.__version__, ")", flush=True)
-    else:
-        cfg.extra.update({"max_prompt_length": args.max_prompt_len,
-                          "max_length": args.max_prompt_len + 128})
-        print("[dpo] trl lacks use_logits_to_keep -> caps", args.max_prompt_len, flush=True)
-    # rpo_alpha adds an NLL term on the chosen completion (DPO+SFT hybrid): the model
-    # must LEARN to produce the verified fix, not merely prefer it. v1 (pure DPO,
-    # 30 steps) failed to move pool accuracy at all — this is the main v2 fix.
-    if args.rpo_alpha and "rpo_alpha" in cfg_keys:
-        cfg.extra["rpo_alpha"] = args.rpo_alpha
-        print("[dpo] rpo_alpha =", args.rpo_alpha, flush=True)
-    elif args.rpo_alpha:
-        print("[dpo] WARNING: trl lacks rpo_alpha; pure DPO", flush=True)
-
-    lens = sorted(len(tok(templ(p["prompt_user_content"]))["input_ids"]) for p in healed)
-    print(f"[dpo] prompt tokens: median={lens[len(lens)//2]} p90={lens[int(len(lens)*.9)]} "
-          f"max={lens[-1]} | cap={cfg.extra['max_prompt_length']}", flush=True)
-
-    trainer = build_retrainer(examples, cfg)
-    stats = trainer.train()
-    print("[dpo] train() returned:", str(stats)[:500], flush=True)
-
-    # Save the adapter explicitly (find the underlying HF trainer/model).
-    os.makedirs(cfg.output_dir, exist_ok=True)
-    inner = None
-    for attr in ("trainer", "_trainer", "dpo_trainer"):
-        inner = getattr(trainer, attr, None)
-        if inner is not None and hasattr(inner, "model"):
-            break
-    model_obj = getattr(inner, "model", None) or getattr(trainer, "model", None)
-    if model_obj is None:
-        raise RuntimeError(
-            f"cannot locate trained model on {type(trainer).__name__}; attrs: "
-            + ", ".join(a for a in dir(trainer) if not a.startswith("__"))[:800]
-        )
-    model_obj.save_pretrained(cfg.output_dir)  # PeftModel -> adapter only
-    print("[dpo] adapter saved:", os.listdir(cfg.output_dir), flush=True)
-
-    with open(os.path.join(args.out, "train_stats.json"), "w", encoding="utf-8") as f:
-        json.dump({"n_pairs": len(healed), "stats": str(stats)[:2000],
-                   "config": {k: v for k, v in vars(args).items()}}, f, indent=2)
-
-    api = HfApi()
-    if not args.no_push_adapter:
-        api.create_repo(args.adapter_repo, private=True, exist_ok=True)
-        api.upload_folder(repo_id=args.adapter_repo, folder_path=cfg.output_dir,
-                          commit_message=f"DPO LoRA ({len(healed)} self-heal pairs)")
-        print("[dpo] adapter pushed:", args.adapter_repo, flush=True)
-    if args.push_repo:
-        api.upload_folder(repo_id=args.push_repo, repo_type="dataset",
-                          folder_path=args.out, path_in_repo=args.push_path or "dpo",
-                          commit_message="dpo training stats")
-
-
-if __name__ == "__main__":
-    main()
+"""DPO retrain via AgentTune's closed-loop retrain path (run inside an HF GPU Job).
+
+Consumes the healer's verified pairs, builds agenttune TrainingExample objects, and
+trains an adapter-only LoRA through agenttune's own
+  examples_to_dpo_dataset -> build_retrainer (create_agentic_trainer) -> train()
+exactly as the closed loop's run_retrain does (we call build_retrainer so we can save
+the adapter explicitly afterwards).
+
+Train/eval consistency: the student evaluates with the Cohere chat template
+(student.build_prompt). examples_to_dpo_dataset flattens prompt messages as
+"role: content", BUT with role="" it emits content verbatim - so we pre-apply the
+chat template here and hand it the final string. Training sees exactly the text the
+model sees at inference.
+
+Launched by launch_job.py, which injects --data/--out/--push-repo/--push-path.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import sys
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--data", default=None)  # injected by launcher; unused
+    ap.add_argument("--out", required=True)
+    ap.add_argument("--push-repo", default=None)
+    ap.add_argument("--push-path", default=None)
+    ap.add_argument("--pairs-repo", default="HideIron/selfheal-sql-runs")
+    ap.add_argument("--pairs-path", default="heal/baseline-pool600/pairs.jsonl")
+    ap.add_argument("--model", default="CohereLabs/tiny-aya-global")
+    ap.add_argument("--adapter-repo", default="HideIron/tiny-aya-global-sql-selfheal")
+    ap.add_argument("--epochs", type=float, default=3.0)
+    ap.add_argument("--lr", type=float, default=1e-5)
+    ap.add_argument("--beta", type=float, default=0.1)
+    ap.add_argument("--lora-r", type=int, default=16)
+    ap.add_argument("--batch-size", type=int, default=2)
+    ap.add_argument("--grad-accum", type=int, default=8)
+    ap.add_argument("--max-steps", type=int, default=None, help="smoke-test cap")
+    ap.add_argument("--max-prompt-len", type=int, default=1408)
+    ap.add_argument("--rpo-alpha", type=float, default=0.0)
+    ap.add_argument("--algorithm", default="dpo", choices=["dpo", "bco"])
+    ap.add_argument("--replay-run", default=None,
+                    help="BCO only: student run whose CORRECT answers are added as "
+                         "positive (label=True) replay rows against forgetting")
+    ap.add_argument("--no-push-adapter", action="store_true")
+    args = ap.parse_args()
+
+    from huggingface_hub import HfApi, hf_hub_download
+    from transformers import AutoTokenizer
+
+    from agenttune.decide.closed_loop.contracts import TrainingExample
+    from agenttune.decide.closed_loop.retrain_config import RetrainConfig, build_retrainer
+
+    pp = hf_hub_download(args.pairs_repo, args.pairs_path, repo_type="dataset")
+    pairs = [json.loads(l) for l in open(pp, encoding="utf-8")]
+    healed = [p for p in pairs if p.get("healed") and p.get("chosen_sql")]
+    print(f"[dpo] {len(healed)} verified pairs of {len(pairs)} failures", flush=True)
+
+    tok = AutoTokenizer.from_pretrained(args.model)
+
+    def templ(user_content: str) -> str:
+        return tok.apply_chat_template(
+            [{"role": "user", "content": user_content}],
+            add_generation_prompt=True, tokenize=False,
+        )
+
+    examples = [
+        TrainingExample(
+            trajectory_id=p["qid"],
+            original_failure_type=p["status_before"],
+            root_cause=p["root_cause"],
+            # role "" -> _messages_to_text emits content verbatim (no "user: " prefix)
+            prompt=[{"role": "", "content": templ(p["prompt_user_content"])}],
+            chosen=[{"role": "", "content": p["chosen_sql"]}],
+            rejected=[{"role": "", "content": p["rejected_sql"]}],
+        )
+        for p in healed
+    ]
+
+    if args.replay_run:
+        assert args.algorithm == "bco", "success replay needs BCO's independent rows"
+        import re as _re
+
+        def _extract(t):  # minimal extract_sql (fences + first statement)
+            t = _re.sub(r"<think>.*?</think>", "", t, flags=_re.S)
+            m = _re.search(r"```(?:sql)?\s*(.*?)```", t, _re.S | _re.I)
+            if m:
+                t = m.group(1)
+            m = _re.search(r"\b(select|with)\b", t, _re.I)
+            return t[m.start():].split(";")[0].strip() if m else t.strip()
+
+        sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+        from student import PROMPT  # the exact student prompt builder
+
+        split_rows = {}
+        sp = os.path.join(args.data or "/data", "splits", "pool.jsonl")
+        for s in map(json.loads, open(sp, encoding="utf-8")):
+            split_rows[s["qid"]] = s
+        rp2 = hf_hub_download(args.pairs_repo, f"runs/{args.replay_run}/results.jsonl",
+                              repo_type="dataset")
+        n0 = len(examples)
+        for r in map(json.loads, open(rp2, encoding="utf-8")):
+            if r["status"] == "correct" and r["qid"] in split_rows:
+                s = split_rows[r["qid"]]
+                examples.append(TrainingExample(
+                    trajectory_id=r["qid"] + "-replay",
+                    original_failure_type="success_replay", root_cause="none",
+                    prompt=[{"role": "", "content": templ(
+                        PROMPT.format(schema=s["schema"], question=s["question"]))}],
+                    chosen=[{"role": "", "content": _extract(r["raw_output"])}],
+                ))
+        print(f"[dpo] +{len(examples)-n0} success-replay rows (label=True)", flush=True)
+
+    cfg = RetrainConfig(
+        model=args.model,
+        algorithm=args.algorithm,
+        output_dir=os.path.join(args.out, "adapter"),
+        lora_r=args.lora_r,
+        lora_alpha=args.lora_r * 2,
+        lora_dropout=0.05,
+        lora_target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
+                             "gate_proj", "up_proj", "down_proj"],
+        num_train_epochs=args.epochs,
+        per_device_train_batch_size=args.batch_size,
+        gradient_accumulation_steps=args.grad_accum,
+        learning_rate=args.lr,
+        beta=args.beta,
+        max_steps=args.max_steps,
+        extra={
+            "bf16": True,
+            "logging_steps": 2,
+            "save_strategy": "no",
+            "report_to": "none",
+            "gradient_checkpointing": True,
+        },
+    )
+    if args.algorithm == "dpo":
+        # The BCO wrapper instantiates the model itself and rejects model_init_kwargs.
+        cfg.extra["model_init_kwargs"] = {"dtype": "bfloat16"}
+
+    # Cohere vocab is ~256k: full-sequence fp32 logits are ~4 GiB per copy and OOM a
+    # 24 GB card. use_logits_to_keep restricts logits to completion tokens (~60 here).
+    # If this TRL lacks it, shrink sequence caps hard instead.
+    import inspect
+
+    import trl
+
+    cfg_keys = set(inspect.signature(trl.DPOConfig.__init__).parameters)
+    if "use_logits_to_keep" in cfg_keys:
+        cfg.extra.update({"use_logits_to_keep": True,
+                          "max_prompt_length": 2944, "max_length": 3072})
+        print("[dpo] use_logits_to_keep=True (trl", trl.__version__, ")", flush=True)
+    else:
+        cfg.extra.update({"max_prompt_length": args.max_prompt_len,
+                          "max_length": args.max_prompt_len + 128})
+        print("[dpo] trl lacks use_logits_to_keep -> caps", args.max_prompt_len, flush=True)
+    # rpo_alpha adds an NLL term on the chosen completion (DPO+SFT hybrid): the model
+    # must LEARN to produce the verified fix, not merely prefer it. v1 (pure DPO,
+    # 30 steps) failed to move pool accuracy at all - this is the main v2 fix.
+    if args.rpo_alpha and "rpo_alpha" in cfg_keys:
+        cfg.extra["rpo_alpha"] = args.rpo_alpha
+        print("[dpo] rpo_alpha =", args.rpo_alpha, flush=True)
+    elif args.rpo_alpha:
+        print("[dpo] WARNING: trl lacks rpo_alpha; pure DPO", flush=True)
+
+    lens = sorted(len(tok(templ(p["prompt_user_content"]))["input_ids"]) for p in healed)
+    print(f"[dpo] prompt tokens: median={lens[len(lens)//2]} p90={lens[int(len(lens)*.9)]} "
+          f"max={lens[-1]} | cap={cfg.extra['max_prompt_length']}", flush=True)
+
+    trainer = build_retrainer(examples, cfg)
+    stats = trainer.train()
+    print("[dpo] train() returned:", str(stats)[:500], flush=True)
+
+    # Save the adapter explicitly (find the underlying HF trainer/model).
+    os.makedirs(cfg.output_dir, exist_ok=True)
+    inner = None
+    for attr in ("trainer", "_trainer", "dpo_trainer"):
+        inner = getattr(trainer, attr, None)
+        if inner is not None and hasattr(inner, "model"):
+            break
+    model_obj = getattr(inner, "model", None) or getattr(trainer, "model", None)
+    if model_obj is None:
+        raise RuntimeError(
+            f"cannot locate trained model on {type(trainer).__name__}; attrs: "
+            + ", ".join(a for a in dir(trainer) if not a.startswith("__"))[:800]
+        )
+    model_obj.save_pretrained(cfg.output_dir)  # PeftModel -> adapter only
+    print("[dpo] adapter saved:", os.listdir(cfg.output_dir), flush=True)
+
+    with open(os.path.join(args.out, "train_stats.json"), "w", encoding="utf-8") as f:
+        json.dump({"n_pairs": len(healed), "stats": str(stats)[:2000],
+                   "config": {k: v for k, v in vars(args).items()}}, f, indent=2)
+
+    api = HfApi()
+    if not args.no_push_adapter:
+        api.create_repo(args.adapter_repo, private=True, exist_ok=True)
+        api.upload_folder(repo_id=args.adapter_repo, folder_path=cfg.output_dir,
+                          commit_message=f"DPO LoRA ({len(healed)} self-heal pairs)")
+        print("[dpo] adapter pushed:", args.adapter_repo, flush=True)
+    if args.push_repo:
+        api.upload_folder(repo_id=args.push_repo, repo_type="dataset",
+                          folder_path=args.out, path_in_repo=args.push_path or "dpo",
+                          commit_message="dpo training stats")
+
+
+if __name__ == "__main__":
+    main()