arborist/bench/three_way_bench.py

"""3-way bench: legacy local vs corpus-query vs cloud-query.

Goal: every fixture question produces the SAME primary-source citation
across all three paths. Diverging primary-source picks are real
regressions even when audit_mode happens to agree.

  legacy        arborist --shards-dir … query --json --burn --answer-mode claim_lattice ...
  corpus-query  arborist --shards-dir … corpus-query --json ...   (Corpus → run_query)
  cloud-query   arborist cloud query --bucket-url <manifest> --json ...   (Sidecar → run_query)

JSONL output: one row per question with all three path results +
agreement flags so iterations can compare across bench runs.
"""
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_BUCKET = (
    "https://nyc3.digitaloceanspaces.com/arborist/clones/manifest-sidecar.json"
)
DEFAULT_QWEN_ENDPOINT = "https://qwen.ai.unturf.com/v1"
DEFAULT_QWEN_MODEL = "Qwen3.6-27B-UD-Q4_K_XL.gguf"


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


def _run(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[:200],
                "_elapsed_s": round(dt, 2)}
    d["_elapsed_s"] = round(dt, 2)
    return d


def _legacy(q, *, shards_dir):
    return _run([
        str(ARBORIST), "--shards-dir", str(shards_dir),
        "query", "--top-k", "8", "--json", "--burn",
        "--answer-mode", "claim_lattice", q,
    ], timeout_s=120)


def _corpus(q, *, shards_dir, endpoint, model):
    return _run([
        str(ARBORIST), "--shards-dir", str(shards_dir),
        "corpus-query", q,
        "--top-k", "4", "--json",
        "--endpoint", endpoint, "--model", model,
    ], timeout_s=120)


def _cloud(q, *, bucket_url, endpoint, model):
    return _run([
        str(ARBORIST), "cloud", "query", q,
        "--bucket-url", bucket_url,
        "--endpoint", endpoint, "--model", model,
        "--top-k", "4", "--json",
    ], timeout_s=240)


def _primary(r: dict) -> dict:
    srcs = r.get("sources") or []
    if not srcs:
        return {"title": "", "uri": ""}
    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 ""}


def main():
    p = argparse.ArgumentParser()
    p.add_argument(
        "--fixture", type=Path,
        default=REPO / "bench" / "qa_questions_smoke.txt",
    )
    p.add_argument("--shards-dir", type=Path, default=DEFAULT_SHARDS)
    p.add_argument("--bucket-url", default=DEFAULT_BUCKET)
    p.add_argument("--endpoint", default=DEFAULT_QWEN_ENDPOINT)
    p.add_argument("--model", default=DEFAULT_QWEN_MODEL)
    p.add_argument(
        "--out-dir", type=Path,
        default=REPO / "bench" / "three_way_results",
    )
    p.add_argument("--skip-legacy", action="store_true")
    p.add_argument("--skip-corpus", action="store_true")
    p.add_argument("--skip-cloud", action="store_true")
    args = p.parse_args()

    questions = _load_fixture(args.fixture)
    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(questions)} questions from {args.fixture}", file=sys.stderr)
    print(f"# shards: {args.shards_dir}", file=sys.stderr)
    print(f"# bucket: {args.bucket_url}", file=sys.stderr)
    print(f"# llm:    {args.endpoint} / {args.model}", file=sys.stderr)
    print(file=sys.stderr)

    n_full_agreement = 0
    n_source_disagree = 0

    with open(out_path, "w") as f:
        for i, q in enumerate(questions, 1):
            print(f"[{i}/{len(questions)}] {q}", file=sys.stderr)
            legacy = {} if args.skip_legacy else _legacy(q, shards_dir=args.shards_dir)
            corpus = {} if args.skip_corpus else _corpus(
                q, shards_dir=args.shards_dir,
                endpoint=args.endpoint, model=args.model,
            )
            cloud = {} if args.skip_cloud else _cloud(
                q, bucket_url=args.bucket_url,
                endpoint=args.endpoint, model=args.model,
            )

            row = {"question": q}
            for label, r in (("legacy", legacy), ("corpus", corpus), ("cloud", cloud)):
                row[label] = {
                    "audit_mode": r.get("audit_mode"),
                    "n_verified": r.get("n_verified"),
                    "n_quotes": r.get("n_quotes"),
                    "primary": _primary(r),
                    "elapsed_s": r.get("_elapsed_s"),
                    "error": r.get("_error"),
                }
            # Agreement: same primary URI across all PRESENT paths
            uris = {
                lbl: row[lbl]["primary"]["uri"]
                for lbl in ("legacy", "corpus", "cloud")
                if row[lbl]["primary"]["uri"]
            }
            distinct = set(uris.values())
            row["agreement"] = {
                "n_paths_with_source": len(uris),
                "distinct_sources": len(distinct),
                "all_agree": len(distinct) <= 1,
            }
            if row["agreement"]["all_agree"]:
                n_full_agreement += 1
            else:
                n_source_disagree += 1

            f.write(json.dumps(row, ensure_ascii=False) + "\n")
            f.flush()

            for label in ("legacy", "corpus", "cloud"):
                r = row[label]
                badge = " ←" if r["primary"]["uri"] and len(distinct) > 1 and r["primary"]["uri"] != min(uris.values()) else ""
                err = f"  ERR={r['error']}" if r["error"] else ""
                print(f"   {label:7s}  {(r['audit_mode'] or '?'):<25} "
                      f"{r['primary']['title']}{badge}{err}",
                      file=sys.stderr)
            if not row["agreement"]["all_agree"]:
                print(f"   ⚠ SOURCE DISAGREEMENT: {len(distinct)} distinct primaries",
                      file=sys.stderr)
            print(file=sys.stderr)

    print(f"=== SUMMARY ===", file=sys.stderr)
    print(f"  {len(questions)} questions", file=sys.stderr)
    print(f"  {n_full_agreement} all-paths agree on primary source", file=sys.stderr)
    print(f"  {n_source_disagree} paths-disagree", file=sys.stderr)
    print(f"  results: {out_path}", file=sys.stderr)
    return 0


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