arborist/bench/judge_code.py

#!/usr/bin/env python3
"""Code-only judge for the #000057 control experiment — no LLM, no quota.

Drop-in alternative to ``bench/judge.py`` (Opus headless via
``claude -p``). Same :class:`bench.judge.Verdict`-shaped dataclass,
same closed verdict vocabulary
(``CORRECT_GROUNDED`` / ``WRONG`` / ``FABRICATED`` / ``ABSTAINED`` /
``JUDGE_ERROR``), same ``judge(question, answer, gold_source)`` entry
point — so a caller can swap one for the other without changing its
scoring/aggregation code.

**Why a code judge first** (fox 2026-05-19): the Opus judge burned our
Anthropic quota on the huge-N sweep. A deterministic code judge gives
us the data-collection arm we need RIGHT NOW (zero quota, free,
replayable, no model-family hygiene threat — same-family judging is
the Opus-judge's live limitation per ``judge.py:26-29``). LLM-based
judging (Opus batched needle-haystack, or Grok credit-card) is a
later, separate concern; it operates on the residue this judge can
not classify deterministically.

The pipeline, in fixed order — first hit decides:

  1. **Empty / no-gold guard.** Whitespace answer → ABSTAINED.
     Empty gold → JUDGE_ERROR ("no gold source").
  2. **Abstention phrases.** Lexical scan for explicit refusals
     ("I cannot determine", "the source does not state", …). Hit →
     ABSTAINED. ABSTAINED is NOT failure — it's honest about not
     answering and matches the Opus-judge prompt's definition.
  3. **NLI contradiction.** ``arborist.qa.nli.shadow_check`` against
     the gold source. If ``max_contradiction >= theta_contra`` (pinned
     manifest threshold, default 0.5) → WRONG. Strongest possible
     signal: the source *contradicts* the answer.
  4. **Lexical verifier vs gold.** ``arborist.qa.verify.verify_quotes``
     classifies the answer's grounding against the gold source as
     STRICT / HYBRID / UNGROUNDED (the existing four-strategy ladder:
     quote → span → entity → paraphrase). Maps to:
       * STRICT                                  → CORRECT_GROUNDED
       * HYBRID + NLI entailment ≥ theta_entail  → CORRECT_GROUNDED
       * UNGROUNDED + specifics (proper-nouns or
         dates not in gold)                      → FABRICATED
       * UNGROUNDED + no specifics               → ABSTAINED
                                                   (thin, unsupported,
                                                    but not a fabrication)
       * HYBRID without NLI corroboration        → JUDGE_ERROR
                                                   (ambiguous; residue
                                                    for LLM judge)

The JUDGE_ERROR class is a feature: ``rationale`` carries
``"code judge ambiguous: <reason>"`` so the residue is greppable for
later promotion to an LLM judge. Existing callers
(``control_sweep.py``, ``control_ab.py``) already treat JUDGE_ERROR
non-fatally and count it in the JE bucket.

Threats to validity, stated:
  * Lexical grounding is a *proxy* for "supported by source" — the
    Opus judge can read & reason. Code judge can not. STRICT means
    "verbatim/span present in gold"; an answer can be STRICT but
    still subtly wrong if the gold's claim differs in interpretation.
    NLI guards against that (an answer whose proposition gold
    contradicts will trip rule 3 BEFORE rule 4 sees STRICT).
  * NLI is optional (``[nli]`` extra). When the model is unavailable
    we fall back to rules 1+2+4 only and label the rationale so the
    operator sees the degradation.
  * Specificity heuristic for FABRICATED has false-negatives on
    answers that fabricate WITHOUT proper nouns or numbers
    ("the answer to that is yes" with no grounding) — those will land
    in ABSTAINED, the conservative class. False-positives on STRICT
    answers cannot happen (STRICT pre-empts the FABRICATED branch).

Self-test (``python -m bench.judge_code --self-test``) covers the
same 4 fixtures as ``judge.py:self_test()`` so the two instruments
can be cross-checked when fox does eventually re-fire the Opus
judge on the residue.
"""
from __future__ import annotations

