arborist/tests/test_verify_json.py

"""JSON-mode claim-lattice verifier (`answer_mode="claim_lattice"`).

Pairs with grammar-constrained inference (vLLM guided_json, Claude/GPT-4
native JSON, Qwen 3.6 reasoner). The lenient pre-parser keeps the path
survivable on inference paths where the model emits non-strict JSON
(markdown fences, prose preamble, curly quotes, trailing commas).
"""

from __future__ import annotations

import json

import pytest

from aborist.qa.evidence import EvidenceObject
from aborist.qa.verify import _lenient_json_parse, verify_claim_lattice_json


# ---------------------------------------------------------------- lenient parser


def test_lenient_strict_passes_through():
    obj, fixups = _lenient_json_parse('{"a": 1}')
    assert obj == {"a": 1}
    assert fixups == []


def test_lenient_strips_markdown_fence():
    raw = '```json\n{"claims": []}\n```'
    obj, fixups = _lenient_json_parse(raw)
    assert obj == {"claims": []}
    assert "fence" in fixups


def test_lenient_strips_unlabeled_fence():
    raw = '```\n{"claims": []}\n```'
    obj, fixups = _lenient_json_parse(raw)
    assert obj == {"claims": []}
    assert "fence" in fixups


def test_lenient_trims_preamble_and_suffix():
    raw = 'Here is the JSON:\n{"claims": []}\nLet me know if you need more.'
    obj, fixups = _lenient_json_parse(raw)
    assert obj == {"claims": []}
    assert "prose_trim" in fixups


def test_lenient_normalizes_curly_quotes():
    raw = '{“claims”: []}'
    obj, fixups = _lenient_json_parse(raw)
    assert obj == {"claims": []}
    assert "curly_quotes" in fixups


def test_lenient_fixes_trailing_comma():
    raw = '{"claims": [],}'
    obj, fixups = _lenient_json_parse(raw)
    assert obj == {"claims": []}
    assert "trailing_comma" in fixups


def test_lenient_combines_multiple_fixups():
    raw = '```json\nHere:\n{“claims”: [],}\n```'
    obj, fixups = _lenient_json_parse(raw)
    assert obj == {"claims": []}
    assert "fence" in fixups
    # At least one of the inner fixups also fired.
    assert any(f in fixups for f in ("curly_quotes", "trailing_comma", "prose_trim"))


def test_lenient_raises_on_truly_broken():
    with pytest.raises(json.JSONDecodeError):
        _lenient_json_parse("not json at all { ] [")


# ---------------------------------------------------------------- JSON verifier


def _ev(
    eid: str,
    span: str,
    role: str = "primary_answer_source",
    pointer_id: str | None = None,
) -> EvidenceObject:
    """Stub evidence object with deterministic eid for the test.

    JSON-mode prompt-facing surface uses ``pointer_id`` (E1, E2, …)
    since 2026-04-30; if the test wants the verifier to resolve a
    citation, it must set ``pointer_id`` explicitly. ``evidence_id``
    stays content-addressed for the cache/run-DAG handle.
    """
    return EvidenceObject(
        evidence_id=eid,
        source_root="00" * 32,
        document_uri="test://doc",
        title="Test Doc",
        chunk_idx=0,
        chunk_root="11" * 32,
        offset_start=0,
        offset_end=len(span),
        source_role=role,
        text_hash="22" * 32,
        span=span,
        pointer_id=pointer_id,
    )


def test_verify_json_strict_when_all_claims_resolve():
    evidence = [
        _ev("E1f8e4c2a", "Brachiosaurus appears in the Jurassic Park film as a herbivore.", pointer_id="E1"),
        _ev("E2c9d7b3f", "Velociraptor is featured prominently throughout Jurassic Park.", pointer_id="E2"),
    ]
    answer = json.dumps({
        "claims": [
            {"text": "Brachiosaurus appears in the film", "evidence_ids": ["E1"]},
            {"text": "Velociraptor is featured", "evidence_ids": ["E2"]},
        ]
    })
    v = verify_claim_lattice_json(answer, evidence)
    assert v["audit_mode"] == "STRICT"
    # Same verifier_method as the pointer variant — disambiguated downstream
    # via `answer_mode` on the run-DAG and `json_fixups` on the verdict.
    assert v["verifier_method"] == "claim_lattice"
    assert v["n_verified"] == 2
    assert v["violations"] == []


