arborist/bench/scripts/real_shard_baseline.py

"""Real-shard workload baseline (#000026 Phase 2).

Runs a fixed question set against the operator's real shard
directory and emits one JSON artifact + a short markdown summary.

Usage:

    python -m bench.scripts.real_shard_baseline \
        --shards-dir ~/.arborist/shards \
        --fixtures bench/fixtures/real-shard-baseline-v1.jsonl \
        --out bench/results/real-shard-baseline.json

Or via the Makefile:

    make bench-real-shard

Outputs:

    bench/results/real-shard-baseline.json   one row per question
    bench/results/real-shard-baseline.md     markdown summary

Captured per question:

    - question, audit_mode, verifier_method, n_verified, n_quotes
    - wall_ms, search_ms, llm_ms, canonical_ms (preflight short-circuit)
    - primary_source_used, sources_count
    - status (canonical_projection | cache_hit | cache_miss_then_written
              | no_sources | broad_quantifier_rejected | ...)

This is observation, not enforcement. Baselines NEVER gate CI per the
ticket's hard constraint. The artifact is for: (a) confirming a fix
moved the needle, (b) seeding ForkScore comparisons, (c) noting
warrant-quality findings worth tickets.
"""

from __future__ import annotations

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


def _shard_fingerprint(shards_dir: Path) -> str:
    """SHA-256 over (path, size, mtime_ns) for every *.db in shards_dir.
    Stable across re-opens; changes when any shard mutates."""
    if not shards_dir.is_dir():
        return ""
    parts = []
    for p in sorted(shards_dir.glob("*.db")):
        st = p.stat()
        parts.append(f"{p.name}\t{st.st_size}\t{st.st_mtime_ns}")
    return hashlib.sha256("\n".join(parts).encode("utf-8")).hexdigest()


def _commit_sha() -> str:
    try:
        return subprocess.check_output(
            ["git", "rev-parse", "HEAD"], cwd=Path(__file__).parent
        ).decode().strip()
    except Exception:
        return ""


def _git_dirty() -> bool:
    try:
        out = subprocess.check_output(
            ["git", "status", "--porcelain"], cwd=Path(__file__).parent
        ).decode()
        return bool(out.strip())
    except Exception:
        return False


def _read_fixtures(path: Path) -> list[dict]:
    rows = []
    for line in path.read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        obj = json.loads(line)
        if "_meta" in obj:
            continue
        rows.append(obj)
    return rows


def _run_one(
    *,
    question: str,
    shards_dir: Path,
    qa_db: Path,
    chat_client,
    model_id: str,
    burn: bool,
) -> dict:
    """Execute one query through the real query() entry point and
    extract a small, stable summary dict."""
    from arborist.qa.query import query as _query
    from arborist.qa.client import StubClient  # noqa: F401 — type hint use

    t0 = time.monotonic()
    result = _query(
        question=question,
        qa_db=qa_db,
        chat_client=chat_client,
        model_id=model_id,
        shards_dir=shards_dir,
        burn_existing=burn,
    )
    elapsed_ms = round((time.monotonic() - t0) * 1000, 1)

    timings = result.get("timings") or {}
    sources = result.get("sources") or []
    primary_used = any(
        s.get("source_role") in (
            "primary", "primary_answer", "primary_answer_source"
        ) and s.get("used") is not False
        for s in sources
    )
    return {
        "question": question,
        "status": result.get("status"),
        "audit_mode": result.get("audit_mode"),
        "verifier_method": result.get("verifier_method"),
        "n_verified": result.get("n_verified"),
        "n_quotes": result.get("n_quotes"),
        "pi_star_ref": result.get("pi_star_ref"),
        "lookup_path": result.get("lookup_path"),
        "wall_ms": elapsed_ms,
        "total_ms": timings.get("total_ms"),
        "search_ms": timings.get("search_ms"),
        "llm_ms": timings.get("llm_ms"),
        "verify_ms": timings.get("verify_ms"),
        "canonical_preflight_ms": timings.get("canonical_preflight_ms"),
        "sources_count": len(sources),
        "primary_source_used": primary_used,
        "answer_text": (result.get("answer_text") or "")[:240],
    }