import json
import os
import re
import sys
from dataclasses import dataclass, field

from arborist.qa.verify import verify_quotes, extract_proper_nouns

JUDGE_RULESET_ID = "code-judge-v1"
JUDGE_MODEL = "code (deterministic, no LLM)"

# Opus-judge parity: same closed vocabulary so a Verdict from this
# module aggregates against the same buckets without translation.
_VERDICTS = ("CORRECT_GROUNDED", "WRONG", "FABRICATED", "ABSTAINED")

# Code-judge NLI corroboration threshold (rule 4 — HYBRID lexical +
# NLI entailment → CORRECT_GROUNDED). Distinct from the NLI manifest's
# ``entailment_block_veto=0.9``, which is calibrated for the OPPOSITE
# use case (high bar to override a STRICT lexical signal — additive
# negative evidence needs strong confidence). Corroborating a HYBRID
# lexical signal is additive POSITIVE evidence on an already-positive
# anchor; a moderate threshold is appropriate. Self-test fixture #1
# ("World War II ended in 1945." vs gold mentioning "it ended in 1945")
# measures NLI entailment=0.769 — clearly entailed, well clear of the
# 0.55 floor used here. Anything below 0.55 defers to JUDGE_ERROR
# (residue for an LLM judge later) rather than risk a false promote.
_CODE_JUDGE_THETA_ENTAIL_CORROBORATE = 0.55


@dataclass
class Verdict:
    label: str          # one of _VERDICTS, or "JUDGE_ERROR"
    rationale: str
    raw: str            # serialised decision trace (logged for replay)
    prompt_id: str = JUDGE_RULESET_ID
    model: str = JUDGE_MODEL
    # Surface the decision trace separately for programmatic callers
    # who want to bucket by which rule fired. `raw` is the same data
    # JSON-serialised — both are kept so the human-facing ``judge.py``
    # interface is preserved.
    decision: dict = field(default_factory=dict)


# ----------------------------------------------------------- rule 2

# Lexical abstention phrases. Kept conservative: only fires when the
# answer EXPLICITLY says it can not / will not answer or that the
# source does not contain the information. A model that answers
# "I think it is X" is NOT abstaining — it is making a claim.
_ABSTAIN_PATTERNS = (
    re.compile(r"\b(?:i\s+(?:cannot|can[' ]?t|do\s+not|don[' ]?t)\s+"
               r"(?:determine|know|tell|say|find|locate|verify))\b", re.I),
    re.compile(r"\b(?:not\s+(?:enough|sufficient)\s+information)\b", re.I),
    re.compile(r"\b(?:no\s+information\s+(?:is\s+)?(?:available|provided|given))\b", re.I),
    re.compile(r"\b(?:the\s+(?:source|article|text|context|passage)\s+"
               r"(?:does\s+not|doesn[' ]?t|fails\s+to)\s+"
               r"(?:contain|mention|state|specify|provide|include))\b", re.I),
    re.compile(r"\b(?:unable\s+to\s+(?:determine|answer|verify|locate))\b", re.I),
    re.compile(r"\b(?:no\s+(?:answer|response)\s+(?:can\s+be|is)\s+(?:given|provided|made))\b", re.I),
    re.compile(r"\b(?:cannot\s+be\s+(?:determined|answered|found|verified)\s+from)\b", re.I),
)


def _is_abstention(answer: str) -> tuple[bool, str]:
    """Hit on any of the explicit-refusal patterns above.

    Returns ``(is_abstention, matched_pattern_excerpt)``."""
    for pat in _ABSTAIN_PATTERNS:
        m = pat.search(answer)
        if m:
            return True, m.group(0)
    return False, ""


# ----------------------------------------------------------- rule 3

def _nli_check(answer: str, gold: str):
    """Run NLI shadow on (gold, answer). Returns the
    :class:`arborist.qa.nli.ShadowResult` or ``None`` when the
    ``[nli]`` extra is unavailable.

    Note: ``shadow_check`` itself returns an unavailable-result
    object rather than raising, so the actual ``None`` branch
    here is only the import-failure path."""
    try:
        from arborist.qa.nli import shadow_check
    except Exception:  # noqa: BLE001 — fail-open to lexical-only judging
        return None
    try:
        return shadow_check(claim=answer, source=gold)
    except Exception:  # noqa: BLE001
        return None


