arborist/bench/legacy_vs_providence_bench.py

"""Head-to-head: arborist query (default, providence_query path) vs
arborist query --legacy (legacy 2000-line query() path).

For each themed bench subset, sample N questions. Run each through
both CLI paths. Compare primary source + audit_mode + n_verified.

Output: bench/legacy_vs_providence_results/<utc>.jsonl + a stderr
summary table per theme showing where the two paths agree / disagree.

The point: which themed subsets regress when the CLI default flips
from legacy to providence_query (#000072 Phase 2 step 3). Drives
porting priority for the 12 missing pre/post gates.

Uses --burn on both paths so cache misses are forced (we measure
the fresh-inference quality, not warm-cache replay).
"""
from __future__ import annotations

import argparse
import datetime as _dt
import json
import os
import subprocess
import sys
import time
from pathlib import Path

REPO = Path(__file__).resolve().parent.parent
ARBORIST = REPO / ".venv" / "bin" / "arborist"
DEFAULT_SHARDS = Path.home() / ".arborist" / "shards"
DEFAULT_QA_DB = Path.home() / ".arborist" / "qa.db"

# Themed subsets to sample. Each names a gate the legacy pipeline has
# that providence_query DOESN'T port yet — regression here points at
# which gate to port next.
THEMES = [
    ("es", "crosslang sandwich (es→en)"),
    ("accent", "_accent_fold_variants in _title_query_tokens"),
    ("hyphen", "_hyphen_fold_variants"),
    ("honorific", "_honorific_fold_variants"),
    ("brit", "_brit_fold_variants"),
    ("numeral", "_numeral_fold_variants"),
    ("quantifier_subset", "quantifier preflight guard"),
    ("metacog_subset", "metacog preflight"),
    ("warrant_chain_probe", "warrant-resolver post-step"),
]


def _load_fixture(path: Path, n: int) -> list[str]:
    out = []
    for line in path.read_text().splitlines():
        s = line.strip()
        if not s or s.startswith("#"):
            continue
        out.append(s)
        if len(out) >= n:
            break
    return out


def _run_cli(cmd: list[str], *, timeout_s: int) -> dict:
    t0 = time.time()
    try:
        proc = subprocess.run(
            cmd, capture_output=True, text=True, timeout=timeout_s,
            env={**os.environ, "ARBORIST_PROGRESS": "0"},
        )
    except subprocess.TimeoutExpired:
        return {"_error": f"timeout {timeout_s}s",
                "_elapsed_s": timeout_s}
    dt = time.time() - t0
    if proc.returncode != 0:
        return {"_error": f"exit {proc.returncode}",
                "_stderr": proc.stderr[-300:],
                "_elapsed_s": round(dt, 2)}
    try:
        d = json.loads(proc.stdout)
    except json.JSONDecodeError as e:
        return {"_error": f"parse: {e}",
                "_stdout_head": proc.stdout[:300],
                "_elapsed_s": round(dt, 2)}
    d["_elapsed_s"] = round(dt, 2)
    return d


def _query(q: str, *, shards_dir: Path, qa_db: Path, legacy: bool) -> dict:
    cmd = [
        str(ARBORIST), "--shards-dir", str(shards_dir),
        "query", q,
        "--qa-db", str(qa_db),
        "--burn", "--json",
        "--top-k", "4",
    ]
    if legacy:
        cmd.append("--legacy")
    return _run_cli(cmd, timeout_s=120)


