scripts/eval_demo.py

"""Evaluate the pipeline on the curated demo benchmark.

Unlike `eval_baseline.py`, which measures BIRD Mini-Dev (research baseline
~50% ceiling), this script measures a *product workload*: 30 realistic
business questions on Chinook where we target ≥90% Execution Accuracy.

Usage:
    uv run python scripts/eval_demo.py
    uv run python scripts/eval_demo.py --benchmark eval/demo_benchmark.json
"""

from __future__ import annotations

import argparse
import json
import sys
import time
from collections import defaultdict
from pathlib import Path
from typing import Any

import chromadb

from nl_sql.agent import PipelineConfig, build_pipeline, run_pipeline
from nl_sql.config import get_settings
from nl_sql.db.connection import Dialect, execute_readonly
from nl_sql.db.registry import get_default_registry
from nl_sql.eval.metrics.execution_accuracy import compare_results
from nl_sql.llm.cache import CachingEmbeddingProvider, CachingLLMProvider
from nl_sql.llm.providers.mistral import MistralProvider
from nl_sql.schema_index.indexer import SchemaIndex

DEFAULT_BENCHMARK = Path(__file__).parent.parent / "eval" / "demo_benchmark.json"


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--benchmark",
        type=Path,
        default=DEFAULT_BENCHMARK,
        help=f"path to benchmark JSON (default: {DEFAULT_BENCHMARK})",
    )
    parser.add_argument("--persist", default="chroma_data")
    parser.add_argument("--no-cache", action="store_true")
    parser.add_argument(
        "--report",
        type=Path,
        help="optional output JSON path for the full per-question report",
    )
    args = parser.parse_args(argv)

    settings = get_settings()
    if not settings.mistral_api_key:
        print("[error] MISTRAL_API_KEY not set in .env", file=sys.stderr)
        return 2

    bench = json.loads(args.benchmark.read_text(encoding="utf-8"))
    db_id = bench["db_id"]
    dialect: Dialect = bench.get("dialect", "sqlite")
    questions = bench["questions"]
    print(f"[info] benchmark: {bench['name']} ({len(questions)} questions on {db_id!r})")

    persist = Path(args.persist)
    if not persist.is_dir():
        print(f"[error] index not found at {persist}; run scripts/build_index.py first")
        return 3

    client = chromadb.PersistentClient(path=str(persist))
    raw_embed = MistralProvider(
        api_key=settings.mistral_api_key,
        gen_model=settings.mistral_gen_model,
        embed_model=settings.mistral_embed_model,
        base_url=settings.mistral_base_url,
    )
    raw_sql = MistralProvider(
        api_key=settings.mistral_api_key,
        gen_model=settings.mistral_gen_model,
        embed_model=settings.mistral_embed_model,
        base_url=settings.mistral_base_url,
    )
    raw_explain = MistralProvider(
        api_key=settings.mistral_api_key,
        gen_model="mistral-large-latest",
        embed_model=settings.mistral_embed_model,
        base_url=settings.mistral_base_url,
    )
    embedder = (
        raw_embed
        if args.no_cache
        else CachingEmbeddingProvider(
            raw_embed,
            cache_dir=settings.llm_cache_dir,
            size_limit_gb=settings.llm_cache_size_limit_gb,
        )
    )
    sql_provider = (
        raw_sql
        if args.no_cache
        else CachingLLMProvider(
            raw_sql,
            cache_dir=settings.llm_cache_dir,
            size_limit_gb=settings.llm_cache_size_limit_gb,
        )
    )
    explain_provider = (
        raw_explain
        if args.no_cache
        else CachingLLMProvider(
            raw_explain,
            cache_dir=settings.llm_cache_dir,
            size_limit_gb=settings.llm_cache_size_limit_gb,
        )
    )

    index = SchemaIndex(persist_dir=persist, embedder=embedder, client=client)
    registry = get_default_registry()
    pipeline = build_pipeline(
        PipelineConfig(
            sql_provider=sql_provider,
            explain_provider=explain_provider,
            schema_index=index,
            registry=registry,
            schema_top_k=5,
            fewshot_top_k=0,
            fk_hops=1,
            table_budget=12,
            sort_schema_block=True,
            primary_sample_size=3,
        )
    )

    spec = registry.get(db_id)
    gold_engine = spec.make_engine()

    records: list[dict[str, Any]] = []
    started_all = time.perf_counter()
    try:
        for i, q in enumerate(questions, start=1):
            t0 = time.perf_counter()
            try:
                result = run_pipeline(
                    pipeline, question=q["question"], db_id=db_id, dialect=dialect
                )
            except Exception as exc:
                elapsed = (time.perf_counter() - t0) * 1000
                records.append(
                    {
                        "id": q["id"],
                        "category": q.get("category", ""),
                        "difficulty": q.get("difficulty", ""),
                        "split": q.get("split", "dev"),
                        "question": q["question"],
                        "gold_sql": q["gold_sql"],
                        "pred_sql": "",
                        "match": False,
                        "reason": f"pipeline raised: {exc!r}",
                        "latency_ms": elapsed,
                    }
                )
                print(f"  [{i:>2}/{len(questions)}] EXCEPTION {q['id']}: {exc}")
                continue

            with execute_readonly(
                gold_engine, q["gold_sql"], statement_timeout_ms=30_000, row_cap=10_000
            ) as gold:
                gold_rows = list(gold.rows)

            if result.outcome is not None and result.outcome.result is not None:
                cmp = compare_results(gold_rows, result.outcome.result.rows, gold_sql=q["gold_sql"])
                match = cmp.match
                reason = cmp.reason
            else:
                match = False
                reason = (
                    f"pred failed: {result.error_kind.value if result.error_kind else 'unknown'}"
                )

            elapsed = (time.perf_counter() - t0) * 1000
            flag = "OK  " if match else "MISS"
            print(
                f"  [{i:>2}/{len(questions)}] {flag} ({elapsed:5.0f}ms) {q['id']} — {q['question'][:70]}"
            )
            if not match:
                print(f"        gold: {q['gold_sql'][:140]}")
                print(f"        pred: {result.sql[:140]}")
                print(f"        why:  {reason}")
            records.append(
                {
                    "id": q["id"],
                    "category": q.get("category", ""),
                    "difficulty": q.get("difficulty", ""),
                    "split": q.get("split", "dev"),
                    "question": q["question"],
                    "gold_sql": q["gold_sql"],
                    "pred_sql": result.sql,
                    "match": match,
                    "reason": reason,
                    "latency_ms": elapsed,
                }
            )
    finally:
        gold_engine.dispose()

    elapsed_total = time.perf_counter() - started_all
    matches = sum(1 for r in records if r["match"])
    ea = matches / len(records) if records else 0.0
    print()
    print("=" * 78)
    print(f"Demo benchmark: {bench['name']}")
    print(f"DB:             {db_id} ({dialect})")
    print(f"Questions:      {len(records)}")
    print(f"Match:          {matches}/{len(records)} = {ea * 100:.1f}%")
    by_cat: dict[str, list[bool]] = defaultdict(list)
    by_diff: dict[str, list[bool]] = defaultdict(list)
    by_split: dict[str, list[bool]] = defaultdict(list)
    for r in records:
        by_cat[r["category"]].append(r["match"])
        by_diff[r["difficulty"]].append(r["match"])
        by_split[r["split"]].append(r["match"])
    print("per category:")
    for cat, ms in sorted(by_cat.items()):
        print(f"  {cat:14s} {sum(ms):>2}/{len(ms):<2} ({sum(ms) / len(ms) * 100:5.1f}%)")
    print("per difficulty:")
    for d, ms in sorted(by_diff.items()):
        print(f"  {d:14s} {sum(ms):>2}/{len(ms):<2} ({sum(ms) / len(ms) * 100:5.1f}%)")
    print("per split:")
    for s, ms in sorted(by_split.items()):
        print(f"  {s:14s} {sum(ms):>2}/{len(ms):<2} ({sum(ms) / len(ms) * 100:5.1f}%)")
    print(f"Wall time: {elapsed_total:.1f}s")

    if args.report:
        args.report.parent.mkdir(parents=True, exist_ok=True)
        args.report.write_text(
            json.dumps(
                {
                    "benchmark": bench["name"],
                    "db_id": db_id,
                    "dialect": dialect,
                    "n": len(records),
                    "matches": matches,
                    "ea": ea,
                    "records": records,
                },
                indent=2,
                ensure_ascii=False,
            ),
            encoding="utf-8",
        )
        print(f"[report] {args.report}")

    return 0 if ea >= 0.9 else 1


if __name__ == "__main__":
    sys.exit(main())