arborist/tests/test_nli_shadow.py

"""Tests for the #000049 Phase-2 NLI shadow scaffold.

These run in the default suite — i.e. WITHOUT the ``[nli]`` extra
installed — so they exercise the pure-Python parts (clause splitter,
manifest, label-index resolution, the dataclass contract) and the
graceful-degradation path (deps missing → ``available=False``, never an
exception). When ``[nli]`` *is* installed they also cover the real
model path; we don't assert specific probabilities (those belong in the
``arborist-nli-bench`` scorecard, not the unit suite).
"""
from __future__ import annotations

import json
from pathlib import Path

import pytest

from arborist.qa.nli import ShadowNLI, ShadowResult, shadow_check, load_manifest
from arborist.qa.nli.shadow import clauses, candidate_clauses, _content_tokens, _resolve_label_indices


# --- clause splitter -------------------------------------------------------

def test_clauses_splits_on_sentence_punctuation():
    assert clauses("Jupiter is the largest. Mercury is the smallest.") == [
        "Jupiter is the largest.",
        "Mercury is the smallest.",
    ]


def test_clauses_handles_semicolons():
    assert clauses("A is true; B is false") == ["A is true;", "B is false"]


def test_clauses_single_sentence_is_one_clause():
    assert clauses("Just one sentence with no breaks") == ["Just one sentence with no breaks"]


def test_clauses_empty_or_blank():
    assert clauses("") == []
    assert clauses("   ") == []
    assert clauses(None) == []  # type: ignore[arg-type]


# --- candidate-clause restriction (§7 #5 step 3 / §7 #20) ------------------

def test_content_tokens_filters_short_and_stopwords():
    toks = _content_tokens("The Mercury is the largest planet in the Solar System.")
    assert "mercury" in toks and "largest" in toks and "planet" in toks and "solar" in toks
    assert "the" not in toks and "is" not in toks and "in" not in toks  # stopword / too short


def test_candidate_clauses_ranks_by_overlap_and_caps():
    claim = "Mercury is the largest planet in the Solar System."
    source = ("Jupiter is the largest planet in the Solar System. "
              "Mercury is the smallest planet in the Solar System. "
              "The Roman god Mercury was the messenger of the gods. "
              "Bananas are yellow.")
    cand = candidate_clauses(claim, source, k=6)
    # the banana clause shares no content tokens → dropped
    assert all("Bananas" not in c for c, _ in cand)
    # the two planet clauses are the highest-overlap ones
    assert "Jupiter is the largest planet" in cand[0][0] or "Mercury is the smallest planet" in cand[0][0]
    assert cand[0][1] > cand[-1][1] or len(cand) == 1
    # cap respected
    assert len(candidate_clauses(claim, source, k=1)) == 1


def test_candidate_clauses_empty_when_no_overlap():
    assert candidate_clauses("Bananas are yellow.", "Mercury is the smallest planet.", k=6) == []
    assert candidate_clauses("", "anything at all here", k=6) == []


def test_check_restricts_to_candidate_clauses_not_whole_haystack():
    # a long source where only 2 clauses touch the claim — the result must
    # report n_candidate_clauses << n_clauses (this is the §7 #20 fix).
    haystack = " ".join(f"Unrelated fact number {i} about pottery." for i in range(40))
    source = "Jupiter is the largest planet. " + haystack + " Mercury is the smallest planet."
    res = shadow_check("Mercury is the largest planet.", source)
    assert res.n_clauses >= 40
    assert res.n_candidate_clauses <= 6
    assert res.n_candidate_clauses < res.n_clauses


def test_check_no_candidate_clauses_is_no_demote():
    res = shadow_check("Mercury is the largest planet.",
                       "Bananas are yellow. Pottery is ancient. The weather is fine.")
    assert res.would_demote is False
    if res.available:
        assert res.reason == "no_candidate_clauses" and res.n_candidate_clauses == 0


def test_manifest_has_candidate_clause_cap():
    assert load_manifest()["max_candidate_clauses"] == 6


# --- manifest --------------------------------------------------------------

def test_manifest_has_required_fields():
    m = load_manifest()
    for key in ("nli_model_version", "hf_repo", "pinned_revision", "license", "thresholds"):
        assert key in m, key
    assert "contradiction_veto" in m["thresholds"]
    assert "entailment_block_veto" in m["thresholds"]
    # shadow-mode thresholds are the bench-validated point
    assert m["thresholds"]["contradiction_veto"] == 0.5
    assert m["thresholds"]["entailment_block_veto"] == 0.9


def test_manifest_is_valid_json_file():
    p = Path(__file__).resolve().parents[1] / "arborist" / "qa" / "nli" / "manifest.json"
    json.loads(p.read_text())  # raises on malformed


# --- label-index resolution (the cross-encoder vs MNLI vs BART orderings) --

@pytest.mark.parametrize("id2label,expect", [
    ({0: "entailment", 1: "neutral", 2: "contradiction"}, (0, 1, 2)),   # MNLI/MoritzLaurer order
    ({0: "contradiction", 1: "entailment", 2: "neutral"}, (1, 2, 0)),   # cross-encoder order
    ({0: "contradiction", 1: "neutral", 2: "entailment"}, (2, 1, 0)),   # BART order
    ({0: "ENTAILMENT", 1: "NEUTRAL", 2: "CONTRADICTION"}, (0, 1, 2)),   # case-insensitive
])
def test_resolve_label_indices(id2label, expect):
    assert _resolve_label_indices(id2label) == expect


