"""Targeted grounded-critique retry on baseline failures. Re-runs the production G pipeline (codestral + fewshot + verify-retry) BUT with `enable_grounded_critique=True`, ONLY on the questions where the multi-vote baseline still failed. The grounded_critique node detects row-shape mismatches (e.g. question asks "how many X" expecting 1 row but SQL returns 12) and re-prompts with the shape feedback as a hint. Output is a voting-shaped report so `merge_voting_rescues.py` can fold the rescues back into the multi-vote baseline. Usage: uv run python scripts/run_critique_retry.py \ --baseline eval/reports/2026-05-13/hybrid+multi-vote-v3.json \ --bird-root data/bird_mini_dev/MINIDEV \ --out eval/reports/2026-05-13/critique-retry.json uv run python scripts/run_critique_retry.py \ --baseline eval/reports/2026-05-22/v20-kimi-k2-thinking-merged.json \ --out eval/reports/2026-05-22/critique-qid1399.json --only-qids 1399 """ from __future__ import annotations import argparse import json import sys import time from pathlib import Path from nl_sql.agent.graph import PipelineConfig, build_pipeline, run_pipeline from nl_sql.config import get_settings from nl_sql.db.registry import get_default_registry from nl_sql.eval.dataset import load_bird_mini_dev from nl_sql.eval.metrics.execution_accuracy import compare_results from nl_sql.eval.runner import _compose_question, _execute_gold from nl_sql.llm.cache import CachingEmbeddingProvider, CachingLLMProvider from nl_sql.llm.providers.groq import GroqProvider from nl_sql.llm.providers.mistral import MistralProvider from nl_sql.schema_index.indexer import SchemaIndex def main() -> int: p = argparse.ArgumentParser(description=__doc__) p.add_argument("--baseline", type=Path, required=True) p.add_argument("--bird-root", type=Path, default=Path("data/bird_mini_dev/MINIDEV")) p.add_argument("--out", type=Path, required=True) p.add_argument("--max-cases", type=int, default=200) p.add_argument( "--only-qids", default="", help="comma-separated baseline failure qids to retry exactly, preserving argument order", ) p.add_argument( "--fewshot-top-k", type=int, default=3, help="PipelineConfig.fewshot_top_k (default 3 = G prod). " "Use 5 for P2.B selective expansion experiment.", ) p.add_argument( "--gen-model", type=str, default="codestral-latest", help="Mistral gen model id (default codestral-latest = G prod). " "Use mistral-large-latest for cross-model voting on residue.", ) p.add_argument( "--sleep-between", type=float, default=0.0, help="Sleep N seconds between cases — required for mistral-large " "on free tier (rate-limited ~2 req/s).", ) p.add_argument( "--provider", type=str, choices=("mistral", "groq"), default="mistral", help="SQL provider: mistral (default, uses --gen-model) or groq " "(uses --gen-model as Groq model id, e.g. qwen/qwen3-32b).", ) p.add_argument( "--base-url", type=str, default=None, help="Override OpenAI-compatible base_url for the SQL provider. " "Use with --provider groq to redirect to OpenRouter " "(https://openrouter.ai/api/v1) or Gemini OpenAI compat " "(https://generativelanguage.googleapis.com/v1beta/openai). " "Requires GROQ_API_KEY env to actually hold the alt-provider key.", ) p.add_argument( "--api-key", type=str, default=None, help="Override API key for the SQL provider (otherwise read from " "settings.groq_api_key / settings.mistral_api_key).", ) args = p.parse_args() baseline = json.loads(args.baseline.read_text(encoding="utf-8")) fails = [r for r in baseline["records"] if not r.get("match")] try: only_qids = [int(x) for x in args.only_qids.split(",") if x.strip()] except ValueError: print("[error] invalid --only-qids: expected comma-separated integers", file=sys.stderr) return 3 if only_qids: fails_by_qid = {int(r["question_id"]): r for r in fails} missing_qids = [qid for qid in only_qids if qid not in fails_by_qid] if missing_qids: print(f"[error] qids not found in baseline failures: {missing_qids}", file=sys.stderr) return 3 fails = [fails_by_qid[qid] for qid in only_qids] fails = fails[: args.max_cases] print(f"[info] {len(fails)} failures to retry with grounded_critique", file=sys.stderr) settings = get_settings() examples = {e.question_id: e for e in load_bird_mini_dev(args.bird_root)} registry = get_default_registry() if args.provider == "mistral": gen_provider = MistralProvider(api_key=settings.mistral_api_key, gen_model=args.gen_model) else: groq_kwargs = { "api_key": args.api_key or settings.groq_api_key, "model": args.gen_model, } if args.base_url: groq_kwargs["base_url"] = args.base_url gen_provider = GroqProvider(**groq_kwargs) sql_prov = CachingLLMProvider(gen_provider, cache_dir=settings.llm_cache_dir) expl_prov = sql_prov # same provider for explain emb = CachingEmbeddingProvider( MistralProvider(api_key=settings.mistral_api_key), cache_dir=settings.llm_cache_dir ) idx = SchemaIndex(persist_dir="chroma_data", embedder=emb) cfg = PipelineConfig( sql_provider=sql_prov, explain_provider=expl_prov, schema_index=idx, registry=registry, fewshot_top_k=args.fewshot_top_k, sort_schema_block=True, cross_db_fewshot=True, verify_retry_on_empty=True, enable_grounded_critique=True, ) pipeline = build_pipeline(cfg) records = [] rescued = 0 regressed = 0 same = 0 for i, br in enumerate(fails, 1): qid = br["question_id"] ex = examples.get(qid) if ex is None: continue spec = registry.get(ex.registry_db_id) engine = spec.make_engine() try: t0 = time.perf_counter() try: alt = run_pipeline( pipeline, question=_compose_question(ex), db_id=ex.registry_db_id, dialect="sqlite", verify_retry_on_empty=True, ) except Exception as exc: print(f"[{i:3d}/{len(fails)}] EXC qid={qid}: {exc}", file=sys.stderr) continue elapsed = (time.perf_counter() - t0) * 1000.0 # Execute alt's pred against the DB and compare with gold. alt_rows = [] if alt.outcome and alt.outcome.result: alt_rows = list(alt.outcome.result.rows) try: gold_rows, _ = _execute_gold( engine, ex.sql, statement_timeout_ms=30_000, row_cap=10_000 ) except Exception: gold_rows = [] alt_cmp = compare_results(gold_rows, alt_rows, gold_sql=ex.sql) alt_match = bool(alt_cmp.match) if alt_match and not br.get("match"): rescued += 1 elif br.get("match") and not alt_match: regressed += 1 else: same += 1 records.append( { "question_id": qid, "db_id": ex.db_id, "difficulty": ex.difficulty, "question": ex.question, "gold_sql": ex.sql, "baseline_pred": br["pred_sql"], "alt_pred": alt.sql, "alt_confidence": getattr(alt, "confidence", None), "baseline_match": bool(br.get("match")), "alt_match": alt_match, # Shape: rescue is what merge_voting_rescues.py looks for. "vote_match": alt_match, "vote_source": "critique-retry", "elapsed_ms": elapsed, } ) tag = ( "RESCUE" if (alt_match and not br.get("match")) else ("regression" if (br.get("match") and not alt_match) else "same") ) print( f"[{i:3d}/{len(fails)}] qid={qid} {ex.difficulty:11s} {tag} ({elapsed:.0f}ms)", file=sys.stderr, ) finally: engine.dispose() if args.sleep_between > 0: time.sleep(args.sleep_between) print("\n=== critique-retry summary ===", file=sys.stderr) print(f" cases: {len(records)}", file=sys.stderr) print(f" rescued: {rescued}", file=sys.stderr) print(f" regressed: {regressed}", file=sys.stderr) print(f" same: {same}", file=sys.stderr) args.out.parent.mkdir(parents=True, exist_ok=True) args.out.write_text( json.dumps( { "alt_model": f"{args.provider}:{args.gen_model}+grounded_critique+fewshot{args.fewshot_top_k}", "summary": { "voted_better": rescued, "voted_worse": regressed, "voted_same": same, }, "records": records, }, indent=2, ), encoding="utf-8", ) return 0 if __name__ == "__main__": raise SystemExit(main())