def _primary(r: dict) -> dict:
    srcs = r.get("sources") or []
    if not srcs:
        return {"title": "", "uri": "", "used": False}
    pri = next(
        (s for s in srcs if s.get("source_role") == "primary_answer_source"),
        srcs[0],
    )
    return {
        "title": (pri.get("title") or "")[:55],
        "uri": pri.get("document_uri") or "",
        "used": bool(pri.get("used")),
    }


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--shards-dir", type=Path, default=DEFAULT_SHARDS)
    p.add_argument("--qa-db", type=Path,
                   help="qa.db path. Default: $TMP/qa_legacy_vs_providence.db "
                        "(fresh per-run so cache doesn't bleed across themes)")
    p.add_argument("--n-per-theme", type=int, default=5)
    p.add_argument("--themes", nargs="*", default=None,
                   help="subset of theme names (default: all)")
    p.add_argument("--out-dir", type=Path,
                   default=REPO / "bench" / "legacy_vs_providence_results")
    args = p.parse_args()

    if args.qa_db is None:
        import tempfile
        args.qa_db = Path(tempfile.mkdtemp()) / "qa_legacy_vs_providence.db"

    selected = THEMES if args.themes is None else [
        (k, d) for k, d in THEMES if k in args.themes
    ]
    args.out_dir.mkdir(parents=True, exist_ok=True)
    stamp = _dt.datetime.now(_dt.timezone.utc).strftime("%Y-%m-%dT%H-%M-%SZ")
    out_path = args.out_dir / f"{stamp}.jsonl"

    print(f"# {len(selected)} themes × {args.n_per_theme} questions × 2 paths",
          file=sys.stderr)
    print(f"# shards: {args.shards_dir}", file=sys.stderr)
    print(f"# qa.db:  {args.qa_db}", file=sys.stderr)
    print(f"# llm:    $ARBORIST_LLM_ENDPOINT or hermes default", file=sys.stderr)
    print(f"# out:    {out_path}", file=sys.stderr)
    print(file=sys.stderr)

    theme_summary: dict[str, dict] = {}
    with open(out_path, "w") as f:
        for theme_key, theme_desc in selected:
            fixture = REPO / "bench" / f"qa_questions_{theme_key}.txt"
            if not fixture.exists():
                print(f"!! skip {theme_key}: {fixture} missing", file=sys.stderr)
                continue
            questions = _load_fixture(fixture, args.n_per_theme)
            print(f"=== {theme_key:25s}  ({theme_desc})", file=sys.stderr)
            agree = 0
            disagree = 0
            audit_better = 0
            audit_worse = 0
            for i, q in enumerate(questions, 1):
                default = _query(
                    q, shards_dir=args.shards_dir, qa_db=args.qa_db,
                    legacy=False,
                )
                legacy = _query(
                    q, shards_dir=args.shards_dir, qa_db=args.qa_db,
                    legacy=True,
                )
                pri_def = _primary(default)
                pri_leg = _primary(legacy)
                same_primary = pri_def["uri"] == pri_leg["uri"]
                amode_def = default.get("audit_mode") or "?"
                amode_leg = legacy.get("audit_mode") or "?"
                # Rank: STRICT > HYBRID > UNGROUNDED for the "did the
                # path produce a stronger verdict" axis.
                rank = {"STRICT": 3, "HYBRID": 2, "UNGROUNDED": 1}
                r_def = rank.get(amode_def, 0)
                r_leg = rank.get(amode_leg, 0)
                row = {
                    "theme": theme_key,
                    "question": q,
                    "default": {
                        "audit_mode": amode_def,
                        "n_verified": default.get("n_verified"),
                        "primary": pri_def,
                        "elapsed_s": default.get("_elapsed_s"),
                        "error": default.get("_error"),
                    },
                    "legacy": {
                        "audit_mode": amode_leg,
                        "n_verified": legacy.get("n_verified"),
                        "primary": pri_leg,
                        "elapsed_s": legacy.get("_elapsed_s"),
                        "error": legacy.get("_error"),
                    },
                    "same_primary": same_primary,
                    "audit_delta": r_def - r_leg,  # +1 default better
                }
                if same_primary:
                    agree += 1
                else:
                    disagree += 1
                if r_def > r_leg:
                    audit_better += 1
                elif r_def < r_leg:
                    audit_worse += 1
                f.write(json.dumps(row, ensure_ascii=False) + "\n")
                f.flush()
                badge = " " if same_primary else "≠"
                print(f"  [{i}/{len(questions)}] {q[:50]:50s}",
                      file=sys.stderr)
                print(f"        default: {amode_def:11s} "
                      f"{pri_def['title']:42s}",
                      file=sys.stderr)
                print(f"        legacy:  {amode_leg:11s} "
                      f"{pri_leg['title']:42s}  {badge}",
                      file=sys.stderr)
            theme_summary[theme_key] = {
                "agree": agree, "disagree": disagree,
                "audit_better_default": audit_better,
                "audit_worse_default": audit_worse,
                "n": len(questions),
            }
            print(file=sys.stderr)

    print("=== SUMMARY ===", file=sys.stderr)
    print(f"  {'theme':25s}  agree  disagree  audit↑  audit↓  n",
          file=sys.stderr)
    for k, s in theme_summary.items():
        print(f"  {k:25s}  {s['agree']:5d}  {s['disagree']:8d}  "
              f"{s['audit_better_default']:6d}  "
              f"{s['audit_worse_default']:6d}  {s['n']:3d}",
              file=sys.stderr)
    print(f"  results: {out_path}", file=sys.stderr)
    return 0


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