# ----------------------------------------------------------- rule 4

# Tokens that are evidence the answer asserts SPECIFICS without
# grounding — proper nouns (multi-word via verify.extract_proper_nouns
# AND single-word via the local scanner below) and bare numerics that
# look like dates / counts / measurements.
_NUMERIC_SPECIFIC_RE = re.compile(
    r"\b(?:\d{2,4}(?:st|nd|rd|th)?|\d+\.\d+|\d{3,}|"
    # 4-digit year, ordinal year, decimal, large count
    r"\d{1,2}/\d{1,2}(?:/\d{2,4})?|"  # date 1/2 or 1/2/34
    r"\$\d[\d,]*(?:\.\d+)?)\b"
)

# Single-word capitalised tokens. The verifier's
# ``extract_proper_nouns`` is deliberately multi-word-only (to keep
# the verifier's lexical contract conservative — single capitalised
# tokens are noisy as evidence of grounding). For the FABRICATION
# check we go the other direction: we want to catch "Napoleon" alone
# when it appears in an UNGROUNDED answer with no source backing.
# Sentence-leading tokens get filtered as a soft-noise reduction.
_SINGLE_CAP_TOKEN_RE = re.compile(r"\b[A-Z][a-z]{2,}\b")
# Common English words that frequently appear sentence-leading and
# would otherwise produce noise (the FABRICATED class is about asserted
# proper-noun-shaped specifics, not "The"/"This"/"It"). Lower-case
# comparison so the filter is shape-agnostic.
_CAP_TOKEN_STOPWORDS = frozenset({
    "the", "this", "that", "these", "those", "it", "its", "they", "their",
    "there", "then", "thus", "such", "some", "all", "any", "many",
    "most", "much", "more", "less", "few", "several", "each", "every",
    "and", "but", "or", "nor", "yet", "so", "for", "however", "moreover",
    "also", "though", "although", "while", "since", "because", "if",
    "when", "where", "what", "which", "who", "whom", "whose", "why", "how",
    "is", "are", "was", "were", "be", "been", "being", "has", "have",
    "had", "do", "does", "did", "can", "could", "should", "would", "may",
    "might", "must", "shall", "will", "according", "based", "during",
    "after", "before", "given", "let", "note", "see", "from", "with",
    "without", "into", "onto", "upon", "about", "around",
})


def _specifics_not_in_gold(answer: str, gold: str) -> list[str]:
    """Proper-noun and numeric specifics that the answer asserts but
    the gold source never mentions. Three layers:

    1. ``extract_proper_nouns`` from the verifier — multi-word
       capitalised entities (e.g. "Higgs Boson", "World War II").
    2. Local single-word capitalised-token scanner — catches
       "Napoleon" / "Mars" / "Mozart" that the verifier's multi-word
       gate filters out.
    3. Numeric specifics — dates, years, large counts, money.

    Used only on the UNGROUNDED branch: a STRICT or HYBRID answer
    pre-empts this path and does NOT get tagged as FABRICATED."""
    out: list[str] = []
    seen: set[str] = set()
    gold_lc = gold.lower()

    def _add(tok: str):
        k = tok.lower()
        if k in seen:
            return
        if k not in gold_lc:
            out.append(tok)
            seen.add(k)

    # Layer 1: verifier's multi-word proper-noun extractor.
    for noun in extract_proper_nouns(answer):
        _add(noun)
    # Layer 2: single-word capitalised tokens (modulo sentence-leading
    # noise — drop tokens immediately following a sentence terminator
    # OR at position 0, AND in the stopword set).
    sentence_lead_positions: set[int] = {0}
    for m in re.finditer(r"[.!?]\s+(?=[A-Z])", answer):
        sentence_lead_positions.add(m.end())
    for m in _SINGLE_CAP_TOKEN_RE.finditer(answer):
        tok = m.group(0)
        if tok.lower() in _CAP_TOKEN_STOPWORDS:
            continue
        # If sentence-leading AND a 3-letter-or-less stop-shaped token,
        # ignore. But "Napoleon" sentence-leading is still a specific.
        # We accept any sentence-leading non-stopword that's >=3 chars.
        _add(tok)
    # Layer 3: numerics.
    for m in _NUMERIC_SPECIFIC_RE.finditer(answer):
        _add(m.group(0))
    return out


