arborist/bench/scripts/witness_sweep.py

"""End-to-end witness sweep against real shards.

Fires a small, deterministic canonical-question set through
``arborist.qa.query.query`` with ``canonical_witness_enabled=True``
against the operator's real shard directory. Each question:

1. Hits the canonical-projection preflight (math/logic/algebra route).
2. Persists a providence_cache row (#000027 path).
3. Runs witness fan-out: kernel + cache + LLM in parallel
   (#000028 path).
4. Writes a ``providence_canonical_witness`` audit event (the new
   wiring landed in commit `70ffc01`).
5. Records LLM cost in the capital ledger (op_type='canonical_witness').

After the sweep, the operator can run::

    make bench-witness-divergence

…to extract any LLM-divergence events as 5F-Falsification
calibration fixtures.

Why this script exists
----------------------
The witness pipeline has unit-test coverage via StubClient
(``tests/test_witness*.py``), but those don't validate against the
real Hermes endpoint. This sweep is the "does it actually work
end-to-end" check — useful before publishing divergence fixtures
to a downstream consumer or before flipping witness=on in production.

Usage::

    python -m bench.scripts.witness_sweep \\
        --shards-dir ~/.arborist/shards \\
        --out bench/results/witness-sweep.json

Or via Makefile::

    make bench-witness-sweep

Default question set (8 canonical-shape questions)
--------------------------------------------------
Mix of arithmetic, logic, and algebra (the three preflight routes):

- Arithmetic: ``0.1 + 0.2``, ``1/3 + 1/6``, ``2**10``
- Logic: ``A IMPL B``, ``(NOT B) IMPL (NOT A)``, ``A OR NOT A``
- Algebra: ``(x+1)**2``, ``x**2 + 2*x + 1``

For each question we record: agreement_label, kernel canonical,
LLM raw, LLM canonicalized bytes, audit-seq, total wall.
"""

from __future__ import annotations

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


_DEFAULT_QUESTIONS = [
    "0.1 + 0.2",
    "1/3 + 1/6",
    "2**10",
    "A IMPL B",
    "(NOT B) IMPL (NOT A)",
    "A OR NOT A",
    "(x+1)**2",
    "x**2 + 2*x + 1",
]


def _run_one(
    *,
    question: str,
    shards_dir: Path,
    qa_db: Path,
    chat_client,
    model_id: str,
) -> dict:
    """Fire one query with witness on; extract a small summary."""
    from arborist.qa.query import DEFAULT_QUERY_POLICY, query
    policy = dict(DEFAULT_QUERY_POLICY)
    policy["canonical_witness_enabled"] = True
    policy["canonical_witness_sample_rate"] = 1.0  # always-on for sweep
    policy["canonical_witness_timeout_s"] = 30.0
    t0 = time.monotonic()
    result = query(
        question=question,
        qa_db=qa_db,
        chat_client=chat_client,
        model_id=model_id,
        shards_dir=shards_dir,
        policy=policy,
    )
    elapsed_ms = round((time.monotonic() - t0) * 1000, 1)
    witness = result.get("witness") or {}
    llm = (witness.get("modalities") or {}).get("llm") or {}
    cache = (witness.get("modalities") or {}).get("cache") or {}
    return {
        "question": question,
        "status": result.get("status"),
        "audit_mode": result.get("audit_mode"),
        "pi_star_ref": result.get("pi_star_ref"),
        "canonical_answer": result.get("answer_text"),
        "agreement_label": witness.get("agreement_label"),
        "llm_raw": llm.get("raw_answer"),
        "llm_canonical_match": (
            llm.get("canonical_bytes") is not None
        ),
        "cache_status": cache.get("error") or "ok",
        "wall_ms": elapsed_ms,
        "audit_event_hash": result.get("audit_event_hash"),
    }


def _summarize(rows: list[dict]) -> dict:
    n = len(rows)
    by_label: dict[str, int] = {}
    diverged = 0
    for r in rows:
        label = r.get("agreement_label") or "—"
        by_label[label] = by_label.get(label, 0) + 1
        if label in ("LLM-DIVERGED", "KERNEL-LLM-DIVERGED", "CACHE-DRIFT"):
            diverged += 1
    walls = [r["wall_ms"] for r in rows if isinstance(r.get("wall_ms"), (int, float))]
    walls_sorted = sorted(walls)
    return {
        "questions": n,
        "by_agreement_label": by_label,
        "divergence_count": diverged,
        "divergence_rate": (diverged / n) if n else 0.0,
        "wall_ms_median": walls_sorted[len(walls_sorted) // 2] if walls_sorted else None,
        "wall_ms_max": max(walls) if walls else None,
    }


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(
        "--out", type=Path,
        default=Path("bench/results/witness-sweep.json"),
        help="JSON output path",
    )
    p.add_argument(
        "--endpoint", default=os.environ.get(
            "ARBORIST_LLM_ENDPOINT", "https://hermes.ai.unturf.com/v1"
        ),
    )
    p.add_argument(
        "--model", default=os.environ.get(
            "ARBORIST_LLM_MODEL",
            "adamo1139/Hermes-3-Llama-3.1-8B-FP8-Dynamic",
        ),
    )
    p.add_argument(
        "--questions", nargs="*", default=None,
        help="override the default canonical-question set",
    )
    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

    from arborist.qa.client import OpenAICompatibleClient

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

    print(
        f"witness sweep: {len(questions)} questions against "
        f"{args.shards_dir} (LLM: {args.endpoint})",
        file=sys.stderr,
    )

    rows = []
    for q in questions:
        print(f"  → {q[:60]}", file=sys.stderr)
        try:
            row = _run_one(
                question=q, shards_dir=args.shards_dir,
                qa_db=qa_db, chat_client=chat_client,
                model_id=args.model,
            )
        except Exception as exc:
            row = {"question": q, "error": str(exc)}
        rows.append(row)

    summary = _summarize(rows)
    artifact = {
        "schema_version": "witness-sweep-v1",
        "timestamp_utc": _dt.datetime.now(_dt.timezone.utc).isoformat(timespec="seconds"),
        "shards_dir": str(args.shards_dir),
        "model": args.model,
        "endpoint": args.endpoint,
        "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))
    print(f"wrote {args.out}", file=sys.stderr)
    print(json.dumps(summary, indent=2))
    return 0


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