PlantAI-rag/tests/test_qa_v1_1_73_agentic_pipeline.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 _item(assertion_id: str, source: str, target: str, answer_type: str) -> EvidenceItem:
    return EvidenceItem(
        evidence_id="",
        kb="test",
        source=source,
        relation="identified_in",
        target=target,
        evidence=f"{source} was identified in {target}.",
        paper_id=f"paper-{assertion_id}",
        relation_id=assertion_id,
        raw={
            "assertion_id": assertion_id,
            "anchor_name": "Litchi chinensis" if "Litchi" in (source + target) else "",
            "answer_entity_type": answer_type,
        },
    )


def _chain(chain_id: str, assertion_id: str, source: str, target: str) -> dict:
    return {
        "chain_id": chain_id,
        "anchor_name": "Litchi chinensis",
        "answer_entity_type": "gene_family",
        "start": "Litchi chinensis",
        "end": source,
        "hops": [
            {
                "assertion_id": assertion_id,
                "source": source,
                "relation": "identified_in",
                "target": target,
                "paper_id": f"paper-{assertion_id}",
            }
        ],
    }


class _ReplanningController:
    max_replans = 2

    def __init__(self):
        self.review_calls = []

    def review(self, question, plan, evidence, chains, *, attempt):
        self.review_calls.append((attempt, plan, evidence, chains))
        if attempt == 0:
            return {
                "decision": "REPLAN",
                "accepted_evidence_ids": [],
                "accepted_chain_ids": [],
                "replan": {"required": True, "reason": "enable incoming traversal"},
                "final_handoff": {"allowed": False},
            }
        litchi_evidence = [item.evidence_id for item in evidence if item.relation_id == "a-litchi"]
        litchi_chains = [chain["chain_id"] for chain in chains if chain["chain_id"].startswith("AG")]
        return {
            "decision": "ACCEPT",
            "accepted_evidence_ids": litchi_evidence,
            "accepted_chain_ids": litchi_chains,
            "replan": {"required": False},
            "final_handoff": {
                "allowed": True,
                "evidence_ids": litchi_evidence,
                "chain_ids": litchi_chains,
                "answer_constraints": ["Only report Litchi-connected gene families"],
            },
        }

    @staticmethod
    def apply_replan(plan, review):
        return {**plan, "round": 1, "relation_plan": {"allow_incoming": True}}


class _RoundGraphRetriever:
    calls = []

    def __init__(self):
        self.calls = _RoundGraphRetriever.calls

    def retrieve(self, plan):
        self.calls.append(plan)
        if plan.get("round") == 1:
            evidence = _item("a-litchi", "bZIP gene family", "Litchi chinensis", "gene_family")
            return [evidence], [_chain("AG1", "a-litchi", evidence.source, evidence.target)], {"status": "ok"}
        return [], [], {"status": "no_matching_assertions"}

    def close(self):
        return None


def test_agentic_loop_replans_and_hands_only_accepted_anchor_evidence_to_answer_llm():
    _RoundGraphRetriever.calls = []
    pipeline = QAV1Pipeline.__new__(QAV1Pipeline)
    pipeline.agentic_controller = _ReplanningController()
    pipeline.agentic_graph_retriever_factory = _RoundGraphRetriever
    unrelated = _item("a-tcp", "TCP gene family", "Arabidopsis thaliana", "gene_family")
    plan = {
        "anchors": [{"canonical_name": "Litchi chinensis"}],
        "graph_plan": {"enabled": True, "required_terminal_types": ["gene_family"]},
    }

    evidence, chains, audit, review = pipeline._run_agentic_review_loop(
        "荔枝中有哪些基因家族已经被研究", plan, [unrelated], []
    )

    assert [item.relation_id for item in evidence] == ["a-litchi"]
    assert [chain["chain_id"] for chain in chains] == ["AG1"]
    assert review["decision"] == "ACCEPT"
    assert audit["attempt_count"] == 2
    assert audit["replan_count"] == 1
    assert len(_RoundGraphRetriever.calls) == 2