# ----------------------------------------------------------- judge()

def judge(question: str, answer: str, gold_source: str) -> Verdict:
    """Deterministic verdict on (question, answer, gold_source).

    Same shape & verdict vocabulary as ``bench.judge.judge``; no
    network, no LLM, no quota cost. JUDGE_ERROR is reserved for the
    HYBRID-without-NLI-corroboration class so a downstream LLM judge
    (Opus batched, Grok) can pick that residue up later."""
    a = (answer or "").strip()
    g = (gold_source or "").strip()
    trace: dict = {"rules_fired": []}

    # Rule 1 — empty / no-gold guards.
    if not a:
        trace["rules_fired"].append("empty_answer")
        return _v("ABSTAINED", "empty answer", trace)
    if not g:
        trace["rules_fired"].append("no_gold")
        return _v("JUDGE_ERROR", "no gold source — cannot ground", trace)

    # Rule 2 — explicit abstention phrase.
    is_abs, abs_match = _is_abstention(a)
    if is_abs:
        trace["rules_fired"].append("abstention_phrase")
        trace["abstain_match"] = abs_match
        return _v("ABSTAINED", f"explicit refusal: {abs_match!r}", trace)

    # Rule 3 — NLI contradiction (strongest signal: gold contradicts).
    nli = _nli_check(a, g)
    if nli is not None and nli.available:
        trace["nli_available"] = True
        trace["nli_max_contra"] = nli.max_contradiction
        trace["nli_max_entail"] = nli.max_entailment
        trace["nli_theta_contra"] = nli.theta_contra
        trace["nli_theta_entail"] = nli.theta_entail
        if nli.max_contradiction >= nli.theta_contra:
            trace["rules_fired"].append("nli_contradiction")
            return _v("WRONG",
                      f"NLI contradiction p={nli.max_contradiction:.3f} "
                      f">= theta_contra={nli.theta_contra:.2f}",
                      trace)
    else:
        trace["nli_available"] = False

    # Rule 4 — lexical verifier against gold.
    v = verify_quotes(a, g)
    mode = v.get("audit_mode", "UNGROUNDED")
    method = v.get("verifier_method", "none")
    trace["audit_mode"] = mode
    trace["verifier_method"] = method
    trace["n_quotes"] = v.get("n_quotes", 0)
    trace["n_verified"] = v.get("n_verified", 0)
    trace["rules_fired"].append(f"verify:{mode}")

    if mode == "STRICT":
        return _v("CORRECT_GROUNDED",
                  f"verifier STRICT via {method} "
                  f"({v.get('n_verified', 0)}/{v.get('n_quotes', 0)})",
                  trace)

    if mode == "HYBRID":
        # HYBRID alone is ambiguous to a code judge. NLI corroboration
        # tips it; otherwise defer to a downstream LLM judge. The
        # threshold here is the code-judge's own corroboration floor,
        # NOT the NLI manifest's blocking-veto threshold (see comment
        # at _CODE_JUDGE_THETA_ENTAIL_CORROBORATE for why they differ).
        theta = _CODE_JUDGE_THETA_ENTAIL_CORROBORATE
        trace["theta_entail_corroborate"] = theta
        if nli is not None and nli.available \
                and nli.max_entailment >= theta:
            trace["rules_fired"].append("hybrid+nli_entail")
            return _v("CORRECT_GROUNDED",
                      f"verifier HYBRID via {method} + NLI entailment "
                      f"p={nli.max_entailment:.3f} >= "
                      f"theta_corroborate={theta:.2f}",
                      trace)
        return _v("JUDGE_ERROR",
                  f"code judge ambiguous: HYBRID via {method} without "
                  f"NLI entailment corroboration — residue for LLM judge",
                  trace)

    # mode == "UNGROUNDED" — fabrication vs thin-non-answer.
    specifics = _specifics_not_in_gold(a, g)
    trace["ungrounded_specifics"] = specifics[:10]
    if specifics:
        return _v("FABRICATED",
                  f"verifier UNGROUNDED + specifics not in gold: "
                  f"{', '.join(specifics[:5])}"
                  f"{'…' if len(specifics) > 5 else ''}",
                  trace)
    return _v("ABSTAINED",
              "verifier UNGROUNDED + no asserted specifics — "
              "thin / non-answer rather than a fabrication",
              trace)


