arborist/bench/cloud_vs_local.py

"""Compare local `arborist query` vs `arborist cloud query` on a fixture.

For each question:
  1. Run local CLI: arborist --shards-dir SHARDS_DIR query --json --burn ...
  2. Run cloud CLI: arborist cloud query --bucket-url ... --json ...
  3. Compare audit_mode, primary source URI, answer text

A regression is any of:
  - local STRICT (or EVIDENCE-WARRANTED) but cloud UNGROUNDED
  - local cited Wikipedia/primary article but cloud cited a tangential one
  - cloud crash / network error

Prints a side-by-side table + a regression summary. JSONL output goes
to bench/cloud_vs_local_results/<utc-iso>.jsonl.
"""
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]:
    questions = []
    for line in path.read_text().splitlines():
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        questions.append(line)
    return questions


def _run_local(question: str, *, shards_dir: Path, timeout_s: int = 120) -> dict:
    """Invoke arborist query --json, parse stdout."""
    t0 = time.time()
    try:
        proc = subprocess.run(
            [
                str(ARBORIST), "--shards-dir", str(shards_dir),
                "query", "--top-k", "8", "--json", "--burn",
                "--answer-mode", "claim_lattice",
                question,
            ],
            capture_output=True, text=True, timeout=timeout_s,
        )
    except subprocess.TimeoutExpired:
        return {"_error": f"timeout after {timeout_s}s", "_elapsed_s": timeout_s}
    elapsed = time.time() - t0
    if proc.returncode != 0:
        return {
            "_error": f"exit {proc.returncode}",
            "_stderr": proc.stderr[:500],
            "_elapsed_s": round(elapsed, 2),
        }
    try:
        data = json.loads(proc.stdout)
    except json.JSONDecodeError as e:
        return {"_error": f"parse: {e}", "_stdout_head": proc.stdout[:200],
                "_elapsed_s": round(elapsed, 2)}
    data["_elapsed_s"] = round(elapsed, 2)
    return data


def _run_cloud(
    question: str, *, bucket_url: str, endpoint: str, model: str,
    timeout_s: int = 240,
) -> dict:
    """Invoke arborist cloud query --json, parse stdout."""
    t0 = time.time()
    try:
        proc = subprocess.run(
            [
                str(ARBORIST), "cloud", "query", question,
                "--bucket-url", bucket_url,
                "--endpoint", endpoint,
                "--model", model,
                "--top-k", "4",
                "--json",
            ],
            capture_output=True, text=True, timeout=timeout_s,
            env={**os.environ, "ARBORIST_PROGRESS": "0"},
        )
    except subprocess.TimeoutExpired:
        return {"_error": f"timeout after {timeout_s}s", "_elapsed_s": timeout_s}
    elapsed = time.time() - t0
    if proc.returncode != 0:
        return {
            "_error": f"exit {proc.returncode}",
            "_stderr": proc.stderr[:500],
            "_elapsed_s": round(elapsed, 2),
        }
    try:
        data = json.loads(proc.stdout)
    except json.JSONDecodeError as e:
        return {"_error": f"parse: {e}", "_stdout_head": proc.stdout[:200],
                "_elapsed_s": round(elapsed, 2)}
    data["_elapsed_s"] = round(elapsed, 2)
    return data


def _primary_source(result: dict) -> dict:
    """Extract primary source = first source marked source_role=primary_answer_source,
    else first source if none marked."""
    srcs = result.get("sources") or []
    if not srcs:
        return {}
    primary = next(
        (s for s in srcs if s.get("source_role") == "primary_answer_source"),
        srcs[0],
    )
    return {
        "title": primary.get("title", "")[:60],
        "uri": primary.get("document_uri", ""),
        "used": primary.get("used"),
    }


def _is_regression(local: dict, cloud: dict) -> tuple[bool, str]:
    """Return (is_regression, reason)."""
    if cloud.get("_error"):
        return True, f"cloud error: {cloud['_error']}"
    if local.get("_error"):
        return False, "local errored too (not a regression)"
    local_audit = local.get("audit_mode")
    cloud_audit = cloud.get("audit_mode")
    rank = {"UNGROUNDED": 0, "HYBRID": 1, "STRICT": 2}
    if rank.get(cloud_audit, -1) < rank.get(local_audit, -1):
        return True, f"audit_mode dropped {local_audit} → {cloud_audit}"
    # Primary-source URI mismatch where both are non-empty
    lp = _primary_source(local)
    cp = _primary_source(cloud)
    if lp.get("uri") and cp.get("uri") and lp["uri"] != cp["uri"]:
        # Tolerable if both contain the same key noun (eg. variant article).
        # Strict check for now — surface and let human judge.
        return True, f"primary uri differs: {lp['uri'][:60]} vs {cp['uri'][:60]}"
    return False, "ok"


def main(argv=None):
    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" / "cloud_vs_local_results",
    )
    p.add_argument("--skip-local", action="store_true",
                   help="only run cloud (e.g. when local results are pre-recorded)")
    args = p.parse_args(argv)

    questions = _load_fixture(args.fixture)
    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)

    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"

    rows: list[dict] = []
    regressions: list[tuple[str, str]] = []

    with open(out_path, "w") as out:
        for i, q in enumerate(questions, 1):
            print(f"[{i}/{len(questions)}] {q}", file=sys.stderr)
            local = {} if args.skip_local else _run_local(q, shards_dir=args.shards_dir)
            cloud = _run_cloud(
                q, bucket_url=args.bucket_url,
                endpoint=args.endpoint, model=args.model,
            )
            row = {
                "question": q,
                "local": {
                    "audit_mode": local.get("audit_mode"),
                    "n_quotes": local.get("n_quotes"),
                    "n_verified": local.get("n_verified"),
                    "primary": _primary_source(local),
                    "elapsed_s": local.get("_elapsed_s"),
                    "error": local.get("_error"),
                },
                "cloud": {
                    "audit_mode": cloud.get("audit_mode"),
                    "n_quotes": cloud.get("n_quotes"),
                    "n_verified": cloud.get("n_verified"),
                    "primary": _primary_source(cloud),
                    "elapsed_s": cloud.get("_elapsed_s"),
                    "error": cloud.get("_error"),
                },
            }
            is_reg, reason = _is_regression(local, cloud)
            row["regression"] = is_reg
            row["reason"] = reason
            rows.append(row)
            out.write(json.dumps(row, ensure_ascii=False) + "\n")
            out.flush()

            # Per-question terminal line.
            print(
                f"   local:  {local.get('audit_mode', '?'):<25} {row['local']['primary'].get('title','')[:55]}",
                file=sys.stderr,
            )
            print(
                f"   cloud:  {cloud.get('audit_mode', '?'):<25} {row['cloud']['primary'].get('title','')[:55]}",
                file=sys.stderr,
            )
            if is_reg:
                regressions.append((q, reason))
                print(f"   ⚠ REGRESSION: {reason}", file=sys.stderr)
            print(file=sys.stderr)

    print(f"=== SUMMARY ===", file=sys.stderr)
    print(f"  {len(questions)} questions, {len(regressions)} regressions", file=sys.stderr)
    for q, r in regressions:
        print(f"  ⚠ {q[:60]} — {r}", file=sys.stderr)
    print(f"  results: {out_path}", file=sys.stderr)
    return 0


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