arborist/tests/test_judge_code.py

"""Tests for the deterministic code-only judge (``bench/judge_code.py``).

The self-test fixture set is the BEHAVIOURAL CONTRACT — these are
exactly the 4 cases the Opus judge (``bench/judge.py``) was hand-built
to grade correctly. A code judge that disagrees on these is not a
trustworthy instrument. CI fails the build if the contract slips.

We deliberately do NOT pin specific NLI probabilities or specific
``rules_fired`` traces — those are tuning knobs we expect to move.
The label is the public contract; the trace is implementation detail.

Runs WITHOUT the ``[nli]`` extra installed too: NLI is optional, the
abstention / lexical / specificity paths give correct labels on the
fixture set even when NLI is unavailable (case 1 then routes via
JUDGE_ERROR — see ``test_self_test_passes_with_nli_or_skips``).
"""
from __future__ import annotations

import pytest

from bench.judge_code import (
    JUDGE_MODEL,
    JUDGE_RULESET_ID,
    Verdict,
    judge,
    judge_batch,
    self_test,
    _is_abstention,
    _specifics_not_in_gold,
    _unwrap_claim_lattice_json,
)


# --- module surface --------------------------------------------------------

def test_module_identifiers_pinned():
    """Audit-trail parity with bench/judge.py: stable model + prompt
    id strings so verdict logs from this judge are sortable / joinable
    against logs from the Opus judge."""
    assert JUDGE_RULESET_ID == "code-judge-v1"
    assert "code" in JUDGE_MODEL.lower()
    assert "no llm" in JUDGE_MODEL.lower()


def test_verdict_dataclass_shape():
    v = judge("Q?", "I cannot determine that.", "gold")
    assert isinstance(v, Verdict)
    # Same fields as bench.judge.Verdict so the two are drop-in compat.
    assert v.label in {"CORRECT_GROUNDED", "WRONG", "FABRICATED",
                       "ABSTAINED", "JUDGE_ERROR"}
    assert v.prompt_id == JUDGE_RULESET_ID
    assert v.model == JUDGE_MODEL
    assert isinstance(v.rationale, str) and v.rationale
    assert isinstance(v.raw, str) and v.raw  # JSON-serialised trace
    assert isinstance(v.decision, dict)


# --- rule 1: empty / no-gold guards ----------------------------------------

def test_empty_answer_is_abstained():
    v = judge("Q?", "", "gold text")
    assert v.label == "ABSTAINED"
    assert "empty_answer" in v.decision["rules_fired"]


def test_empty_gold_is_judge_error():
    v = judge("Q?", "Some claim.", "")
    assert v.label == "JUDGE_ERROR"
    assert "no_gold" in v.decision["rules_fired"]


# --- rule 2: explicit abstention -------------------------------------------

@pytest.mark.parametrize("phrase", [
    "I cannot determine that from the provided source.",
    "I do not know.",
    "I can't tell from the article.",
    "The source does not mention this.",
    "The article doesn't contain that information.",
    "Not enough information is provided.",
    "Unable to determine from the passage.",
    "Cannot be determined from the source.",
])
def test_abstention_phrases_classify_as_abstained(phrase):
    v = judge("Q?", phrase, "irrelevant gold content")
    assert v.label == "ABSTAINED", phrase
    assert "abstention_phrase" in v.decision["rules_fired"]


def test_assertive_answer_is_not_abstention():
    """A model that makes a confident claim is NOT abstaining, even if
    the claim happens to be ungrounded — that's WRONG / FABRICATED,
    not ABSTAINED."""
    is_abs, _ = _is_abstention(
        "The capital of France is Berlin and the population is 12 million."
    )
    assert is_abs is False


# --- rule 4: specificity (the FABRICATED gate) -----------------------------