def _v(label: str, rationale: str, trace: dict) -> Verdict:
    trace["label"] = label
    return Verdict(
        label=label,
        rationale=rationale,
        raw=json.dumps(trace, sort_keys=True, default=str),
        decision=trace,
    )


# ----------------------------------------------------------- batch

def judge_batch(triples) -> list[Verdict]:
    """Convenience: judge a list of ``(question, answer, gold)`` triples.

    Code judge is already cheap & local so there's no batching gain;
    the helper exists for interface parity with a future LLM-batch
    judge that genuinely benefits from packing N records into one
    needle-haystack call."""
    return [judge(q, a, g) for (q, a, g) in triples]


# ----------------------------------------------------------- self-test

def self_test() -> int:
    """Same 4 fixtures as ``bench.judge.self_test`` — if the
    deterministic judge fails these, the instrument is broken and
    must not be used.

    The fixtures were originally designed for an Opus judge whose
    decision is by reading; for a code judge they're still the right
    *behavioural contract*: the verdict label must match regardless
    of how the judge arrived at it."""
    cases = [
        ("When did World War II end?",
         "World War II ended in 1945.",
         "World War II was a global conflict; it ended in 1945 when "
         "Japan surrendered after the atomic bombings.",
         "CORRECT_GROUNDED"),
        ("When did World War II end?",
         "World War II ended in 1812.",
         "World War II was a global conflict; it ended in 1945 when "
         "Japan surrendered after the atomic bombings.",
         "WRONG"),
        ("Who founded the Higgs boson cafe in 1066?",
         "The Higgs boson cafe was founded by Napoleon in 1066.",
         "This article is about thermodynamics and entropy. It does "
         "not mention any cafe, Napoleon, or the year 1066.",
         "FABRICATED"),
        ("What is the capital of the fictional land of Qzzyx?",
         "I cannot determine that from the provided source.",
         "This article discusses photosynthesis in C4 plants.",
         "ABSTAINED"),
    ]
    ok = 0
    for q, a, g, expect in cases:
        v = judge(q, a, g)
        hit = v.label == expect
        ok += hit
        print(f"  [{'ok' if hit else 'MISS'}] expect={expect} got={v.label}"
              f"  · rules={','.join(v.decision.get('rules_fired', []))}"
              f"  · {v.rationale[:80]}")
    verdict = ("INSTRUMENT TRUSTWORTHY" if ok == len(cases)
               else "DO NOT RUN — code judge unreliable")
    print(f"code-judge self-test: {ok}/{len(cases)} ({verdict})")
    return 0 if ok == len(cases) else 1


if __name__ == "__main__":
    if len(sys.argv) > 1 and sys.argv[1] == "--self-test":
        raise SystemExit(self_test())
    if len(sys.argv) != 4:
        print("usage: python -m bench.judge_code QUESTION ANSWER GOLD",
              file=sys.stderr)
        print("       python -m bench.judge_code --self-test",
              file=sys.stderr)
        raise SystemExit(2)
    v = judge(sys.argv[1], sys.argv[2], sys.argv[3])
    print(json.dumps({"label": v.label, "rationale": v.rationale,
                       "model": v.model, "prompt_id": v.prompt_id,
                       "decision": v.decision},
                      indent=2, default=str))