class _AlwaysRejectController:
    max_replans = 2

    def review(self, question, plan, evidence, chains, *, attempt):
        if attempt >= self.max_replans:
            return {
                "decision": "ABSTAIN",
                "accepted_evidence_ids": [],
                "accepted_chain_ids": [],
                "replan": {"required": False},
                "final_handoff": {
                    "allowed": True,
                    "answer_constraints": ["State that evidence is insufficient"],
                },
            }
        return {
            "decision": "REPLAN",
            "accepted_evidence_ids": [],
            "accepted_chain_ids": [],
            "replan": {"required": True},
            "final_handoff": {"allowed": False},
        }

    @staticmethod
    def apply_replan(plan, review):
        return {
            **plan,
            "retry": int(plan.get("retry", 0)) + 1,
            "relation_plan": {"allow_incoming": True, "allow_outgoing": False},
        }


def test_agentic_loop_stops_after_one_replan_and_returns_no_unreviewed_anchor_evidence():
    _RoundGraphRetriever.calls = []
    pipeline = QAV1Pipeline.__new__(QAV1Pipeline)
    pipeline.agentic_controller = _AlwaysRejectController()
    pipeline.agentic_graph_retriever_factory = _RoundGraphRetriever
    unrelated = _item("a-tcp", "TCP gene family", "Arabidopsis thaliana", "gene_family")

    evidence, chains, audit, review = pipeline._run_agentic_review_loop(
        "荔枝中有哪些基因家族已经被研究",
        {
            "anchors": [{"canonical_name": "Litchi chinensis"}],
            "graph_plan": {
                "enabled": True,
                "required_terminal_types": ["gene_family"],
            },
        },
        [unrelated],
        [],
    )

    assert evidence == []
    assert chains == []
    assert review["decision"] == "ABSTAIN"
    assert audit["attempt_count"] == 2
    assert audit["replan_count"] == 1


class _OneResponseLLM:
    available = True

    def complete_json(self, *_args, **_kwargs):
        return {
            "decision": "REPLAN",
            "accepted_evidence_ids": [],
            "accepted_chain_ids": [],
            "replan": {"required": True, "reason": "retry graph"},
            "final_handoff": {"allowed": False},
        }


def test_agentic_loop_stops_noop_replan_and_falls_back_to_structured_vector_evidence():
    _RoundGraphRetriever.calls = []
    pipeline = QAV1Pipeline.__new__(QAV1Pipeline)
    pipeline.agentic_controller = AgenticRetrievalController(_OneResponseLLM(), max_replans=2)
    pipeline.agentic_graph_retriever_factory = _RoundGraphRetriever
    direct = EvidenceItem(
        evidence_id="",
        kb="V16/V15",
        source="R2R3-MYB proteins",
        relation="contain",
        target="R2 and R3 MYB repeats",
        evidence="R2R3-MYB proteins contain the R2 and R3 MYB repeats.",
        doi="10.1000/r2r3",
        relation_id="r-r2r3",
        raw={"answer_entity_type": "domain"},
    )
    plan = {
        "anchors": [],
        "question_analysis": {
            "operation": "lookup",
            "expected_answer_shape": "fact",
        },
        "graph_plan": {
            "enabled": True,
            "required_terminal_types": ["unknown"],
        },
        "vector_plan": {"queries": ["R2R3-MYB domain"]},
    }

    evidence, chains, audit, review = pipeline._run_agentic_review_loop(
        "R2R3-MYB蛋白为什么被称为R2R3-MYB？", plan, [direct], []
    )

    assert [item.relation_id for item in evidence] == ["r-r2r3"]
    assert chains == []
    assert audit["status"] == "fallback_to_verified_base_evidence"
    assert audit["attempt_count"] == 1
    assert audit["replan_count"] == 0
    assert audit["stop_reason"] == "no_effective_plan_change"
    assert len(_RoundGraphRetriever.calls) == 1
    assert review["final_handoff"]["allowed"] is True
    assert review["final_handoff"]["evidence_ids"] == [evidence[0].evidence_id]