def test_specifics_picked_up_when_not_in_gold():
    answer = "The Higgs boson cafe was founded by Napoleon in 1066."
    gold = "This article discusses thermodynamics and entropy."
    s = _specifics_not_in_gold(answer, gold)
    # Should flag the proper noun(s) and the date.
    joined = " ".join(s).lower()
    assert "napoleon" in joined
    assert "1066" in joined


def test_specifics_excluded_when_in_gold():
    """A proper noun / number that ALSO appears in gold is grounded
    by definition (case-insensitive) and must NOT be flagged."""
    answer = "Napoleon led the army in 1812."
    gold = "Napoleon Bonaparte's campaign in 1812 was disastrous."
    s = _specifics_not_in_gold(answer, gold)
    joined = " ".join(s).lower()
    assert "napoleon" not in joined
    assert "1812" not in joined


# --- the canonical behavioural contract: 4-case self-test ------------------

def test_self_test_passes():
    """The same 4 fixtures the Opus judge was hand-built to grade
    correctly. If this fails, the code judge is not an interchangeable
    instrument."""
    assert self_test() == 0


# --- batch parity ----------------------------------------------------------

def test_judge_batch_loops_correctly():
    triples = [
        ("Q1", "I cannot determine.", "gold a"),
        ("Q2", "", "gold b"),
    ]
    out = judge_batch(triples)
    assert [v.label for v in out] == ["ABSTAINED", "ABSTAINED"]
    assert all(isinstance(v, Verdict) for v in out)


# --- claim-lattice JSON unwrap ---------------------------------------------

def test_unwrap_single_claim_lattice():
    """The Arborist arm's `claim_lattice` answer_mode emits a JSON
    envelope; the judge must unwrap to the inner claim prose so the
    verifier's span/entity paths see plain text, not brace syntax."""
    a = ('{"claims":[{"text":"The current president of Iceland is '
         'Ólafur Ragnar Grímsson.","evidence_ids":["E1"]}]}')
    u = _unwrap_claim_lattice_json(a)
    assert u == "The current president of Iceland is Ólafur Ragnar Grímsson."


def test_unwrap_multi_claim_concatenates_with_periods():
    """Multiple claims → concatenated sentences, with periods inserted
    where the claim text doesn't already terminate. extract_claim_spans
    treats each as its own span."""
    a = ('{"claims":[{"text":"X is A.","evidence_ids":["E1"]},'
         '{"text":"Y is B","evidence_ids":["E2"]}]}')
    u = _unwrap_claim_lattice_json(a)
    assert u == "X is A. Y is B."


@pytest.mark.parametrize("plain", [
    "Just a regular answer.",
    "",
    "I cannot determine that.",
    "Napoleon founded the cafe in 1066.",
])
def test_unwrap_passes_plain_prose_through(plain):
    """Non-claim-lattice answers must pass through unchanged so the
    judge's behaviour on solo-arm answers is unaffected."""
    assert _unwrap_claim_lattice_json(plain) == plain


def test_unwrap_passes_malformed_json_through():
    """Malformed JSON is left intact so the downstream verifier can
    see what it gets — no silent rewriting on broken input."""
    a = '{"claims":[broken json'
    assert _unwrap_claim_lattice_json(a) == a


def test_unwrap_passes_unrelated_json_through():
    """JSON that doesn't have the claim-lattice signature (claims +
    text keys) is not the Arborist envelope; pass through."""
    a = '{"answer": "X is A", "confidence": 0.9}'
    assert _unwrap_claim_lattice_json(a) == a


def test_arborist_envelope_correctly_grounded():
    """End-to-end: Arborist claim_lattice envelope around a clearly
    correct claim grades to CG (was ABSTAINED / WRONG before the
    unwrap)."""
    gold = ("{{Infobox Political post |post = President |body = Iceland "
            "|incumbent = [[Ólafur Ragnar Grímsson]] "
            "|incumbentsince = 1 August 1996")
    answer = ('{"claims":[{"text":"The current president of Iceland is '
              'Ólafur Ragnar Grímsson.","evidence_ids":["E1"]}]}')
    v = judge("who is the president of Iceland?", answer, gold)
    assert v.label == "CORRECT_GROUNDED"
    assert v.decision.get("lattice_unwrapped") is True