def test_resolve_label_indices_rejects_incomplete():
    with pytest.raises(RuntimeError):
        _resolve_label_indices({0: "entailment", 1: "neutral"})


# --- ShadowNLI / ShadowResult contract -------------------------------------

def test_shadownli_construction_never_raises_and_starts_unavailable():
    nli = ShadowNLI()
    assert nli.available is False           # not loaded until first check
    assert nli.model_version == load_manifest()["nli_model_version"]
    assert nli.theta_contra == 0.5 and nli.theta_entail == 0.9
    assert nli.device_pref in ("auto", "cpu", "cuda")  # default 'auto'
    assert nli.batch_size >= 1
    assert nli.backend is None              # set at load → "torch" / "onnx" / "onnx-int8"
    assert nli._onnx_dir().name == "onnx"   # conventional export location


def test_shadownli_device_and_onnx_dir_env_overrides(monkeypatch, tmp_path):
    monkeypatch.setenv("ARBORIST_NLI_DEVICE", "cpu")
    monkeypatch.setenv("ARBORIST_NLI_ONNX_DIR", str(tmp_path / "x"))
    monkeypatch.setenv("ARBORIST_NLI_BATCH", "8")
    nli = ShadowNLI()
    assert nli.device_pref == "cpu" and nli.batch_size == 8
    assert nli._onnx_dir() == tmp_path / "x"


def test_nli_batch_matches_single_when_available():
    nli = ShadowNLI()
    nli._ensure_loaded()
    if not nli.available:
        pytest.skip("[nli] extra not installed")
    pairs = [("Jupiter is the largest planet.", "Mercury is the largest planet."),
             ("Paris is the capital of France.", "Paris is the capital of France.")]
    batched = nli._nli_batch(pairs)
    singles = [nli._nli(p, h) for p, h in pairs]
    for b, s in zip(batched, singles):
        # batching pads the shorter sequence; with attention masking that's
        # ~a no-op for fp32, and within quantization noise for int8 ONNX —
        # the operative thing (argmax label, gate decision) must not move.
        assert max(range(3), key=lambda i: b[i]) == max(range(3), key=lambda i: s[i])
        assert all(abs(x - y) < 0.05 for x, y in zip(b, s))


def test_check_returns_shadowresult_and_degrades_gracefully():
    res = shadow_check("Mercury is the largest planet.",
                       "Jupiter is the largest planet. Mercury is the smallest planet.")
    assert isinstance(res, ShadowResult)
    assert res.would_demote in (True, False)
    assert 0.0 <= res.max_contradiction <= 1.0
    assert 0.0 <= res.max_entailment <= 1.0
    assert res.n_clauses == 2
    d = res.as_dict()
    assert set(d) >= {"available", "would_demote", "max_contradiction", "max_entailment",
                      "best_clause", "n_clauses", "model_version", "theta_contra",
                      "theta_entail", "reason"}
    if not res.available:
        # deps-missing path: a clean, non-raising, no-signal result
        assert res.would_demote is False
        assert res.max_contradiction == 0.0 and res.max_entailment == 0.0
        assert "deps_missing" in res.reason or "load_failed" in res.reason


def test_check_empty_input_when_available_is_no_demote():
    nli = ShadowNLI()
    nli._ensure_loaded()
    if not nli.available:
        pytest.skip("[nli] extra not installed — empty-input branch only reachable when loaded")
    res = nli.check("", "Some source text. Another clause.")
    assert res.available is True and res.would_demote is False and res.reason == "empty_input"


# --- bench sweep script ----------------------------------------------------

def test_shadow_sweep_parses_both_record_shapes(tmp_path):
    from bench.scripts.nli_shadow_sweep import _records
    p = tmp_path / "mixed.jsonl"
    p.write_text("\n".join([
        json.dumps({"_meta": {"battery": "5f"}}),
        json.dumps({"id": "a", "answer_text": "X is Y.", "context": "X is Z. Q is Y.", "expected_reason": "UNGROUNDED"}),
        json.dumps({"id": "b", "claim": "P is Q.", "source": "P is R.", "want": "not_contradiction"}),
        json.dumps({"id": "c", "answer_text": "only answer, no context"}),  # incomplete → skipped
    ]))
    recs = list(_records(p))
    assert [r["id"] for r in recs] == ["a", "b"]
    assert recs[0]["bucket"] == "UNGROUNDED" and recs[0]["is_fp_probe"] is False
    assert recs[1]["bucket"] == "not_contradiction" and recs[1]["is_fp_probe"] is True


def test_shadow_sweep_main_smoke(tmp_path):
    from bench.scripts.nli_shadow_sweep import main
    inp = tmp_path / "in.jsonl"
    inp.write_text("\n".join([
        json.dumps({"id": "r1", "claim": "Mercury is the largest planet.",
                    "source": "Jupiter is the largest planet. Mercury is the smallest planet.",
                    "want": "contradiction"}),
        json.dumps({"id": "r2", "claim": "Batman, the alias of Bruce Wayne, lives in Gotham.",
                    "source": "Batman is the alias of Bruce Wayne. Batman lives in Gotham.",
                    "want": "not_contradiction"}),
    ]))
    out = tmp_path / "report.json"
    rc = main(["--input", str(inp), "--out", str(out)])
    assert rc == 0
    rep = json.loads(out.read_text())
    assert rep["n_records"] == 2
    assert "by_bucket" in rep and "false_positive_probe" in rep
    assert set(rep["by_bucket"]) == {"contradiction", "not_contradiction"}
    # when [nli] absent the report still renders, just unmeasured
    assert rep["available"] in (True, False)