PlantAI-rag/tests/test_qa_v1_1_73_q0_soft_fallback.py

from __future__ import annotations

from src.rag.qa_v1_1_73.agentic_controller import AgenticRetrievalController
from src.rag.qa_v1_1_73.pipeline import QAV1Pipeline
from src.rag.qa_v1_1_73.schemas import EvidenceItem


def _controller():
    return AgenticRetrievalController.__new__(AgenticRetrievalController)


def _item(
    index: int,
    *,
    source: str = "ANCHOR",
    relation: str = "regulates",
    target: str | None = None,
    evidence: str | None = None,
    paper_id: str = "paper-1",
    answer_type: str = "unmapped_type",
) -> EvidenceItem:
    resolved_target = target or f"ANSWER{index}"
    return EvidenceItem(
        evidence_id=f"Q0.A1.E{index}",
        kb="V16/V15",
        source=source,
        relation=relation,
        target=resolved_target,
        evidence=evidence or f"ANCHOR {relation} {resolved_target}.",
        paper_id=paper_id,
        relation_id=f"assertion-{index}",
        raw={
            "assertion_id": f"assertion-{index}",
            "answer_entity_type": answer_type,
            "source_type": "gene",
            "target_type": answer_type,
            "anchor_name": "ANCHOR",
        },
    )


def _decomposition():
    return {
        "subquestions": [
            {
                "subquestion_id": "Q0",
                "question": "Which entities are associated with ANCHOR?",
                "role": "required",
                "answer_slot": "entities",
                "anchors": [{"mention": "ANCHOR", "canonical_name": "ANCHOR"}],
                "target_entity_types": ["gene"],
                "acceptance_requirements": {
                    "requires_structured_assertion": True,
                    "requires_paper_provenance": True,
                    "requires_anchor_connection": True,
                    "requires_terminal_type_match": True,
                },
            },
            {
                "subquestion_id": "Q1",
                "question": "What background is known?",
                "role": "supporting",
            },
        ]
    }


def _groups(q0_evidence=None, q0_chains=None, supporting=None):
    q0_evidence = list(q0_evidence or [])
    q0_chains = list(q0_chains or [])
    supporting = list(supporting or [])
    return [
        {
            "subquestion_id": "Q0",
            "question": "Which entities are associated with ANCHOR?",
            "role": "required",
            "status": "COMPLETE" if q0_evidence or q0_chains else "EMPTY",
            "accepted_evidence_ids": [item.evidence_id for item in q0_evidence],
            "accepted_chain_ids": [chain["chain_id"] for chain in q0_chains],
            "evidence": q0_evidence,
            "chains": q0_chains,
        },
        {
            "subquestion_id": "Q1",
            "question": "What background is known?",
            "role": "supporting",
            "status": "COMPLETE" if supporting else "EMPTY",
            "accepted_evidence_ids": [item.evidence_id for item in supporting],
            "accepted_chain_ids": [],
            "evidence": supporting,
            "chains": [],
        },
    ]


def test_soft_fallback_keeps_agent_items_first_and_copies_safe_q0(monkeypatch):
    monkeypatch.setenv("QA_AGENTIC_GATE_POLICY", "soft_compat_v1")
    accepted = _item(1, answer_type="gene")
    fallback = _item(2, answer_type="not_in_planner_vocabulary")
    groups = _groups(q0_evidence=[accepted])

    result = QAV1Pipeline._apply_q0_soft_fallback(
        _decomposition(), groups, [accepted, fallback], _controller(), limit=30
    )
    q0 = result["groups"][0]

    assert [item.evidence_id for item in q0["evidence"]] == [
        accepted.evidence_id,
        fallback.evidence_id,
    ]
    assert q0["evidence"][1] is not fallback
    assert fallback.raw.get("handoff_origin") is None
    assert q0["evidence"][1].raw["handoff_origin"] == "q0_soft_fallback"
    assert result["strict_review_accepted"] is True
    assert result["required_q0_complete"] is True
    assert result["fallback_count"] == 1
    assert result["fallback_original_ranks"] == [2]


def test_soft_fallback_preserves_non_type_safety_checks(monkeypatch):
    monkeypatch.setenv("QA_AGENTIC_GATE_POLICY", "soft_compat_v1")
    no_provenance = _item(2, paper_id="")
    wrong_anchor = _item(
        3,
        source="OTHER",
        evidence="OTHER regulates ANSWER3.",
    )
    cold_for_heat = _item(
        4,
        evidence="ANCHOR responds only to cold stress and freezing.",
    )
    decomposition = _decomposition()
    decomposition["subquestions"][0]["question"] = (
        "Which entities are associated with ANCHOR under heat stress?"
    )

    result = QAV1Pipeline._apply_q0_soft_fallback(
        decomposition,
        _groups(),
        [no_provenance, wrong_anchor, cold_for_heat],
        _controller(),
        limit=30,
    )

    assert result["groups"][0]["evidence"] == []
    assert result["strict_review_accepted"] is False
    assert result["required_q0_complete"] is False


def test_supporting_only_never_counts_as_required_q0_complete(monkeypatch):
    monkeypatch.setenv("QA_AGENTIC_GATE_POLICY", "soft_compat_v1")
    supporting = _item(9, source="BACKGROUND", answer_type="gene")

    result = QAV1Pipeline._apply_q0_soft_fallback(
        _decomposition(), _groups(supporting=[supporting]), [], _controller(), limit=30
    )

    assert result["groups"][1]["evidence"] == [supporting]
    assert result["strict_review_accepted"] is False
    assert result["required_q0_complete"] is False


def test_q0_fallback_deduplicates_and_caps_at_thirty(monkeypatch):
    monkeypatch.setenv("QA_AGENTIC_GATE_POLICY", "soft_compat_v1")
    accepted = [_item(1, answer_type="gene"), _item(2, answer_type="gene")]
    candidates = accepted + [_item(index) for index in range(2, 40)]

    result = QAV1Pipeline._apply_q0_soft_fallback(
        _decomposition(), _groups(q0_evidence=accepted), candidates, _controller(), limit=30
    )
    q0 = result["groups"][0]

    assert len(q0["evidence"]) == 30
    assert len({item.relation_id for item in q0["evidence"]}) == 30
    assert q0["evidence"][:2] == accepted


def test_legacy_policy_is_a_noop(monkeypatch):
    monkeypatch.setenv("QA_AGENTIC_GATE_POLICY", "legacy")
    fallback = _item(2)
    groups = _groups()

    result = QAV1Pipeline._apply_q0_soft_fallback(
        _decomposition(), groups, [fallback], _controller(), limit=30
    )

    assert result["groups"] == groups
    assert result["fallback_count"] == 0
    assert result["strict_review_accepted"] is False


def test_global_handoff_orders_all_agent_items_before_q0_fallback():
    q0_agent = _item(1, answer_type="gene")
    q0_fallback = _item(2)
    q0_fallback.raw["handoff_origin"] = "q0_soft_fallback"
    supporting_agent = _item(9, source="BACKGROUND", answer_type="gene")
    groups = _groups(q0_evidence=[q0_agent, q0_fallback], supporting=[supporting_agent])

    selected = QAV1Pipeline._ordered_agentic_handoff_evidence(groups, limit=3)

    assert [item.evidence_id for item in selected] == [
        q0_agent.evidence_id,
        supporting_agent.evidence_id,
        q0_fallback.evidence_id,
    ]