"""Wider self-consistency POC: 2 prompt variants x 4 temps on v29 residue qids. Hypothesis: current config F (4 temps x 1 prompt) converges on a single shape; adding a BIRD-shape-hint prompt variant introduces alternative aggregation/sort patterns that residue qids need (LIMIT vs WHERE=MAX, AVG vs CAST(SUM)/COUNT, date-format conventions). Standalone -- bypasses LangGraph. For each residue qid: 1. Build context via retrieve_context (same as production C config). 2. Generate 8 candidates (2 variants x 4 temps). 3. Execute each on the live db. 4. Cluster by fingerprint_rows (existing eval.self_consistency helper). 5. Pick plurality cluster; compare winner vs gold. POC scope: 3 BIRD-shape-friendly residue qids first. If lift detected -> scale. """ from __future__ import annotations import argparse import json import sys from collections import defaultdict from pathlib import Path from typing import Any import chromadb from nl_sql.agent.nodes._support import parse_generate_sql_output, render_schema_block from nl_sql.agent.prompts import load_prompt from nl_sql.config import get_settings from nl_sql.db.registry import get_default_registry from nl_sql.eval.dataset import DEFAULT_BIRD_ROOT, load_bird_mini_dev from nl_sql.eval.metrics.execution_accuracy import safe_compare_pred from nl_sql.eval.runner import _execute_gold_with_status from nl_sql.eval.self_consistency import fingerprint_rows from nl_sql.execution.runner import execute_validated from nl_sql.llm.cache import CachingEmbeddingProvider, CachingLLMProvider from nl_sql.llm.providers import build_provider from nl_sql.llm.providers.base import GenerateRequest from nl_sql.schema_index.indexer import SchemaIndex from nl_sql.schema_index.retriever import retrieve_context BIRD_SHAPE_RULES = """ # BIRD-style shape conventions (apply when relevant to the question) These are common shape patterns observed in BIRD gold SQL; if your default choice does not fit one of them, consider the alternative. - "Which/who has the highest/lowest/most X" → BIRD gold often uses `WHERE col = (SELECT MAX(col) FROM ...)` rather than `ORDER BY col DESC LIMIT 1`. Prefer the WHERE=MAX subquery shape unless the question explicitly says "top 1" or "first". - "Average of average X" / "Mean X" in BIRD context → prefer `CAST(SUM(X) AS REAL) / COUNT(*)` over `AVG(X)`. BIRD gold rarely uses AVG(). - "After Y/M/D" / "before Y/M/D" date filters → match the exact format stored in the column. If samples show 'YYYY-MM-DD' literal, use `date_col > 'Y-M-D'` directly (no strftime). If samples show numeric year, cast accordingly. - "Rank N" / "in position N" / "Nth place" → BIRD gold uses `WHERE rank_col = N` rather than `ORDER BY rank_col LIMIT N`. Returns all ties; the LIMIT version silently drops them. - "List all X with maximum/minimum Y" → BIRD gold uses `WHERE Y = (SELECT MAX/MIN(Y))` to return all ties. Do NOT use `ORDER BY Y DESC LIMIT 1` if the question implies tie inclusion. - "Highest scoring" / "best" in european_football_2: BIRD gold occasionally treats lower numeric values as "higher rank" (positional inversion). Consider both ASC and DESC sort orders when the column semantics are ambiguous from the schema. """ def _build_prompt( *, variant: str, context: Any, question: str, dialect: str = "sqlite", ) -> str: """Build the full prompt for a given variant.""" schema_text = render_schema_block(context, sort_alphabetically=True) base = load_prompt( "generate_sql", dialect=dialect, schema_block=schema_text, fewshot_block="", plan_block="(no plan — generate SQL directly from question)", question=question, ) if variant == "bird_shape": # Splice BIRD-shape rules just before the JSON output contract so the # model sees them before formulating SQL. marker = "# Output contract" if marker in base: head, tail = base.split(marker, 1) return head + BIRD_SHAPE_RULES + "\n" + marker + tail return base + "\n" + BIRD_SHAPE_RULES return base def _run_one_qid( *, example: Any, schema_index: SchemaIndex, registry: Any, provider: Any, variants: tuple[str, ...], temperatures: tuple[float, ...], ) -> dict[str, Any]: """Generate 8 candidates, execute, cluster, return winner + diagnostics.""" bundle = retrieve_context( schema_index, example.question, db_id=example.registry_db_id, schema_top_k=5, fewshot_top_k=0, fk_hops=1, table_budget=12, primary_sample_size=3, extended_sample_size=0, cross_db_fewshot=False, ) engine = registry.get(example.registry_db_id).make_engine() try: gold_rows, _gold_cols, gold_failed = _execute_gold_with_status( engine, example.sql, statement_timeout_ms=60_000, row_cap=10_000 ) candidates: list[dict[str, Any]] = [] for variant in variants: prompt = _build_prompt(variant=variant, context=bundle, question=example.question) for temp in temperatures: try: response = provider.generate( GenerateRequest(prompt=prompt, max_tokens=1024, temperature=temp) ) except Exception as exc: candidates.append( { "variant": variant, "temperature": temp, "sql": "", "rows": [], "fingerprint": None, "executed": False, "confidence": 0.0, "error": f"provider: {exc!s}"[:200], } ) continue parsed = parse_generate_sql_output(response.text) if not parsed.sql: candidates.append( { "variant": variant, "temperature": temp, "sql": "", "rows": [], "fingerprint": None, "executed": False, "confidence": parsed.confidence, "error": "parse: empty sql", } ) continue outcome = execute_validated( engine, parsed.sql, dialect="sqlite", statement_timeout_ms=60_000, row_cap=10_000, ) if outcome.ok and outcome.result is not None: rows = list(outcome.result.rows) fp = fingerprint_rows(rows) candidates.append( { "variant": variant, "temperature": temp, "sql": parsed.sql, "rows": rows[:5], "row_count": len(rows), "fingerprint": fp, "executed": True, "confidence": parsed.confidence, "error": "", } ) else: candidates.append( { "variant": variant, "temperature": temp, "sql": parsed.sql, "rows": [], "fingerprint": None, "executed": False, "confidence": parsed.confidence, "error": f"exec: {outcome.error_kind.value if outcome.error_kind else 'unknown'}: {outcome.error_message[:200]}", } ) # Cluster by fingerprint. clusters: dict[str, list[dict[str, Any]]] = defaultdict(list) for c in candidates: if c["fingerprint"] is not None: clusters[c["fingerprint"]].append(c) winner: dict[str, Any] | None = None cluster_summary: list[dict[str, Any]] = [] if clusters: ranked = sorted( clusters.items(), key=lambda kv: ( -len(kv[1]), -max(m["confidence"] for m in kv[1]), min(m["temperature"] for m in kv[1]), ), ) for fp, members in ranked: cluster_summary.append( { "fingerprint": fp[:16], "size": len(members), "row_count": members[0].get("row_count", 0), "variants": sorted({m["variant"] for m in members}), "temps": sorted({m["temperature"] for m in members}), "sample_sql": members[0]["sql"][:200], } ) _winner_cluster_fp, winner_members = ranked[0] winner = max( winner_members, key=lambda c: (c["confidence"], -c["temperature"]), ) # Compare winner vs gold. if winner is None: comparison = safe_compare_pred( gold_rows, [], gold_sql=example.sql, pred_failed=True, gold_failed=gold_failed ) else: comparison = safe_compare_pred( gold_rows, [ tuple(r) if not isinstance(r, tuple) else r for r in ( # winner rows is truncated to 5 in candidates dict for display, # re-execute to get full rowset [] ) ], gold_sql=example.sql, pred_failed=False, gold_failed=gold_failed, ) # Re-execute winner SQL fully to get true rows for comparison. outcome = execute_validated( engine, winner["sql"], dialect="sqlite", statement_timeout_ms=60_000, row_cap=10_000, ) pred_rows = list(outcome.result.rows) if outcome.ok and outcome.result else [] comparison = safe_compare_pred( gold_rows, pred_rows, gold_sql=example.sql, pred_failed=not outcome.ok, gold_failed=gold_failed, ) return { "qid": example.question_id, "db_id": example.registry_db_id, "difficulty": example.difficulty, "question": example.question, "gold_sql": example.sql, "gold_failed": gold_failed, "gold_rows_count": len(gold_rows), "candidates_total": len(candidates), "candidates_executed": sum(1 for c in candidates if c["executed"]), "clusters": cluster_summary, "winner_sql": winner["sql"] if winner else "", "winner_variant": winner["variant"] if winner else None, "winner_temp": winner["temperature"] if winner else None, "winner_confidence": winner["confidence"] if winner else 0.0, "match": comparison.match, "match_reason": comparison.reason if hasattr(comparison, "reason") else "", "all_candidates": candidates, } finally: engine.dispose() def main(argv: list[str] | None = None) -> int: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument( "--qids", default="", help="comma-separated qids to run; default: all v29 residue", ) parser.add_argument( "--baseline", default="eval/reports/2026-05-24/v29-v28-plus-p3f-q1275-merged.json", help="v29 merged baseline to source residue qids", ) parser.add_argument( "--temps", default="0.2,0.4,0.6,0.8", help="comma-separated sampling temperatures", ) parser.add_argument( "--variants", default="default,bird_shape", help="comma-separated prompt variants", ) parser.add_argument( "--out", default="eval/reports/2026-05-25/wider_sc_poc.json", help="output JSON path", ) parser.add_argument("--persist", default="chroma_data", help="chroma persist directory") parser.add_argument("--bird-root", default=str(DEFAULT_BIRD_ROOT), help="MINIDEV/ root") args = parser.parse_args(argv) # Load residue qids. baseline_path = Path(args.baseline) if not baseline_path.is_file(): print(f"[error] baseline not found: {baseline_path}", file=sys.stderr) return 2 baseline_data = json.loads(baseline_path.read_text(encoding="utf-8")) residue_qids = [r["question_id"] for r in baseline_data["records"] if not r["match"]] if args.qids: residue_qids = [int(q) for q in args.qids.split(",") if q.strip()] print(f"[info] residue qids: {residue_qids}") # Load BIRD examples. all_examples = load_bird_mini_dev(Path(args.bird_root)) by_qid = {e.question_id: e for e in all_examples} sample = [by_qid[q] for q in residue_qids if q in by_qid] missing = [q for q in residue_qids if q not in by_qid] if missing: print(f"[warn] qids not found in MINIDEV: {missing}", file=sys.stderr) print(f"[info] running on {len(sample)} qids") # Setup providers + index + registry. settings = get_settings() raw = build_provider("mistral", settings=settings) provider = CachingLLMProvider( raw, cache_dir=settings.llm_cache_dir, size_limit_gb=settings.llm_cache_size_limit_gb ) print(f"[info] provider: mistral (model={raw.model}); cache: {settings.llm_cache_dir}") persist_dir = Path(args.persist) if not persist_dir.is_dir(): print(f"[error] chroma persist dir not found: {persist_dir}", file=sys.stderr) return 2 embed_provider_raw = build_provider("mistral", settings=settings) embed_provider = CachingEmbeddingProvider( embed_provider_raw, cache_dir=settings.llm_cache_dir, size_limit_gb=settings.llm_cache_size_limit_gb, ) client = chromadb.PersistentClient(path=str(persist_dir)) schema_index = SchemaIndex(persist_dir, embedder=embed_provider, client=client) registry = get_default_registry() variants = tuple(v.strip() for v in args.variants.split(",") if v.strip()) temperatures = tuple(float(t) for t in args.temps.split(",") if t.strip()) print( f"[info] variants={variants} x temps={temperatures} = {len(variants) * len(temperatures)} candidates/qid" ) results = [] for idx, ex in enumerate(sample, start=1): print( f"[{idx:>2}/{len(sample)}] qid={ex.question_id} db={ex.registry_db_id} — {ex.question[:80]}" ) try: res = _run_one_qid( example=ex, schema_index=schema_index, registry=registry, provider=provider, variants=variants, temperatures=temperatures, ) except Exception as exc: print(f" [error] {exc!r}") res = { "qid": ex.question_id, "db_id": ex.registry_db_id, "difficulty": ex.difficulty, "question": ex.question, "error": repr(exc), "match": False, } results.append(res) flag = "OK " if res.get("match") else "MISS" winner_var = res.get("winner_variant", "?") n_clusters = len(res.get("clusters", [])) print(f" {flag} | clusters={n_clusters} | winner_variant={winner_var}") out_path = Path(args.out) out_path.parent.mkdir(parents=True, exist_ok=True) out_path.write_text( json.dumps( { "baseline": str(baseline_path), "variants": list(variants), "temperatures": list(temperatures), "total_qids": len(sample), "matches": sum(1 for r in results if r.get("match")), "records": results, }, ensure_ascii=False, indent=2, default=str, ), encoding="utf-8", ) matches = sum(1 for r in results if r.get("match")) print( f"\n[summary] {matches}/{len(results)} matches on residue ({matches / len(results) * 100:.1f}% if N>0)" ) print(f"[summary] saved: {out_path}") return 0 if __name__ == "__main__": raise SystemExit(main())