# --- graceful NLI degradation ----------------------------------------------

def test_judge_handles_nli_unavailable(monkeypatch):
    """When the [nli] extra is missing, shadow_check returns an
    unavailable ShadowResult — the judge must NOT crash; it just
    proceeds without rule-3 (contradiction) signal and without the
    HYBRID corroboration path. Labels stay sound for the explicit
    abstention / UNGROUNDED-specifics paths (which short-circuit
    before NLI is even consulted)."""
    import bench.judge_code as jc

    monkeypatch.setattr(jc, "_nli_check", lambda a, g: None)

    # Abstention path returns at rule 2, BEFORE the NLI block runs —
    # so the trace doesn't (and shouldn't) carry nli_available. The
    # contract here is "no crash + correct label", not "NLI metadata
    # is populated on every path".
    v = jc.judge("Q?", "I cannot determine.", "gold")
    assert v.label == "ABSTAINED"
    assert "abstention_phrase" in v.decision["rules_fired"]
    assert "nli_available" not in v.decision  # rule 2 short-circuit

    # An answer that reaches the NLI block must record nli_available
    # = False when _nli_check returns None.
    v = jc.judge("Q?", "Some claim about something.", "different content here")
    assert v.decision.get("nli_available") is False

    # FABRICATED path still works (no NLI needed).
    v = jc.judge(
        "Who founded the cafe?",
        "Napoleon founded it in 1066.",
        "This article is about thermodynamics.",
    )
    assert v.label == "FABRICATED"


# --- relaxed HYBRID rescue (2026-05-21): tolerate unsourced proper nouns
# (aliases/paraphrase) but still block unsourced numerics (date/count
# errors), gated on a verified quote + on-topic + no contradiction. -----

_EPIRUS_GOLD = ("Alexander I of Epirus was a king of Epirus from 350 to "
                "331 BC, belonging to the Aeacid dynasty. He was the son "
                "of Neoptolemus I and brother of Olympias.")


def test_hybrid_rescued_when_only_unsourced_specific_is_proper_noun(monkeypatch):
    # NLI off so the lexical tier is isolated (not the NLI corroboration
    # path). One verbatim quote (verified) + one quote NOT in gold
    # (-> HYBRID); the only unsourced specifics are proper nouns.
    import bench.judge_code as jc
    monkeypatch.setattr(jc, "_nli_check", lambda a, g: None)
    answer = ('Alexander I of Epirus "was a king of Epirus from 350 to '
              '331 BC". Also called Alexander Molossus, he "led famed '
              'campaigns in southern Italy against the Romans".')
    v = jc.judge("who was Alexander the first of epirus?", answer, _EPIRUS_GOLD)
    assert v.label == "CORRECT_GROUNDED", (v.label, v.rationale)
    assert v.decision.get("audit_mode") == "HYBRID"
    assert "hybrid+verified_quote_on_topic" in v.decision["rules_fired"]


def test_hybrid_stays_residue_when_unsourced_specific_is_numeric(monkeypatch):
    # NLI off. One verified quote + an unverified quote with a WRONG date
    # -> unsourced numeric -> NOT rescued (date/count errors stay residue).
    import bench.judge_code as jc
    monkeypatch.setattr(jc, "_nli_check", lambda a, g: None)
    answer = ('Alexander I of Epirus "was a king of Epirus". He '
              '"reigned from 999 to 888 BC".')
    v = jc.judge("who was Alexander the first of epirus?", answer, _EPIRUS_GOLD)
    assert v.label == "JUDGE_ERROR", (v.label, v.rationale)
    assert "hybrid+verified_quote_on_topic" not in v.decision["rules_fired"]