"""Wide-schema retry on row_count_off failures.

For each row_count_off failure, re-runs the production G pipeline with a

WIDER retrieval budget: schema_top_k=10, fk_hops=2, table_budget=20.

Rationale: row_count_off = the model picked a wrong JOIN structure or

missed a WHERE filter. A common root cause is that the right table was

not in the retrieved schema block (the model can't filter on a column it

hasn't seen). Bumping retrieval budget gives the model more context to

find the missing table / FK chain.

Memory: prior n=200 ablation tied top_k=5↔8 because table_budget=12

saturated. Lifting table_budget=20 is the un-tried regime.

Output is voting-shaped for `merge_voting_rescues.py`.

Usage:

uv run python scripts/run_wide_schema_retry.py \

--baseline eval/reports/2026-05-13/hybrid+multi-vote+critique+selfcon-v5.json \

--out eval/reports/2026-05-13/wide-schema-retry.json

uv run python scripts/run_wide_schema_retry.py \

--baseline eval/reports/2026-05-22/v20-kimi-k2-thinking-merged.json \

--out eval/reports/2026-05-22/wide-schema-qid207.json --only-qids 207

"""

from __future__ import annotations

import argparse

import json

import sys

import time

from pathlib import Path

from typing import Any

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.execution.runner import execute_validated

from nl_sql.llm.cache import CachingEmbeddingProvider, CachingLLMProvider

from nl_sql.llm.providers.mistral import MistralProvider

from nl_sql.schema_index.indexer import SchemaIndex

def _is_row_count_off(r: dict[str, Any]) -> bool:

if r.get("match") or r.get("error_kind"):

return False

gc = r.get("gold_row_count") or 0

pc = r.get("pred_row_count") or 0

return gc != pc

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("--schema-top-k", type=int, default=10)

p.add_argument("--fk-hops", type=int, default=2)

p.add_argument("--table-budget", type=int, default=20)

p.add_argument(

"--only-qids",

default="",

help="comma-separated row_count_off failure qids to retry exactly, preserving argument order",

)

p.add_argument("--out", type=Path, required=True)

args = p.parse_args()

baseline = json.loads(args.baseline.read_text(encoding="utf-8"))

fails = [r for r in baseline["records"] if _is_row_count_off(r)]

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 row_count_off failures: {missing_qids}", file=sys.stderr

)

return 3

fails = [fails_by_qid[qid] for qid in only_qids]

print(

f"[info] {len(fails)} row_count_off fails to retry with "

f"top_k={args.schema_top_k}, hops={args.fk_hops}, budget={args.table_budget}",

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()

mistral = MistralProvider(api_key=settings.mistral_api_key, gen_model="codestral-latest")

sql_prov = CachingLLMProvider(mistral, cache_dir=settings.llm_cache_dir)

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=sql_prov,

schema_index=idx,

registry=registry,

schema_top_k=args.schema_top_k,

fk_hops=args.fk_hops,

table_budget=args.table_budget,

fewshot_top_k=3,

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",

)

except Exception as exc:

print(f"[{i:3d}/{len(fails)}] qid={qid} EXC: {str(exc)[:120]}", file=sys.stderr)

continue

elapsed = (time.perf_counter() - t0) * 1000.0

alt_sql = alt.sql or ""

alt_rows: list[Any] = []

try:

outcome = execute_validated(

engine,

alt_sql,

dialect="sqlite",

statement_timeout_ms=30_000,

row_cap=10_000,

)

if outcome.result:

alt_rows = list(outcome.result.rows)

except Exception:

pass

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

tag = "RESCUE"

elif br.get("match") and not alt_match:

regressed += 1

tag = "regression"

else:

same += 1

tag = "same"

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,

"vote_match": alt_match,

"vote_source": "wide-schema",

"elapsed_ms": elapsed,

}

)

print(

f"[{i:3d}/{len(fails)}] qid={qid} {ex.difficulty:11s} {tag} ({elapsed:.0f}ms)",

file=sys.stderr,

)

finally:

engine.dispose()

print("\n=== wide-schema retry summary ===", file=sys.stderr)

print(

f" cases: {len(records)} rescued: {rescued} regressed: {regressed} same: {same}",

file=sys.stderr,

)

args.out.parent.mkdir(parents=True, exist_ok=True)

args.out.write_text(

json.dumps(

{

"alt_model": f"codestral+wide-schema(top_k={args.schema_top_k},hops={args.fk_hops},budget={args.table_budget})+critique",

"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())