def _summarize(rows: list[dict]) -> dict:
    """Aggregate the per-question rows into headline metrics."""
    n = len(rows)
    walls = [r["wall_ms"] for r in rows if isinstance(r.get("wall_ms"), (int, float))]
    walls_sorted = sorted(walls)
    median = walls_sorted[len(walls_sorted) // 2] if walls_sorted else None

    by_audit: dict[str, int] = {}
    by_status: dict[str, int] = {}
    for r in rows:
        a = r.get("audit_mode") or "—"
        s = r.get("status") or "—"
        by_audit[a] = by_audit.get(a, 0) + 1
        by_status[s] = by_status.get(s, 0) + 1
    return {
        "questions": n,
        "wall_ms_median": median,
        "wall_ms_max": max(walls) if walls else None,
        "wall_ms_min": min(walls) if walls else None,
        "by_audit_mode": by_audit,
        "by_status": by_status,
        "primary_source_used_count": sum(
            1 for r in rows if r.get("primary_source_used")
        ),
        "canonical_projection_count": sum(
            1 for r in rows if r.get("status") == "canonical_projection"
        ),
    }


def _markdown(rows: list[dict], summary: dict, meta: dict) -> str:
    lines = [
        "# Real-shard baseline — " + meta["timestamp_utc"],
        "",
        f"**Commit:** `{meta['commit_sha'][:12]}`"
        + (" (dirty)" if meta.get("git_dirty") else ""),
        f"**Shards fingerprint:** `{meta['shards_fingerprint'][:16]}…`",
        f"**Shards directory:** `{meta['shards_dir']}`",
        "",
        "## Summary",
        "",
        f"- Questions: **{summary['questions']}**",
        f"- Wall median: **{summary['wall_ms_median']} ms**",
        f"- Wall range: {summary['wall_ms_min']} – {summary['wall_ms_max']} ms",
        f"- Primary source used: {summary['primary_source_used_count']}/"
        f"{summary['questions']}",
        f"- Canonical-projection short-circuits: "
        f"{summary['canonical_projection_count']}",
        "",
        "**Audit-mode distribution:** "
        + ", ".join(f"`{k}` {v}" for k, v in sorted(summary["by_audit_mode"].items())),
        "",
        "**Status distribution:** "
        + ", ".join(f"`{k}` {v}" for k, v in sorted(summary["by_status"].items())),
        "",
        "## Per-question",
        "",
        "| # | Question | Audit | Status | Wall (ms) | Method |",
        "|---|----------|-------|--------|-----------|--------|",
    ]
    for i, r in enumerate(rows, 1):
        q = r["question"]
        if len(q) > 60:
            q = q[:57] + "…"
        lines.append(
            f"| {i} | {q} | "
            f"{r.get('audit_mode') or '—'} | "
            f"{r.get('status') or '—'} | "
            f"{r.get('wall_ms') or '—'} | "
            f"{r.get('verifier_method') or r.get('pi_star_ref') or '—'} |"
        )
    return "\n".join(lines) + "\n"


def main(argv: list[str] | None = None) -> int:
    p = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    p.add_argument(
        "--shards-dir", type=Path, required=True,
        help="path to the shard directory (e.g. ~/.arborist/shards)",
    )
    p.add_argument(
        "--fixtures", type=Path,
        default=Path("bench/fixtures/real-shard-baseline-v1.jsonl"),
        help="JSONL question set (default: real-shard-baseline-v1)",
    )
    p.add_argument(
        "--out", type=Path,
        default=Path("bench/results/real-shard-baseline.json"),
        help="JSON output path",
    )
    p.add_argument(
        "--md-out", type=Path,
        default=Path("bench/results/real-shard-baseline.md"),
        help="markdown summary path",
    )
    p.add_argument(
        "--endpoint", default=os.environ.get(
            "ARBORIST_LLM_ENDPOINT", "https://hermes.ai.unturf.com/v1"
        ),
        help="LLM endpoint URL",
    )
    p.add_argument(
        "--model", default=os.environ.get(
            "ARBORIST_LLM_MODEL",
            "adamo1139/Hermes-3-Llama-3.1-8B-FP8-Dynamic",
        ),
        help="model id (default $ARBORIST_LLM_MODEL or hermes-3)",
    )
    p.add_argument(
        "--burn", action="store_true",
        help="force fresh inference per question (skip cache)",
    )
    args = p.parse_args(argv)

    if not args.shards_dir.is_dir():
        print(f"error: shards-dir not found: {args.shards_dir}", file=sys.stderr)
        return 2
    if not args.fixtures.is_file():
        print(f"error: fixtures not found: {args.fixtures}", file=sys.stderr)
        return 2

    # Defer LLM client construction until the args are known.
    from arborist.qa.client import OpenAICompatibleClient

    qa_db = args.shards_dir / "qa.db"
    chat_client = OpenAICompatibleClient(base_url=args.endpoint)

    fixtures = _read_fixtures(args.fixtures)
    print(
        f"running {len(fixtures)} questions against {args.shards_dir}",
        file=sys.stderr,
    )

    rows = []
    for fx in fixtures:
        question = fx["question"]
        print(f"  → {question[:60]}", file=sys.stderr)
        try:
            row = _run_one(
                question=question,
                shards_dir=args.shards_dir,
                qa_db=qa_db,
                chat_client=chat_client,
                model_id=args.model,
                burn=args.burn,
            )
        except Exception as exc:  # one bad question shouldn't kill the run
            row = {
                "question": question,
                "status": "error",
                "error": str(exc),
            }
        # Fold any operator notes from the fixture
        if "expected_routing" in fx:
            row["expected_routing"] = fx["expected_routing"]
        if "note" in fx:
            row["note"] = fx["note"]
        rows.append(row)

    summary = _summarize(rows)
    meta = {
        "schema_version": "real-shard-baseline-v1",
        "timestamp_utc": _dt.datetime.now(_dt.timezone.utc).isoformat(
            timespec="seconds"
        ),
        "commit_sha": _commit_sha(),
        "git_dirty": _git_dirty(),
        "shards_dir": str(args.shards_dir),
        "shards_fingerprint": _shard_fingerprint(args.shards_dir),
        "model": args.model,
        "endpoint": args.endpoint,
        "burn": args.burn,
        "fixtures_path": str(args.fixtures),
    }

    artifact = {**meta, "summary": summary, "queries": rows}

    args.out.parent.mkdir(parents=True, exist_ok=True)
    args.out.write_text(json.dumps(artifact, indent=2, ensure_ascii=False))
    args.md_out.write_text(_markdown(rows, summary, meta))

    print(f"wrote {args.out}", file=sys.stderr)
    print(f"wrote {args.md_out}", file=sys.stderr)
    print(json.dumps(summary, indent=2))
    return 0


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