def test_verify_json_hybrid_when_some_unknown_evidence_id():
    evidence = [
        _ev("E1f8e4c2a", "Brachiosaurus appears in the Jurassic Park film as a herbivore.", pointer_id="E1"),
    ]
    answer = json.dumps({
        "claims": [
            {"text": "Brachiosaurus appears in the film", "evidence_ids": ["E1"]},
            {"text": "Made-up claim", "evidence_ids": ["E99"]},
        ]
    })
    v = verify_claim_lattice_json(answer, evidence)
    assert v["audit_mode"] == "HYBRID"
    assert any(vio["kind"] == "UNKNOWN_EVIDENCE_ID" for vio in v["violations"])


def test_verify_json_ungrounded_on_schema_invalid():
    """Lenient parser fails too → SCHEMA_INVALID → UNGROUNDED."""
    v = verify_claim_lattice_json("not json {[", [])
    assert v["audit_mode"] == "UNGROUNDED"
    assert any(vio["kind"] == "SCHEMA_INVALID" for vio in v["violations"])


def test_verify_json_recovers_from_markdown_fence():
    """JSON-fenced output still parses & verifies; fence fixup logged."""
    evidence = [
        _ev("E1f8e4c2a", "Brachiosaurus appears in the Jurassic Park film as a herbivore.", pointer_id="E1"),
    ]
    answer = (
        '```json\n'
        '{"claims": [{"text": "Brachiosaurus appears in the film", '
        '"evidence_ids": ["E1"]}]}\n'
        '```'
    )
    v = verify_claim_lattice_json(answer, evidence)
    assert v["audit_mode"] == "STRICT"
    assert "fence" in v["json_fixups"]


def test_verify_json_manual_quote_violation():
    """Strict no-double-quote rule — even valid JSON with double quotes
    inside a claim's text field fails MANUAL_QUOTE_VIOLATION."""
    evidence = [
        _ev("E1f8e4c2a", "Brachiosaurus appears in the Jurassic Park film as a herbivore.", pointer_id="E1"),
    ]
    answer = json.dumps({
        "claims": [
            {
                "text": 'Brachiosaurus is "a herbivore" appears in the film',
                "evidence_ids": ["E1"],
            }
        ]
    })
    v = verify_claim_lattice_json(answer, evidence)
    assert any(vio["kind"] == "MANUAL_QUOTE_VIOLATION" for vio in v["violations"])
    assert v["audit_mode"] == "UNGROUNDED"


def test_verify_json_blocks_disallowed_source_role():
    """Evidence resolved but source_role outside the allowlist fails
    SOURCE_ROLE_BLOCKED."""
    evidence = [
        _ev("E1f8e4c2a", "Brachiosaurus content here.", role="noisy_background_source", pointer_id="E1"),
    ]
    answer = json.dumps({
        "claims": [
            {"text": "Brachiosaurus appears", "evidence_ids": ["E1"]},
        ]
    })
    v = verify_claim_lattice_json(answer, evidence)
    assert any(vio["kind"] == "SOURCE_ROLE_BLOCKED" for vio in v["violations"])
    assert v["audit_mode"] == "UNGROUNDED"


# ---------------------------------------------------------------- runner integration


def test_runner_ask_json_mode_passes_guided_json_extra_body(tmp_path):
    """`ask()` in JSON mode forwards `extra_body={"guided_json": SCHEMA}`
    through the chat client. Stub captures the kwargs so we can assert."""
    from typing import Iterator

    from aborist.document import Document
    from aborist.ingest import ingest_source
    from aborist.qa import ask
    from aborist.qa.client import StubClient
    from aborist.qa.runner import DEFAULT_POLICY
    from aborist.source import Source
    from aborist.store import connect

    class _FakeSource(Source):
        source_type = "test"
        def __init__(self, docs): self.docs = docs
        def iter_documents(self) -> Iterator[Document]: yield from self.docs

    db = tmp_path / "qa.db"
    conn = connect(db)
    try:
        ingest_source(conn, _FakeSource([
            Document(
                uri="test://doc",
                content="Cloud Strife is the protagonist. " * 30,
                source_type="test",
                title="Cloud",
            )
        ]))
        root = conn.execute(
            "SELECT document_root FROM documents WHERE document_uri='test://doc'"
        ).fetchone()["document_root"]

        client = StubClient(answer='{"claims":[]}')
        policy = dict(DEFAULT_POLICY)
        policy["answer_mode"] = "claim_lattice"

        ask(
            conn,
            document_root=root,
            question="Who is Cloud?",
            client=client,
            model_id="m",
            policy=policy,
        )
        assert len(client.calls) == 1
        kwargs = client.calls[0]["kwargs"]
        assert "extra_body" in kwargs
        assert kwargs["extra_body"] is not None
        assert "guided_json" in kwargs["extra_body"]
    finally:
        conn.close()