arborist/tests/test_dag.py

"""Per-run Merkle-DAG provenance.

`build_run_dag` produces a deterministic Merkle root over the seven
stages of a query/ask call. `verify_run_dag` recomputes the root from
the persisted node list and confirms it matches.

The DAG is not part of cache_key (cache inputs determine the answer;
the answer determines the DAG — folding it back would create a cycle).
It rides alongside the providence record as `run_dag_root` /
`run_dag_blob` so an auditor can verify the run was constructed
exactly as recorded.
"""

from __future__ import annotations

from aborist.qa.dag import build_run_dag, localize_failure, verify_run_dag


def _kw(**overrides):
    base = dict(
        question_hash="a" * 64,
        sources=[
            {
                "document_root": "b" * 64,
                "source_role": "primary_answer_source",
                "score": 1.0,
                "chunk_idx": 0,
            }
        ],
        context_root="b" * 64,
        conversation_hash="c" * 64,
        answer_text="The cat is on the mat.",
        audit_mode="STRICT",
        verifier_method="quote",
        n_quotes=1,
        n_verified=1,
        claim_statuses=[
            {"text": "the cat", "status": "VERIFIED_QUOTE", "method": "quote"}
        ],
        lookup_path="miss",
    )
    base.update(overrides)
    return base


def test_dag_root_is_deterministic():
    """Same inputs → same root, byte-for-byte."""
    a = build_run_dag(**_kw())
    b = build_run_dag(**_kw())
    assert a["root"] == b["root"]
    assert len(a["nodes"]) == 7


def test_dag_root_changes_when_answer_changes():
    a = build_run_dag(**_kw())
    b = build_run_dag(**_kw(answer_text="Different answer."))
    assert a["root"] != b["root"]


def test_dag_root_changes_when_audit_mode_changes():
    a = build_run_dag(**_kw())
    b = build_run_dag(**_kw(audit_mode="UNGROUNDED"))
    assert a["root"] != b["root"]


def test_dag_root_changes_when_verifier_method_changes():
    a = build_run_dag(**_kw())
    b = build_run_dag(**_kw(verifier_method="paraphrase"))
    assert a["root"] != b["root"]


def test_dag_root_changes_when_question_hash_changes():
    a = build_run_dag(**_kw())
    b = build_run_dag(**_kw(question_hash="d" * 64))
    assert a["root"] != b["root"]


def test_dag_root_changes_when_sources_change():
    a = build_run_dag(**_kw())
    new_sources = [{"document_root": "e" * 64, "source_role": "primary_answer_source",
                    "score": 1.0, "chunk_idx": 0}]
    b = build_run_dag(**_kw(sources=new_sources))
    assert a["root"] != b["root"]


def test_dag_seven_stages_in_order():
    """The seven stages are emitted in fixed order so the root is
    canonical: question / retrieval / context / prompt / answer /
    verify / final_label."""
    out = build_run_dag(**_kw())
    stages = [n["stage"] for n in out["nodes"]]
    assert stages == [
        "question", "retrieval", "context", "prompt",
        "answer", "verify", "final_label",
    ]


def test_verify_dag_round_trip():
    """Recomputing the root from the persisted nodes must match."""
    out = build_run_dag(**_kw())
    assert verify_run_dag(out) is True


def test_verify_dag_detects_node_tampering():
    """Mutating a stage hash makes verify fail."""
    out = build_run_dag(**_kw())
    out["nodes"][3]["hash"] = "f" * 64  # tamper with prompt node
    # Root no longer matches the (mutated) leaves.
    assert verify_run_dag(out) is False


def test_localize_failure_strict_returns_none():
    """STRICT verdict has no failure to localize."""
    assert localize_failure(audit_mode="STRICT", n_sources=3, n_quotes=4, n_verified=4) is None


def test_localize_failure_no_sources_is_retrieval():
    """No admitted sources → retrieval-stage failure (gate over-rejected
    or corpus genuinely lacks the topic)."""
    assert localize_failure(
        audit_mode="UNGROUNDED", n_sources=0, n_quotes=0, n_verified=0
    ) == "retrieval"


def test_localize_failure_no_quotes_is_context():
    """Sources retrieved but no quotes extracted → context-stage failure
    (per-source cap dropped relevant content, or model declined to
    cite). Distinct from answer-stage failures where quotes existed but
    didn't verify."""
    assert localize_failure(
        audit_mode="UNGROUNDED", n_sources=5, n_quotes=0, n_verified=0
    ) == "context"


def test_localize_failure_unverified_quotes_is_answer():
    """Quotes extracted but not all verify → answer-stage failure
    (model fabricated, paraphrased inside quotes, or appended
    citations). The case mechanical_repair targets."""
    assert localize_failure(
        audit_mode="HYBRID", n_sources=3, n_quotes=4, n_verified=3
    ) == "answer"
    assert localize_failure(
        audit_mode="UNGROUNDED", n_sources=3, n_quotes=2, n_verified=0
    ) == "answer"


def test_localize_failure_lands_on_run_dag_verify_node():
    """failure_stage is folded into the verify node's payload so
    auditors reading the persisted DAG can see which stage produced
    a non-STRICT verdict."""
    out = build_run_dag(**_kw(
        audit_mode="HYBRID",
        n_quotes=4,
        n_verified=3,
    ))
    # The verify_node hash should differ between failure_stage='answer'
    # and a STRICT verdict — failure_stage enters the hashed payload.
    out_strict = build_run_dag(**_kw(
        audit_mode="STRICT",
        n_quotes=4,
        n_verified=4,
    ))
    out["root"] != out_strict["root"]  # at minimum, audit_mode differs


def test_verify_dag_accepts_json_string():
    """Sidecar/audit tools persist the DAG as JSON; verify accepts both."""
    import json
    out = build_run_dag(**_kw())
    blob = json.dumps(out, separators=(",", ":"))
    assert verify_run_dag(blob) is True