PlantAI-rag/tests/test_qa_v1_1_73_agentic_graph_retriever.py

from src.rag.qa_v1_1_73.agentic_graph_retriever import AgenticGraphRetriever


class FakeDriver:
    def __init__(self, rows):
        self.rows = rows
        self.calls = []

    def execute_query(self, query, **kwargs):
        self.calls.append((query, kwargs))
        return list(self.rows)


def plan(*, incoming=True, outgoing=True, terminal_types=None):
    return {
        "anchors": [
            {
                "mention": "荔枝",
                "canonical_name": "Litchi chinensis",
                "canonical_id": "cent_litchi",
                "entity_type": "species",
            }
        ],
        "answer_targets": [
            {
                "description": "基因家族",
                "allowed_entity_types": terminal_types or ["gene_family"],
            }
        ],
        "relation_plan": {
            "allow_incoming": incoming,
            "allow_outgoing": outgoing,
            "allowed_relation_categories": [],
            "relation_name_hints": [],
            "preserve_assertion_direction": True,
        },
        "graph_plan": {
            "enabled": True,
            "max_hops": 1,
            "anchor_only_start": True,
            "required_terminal_types": terminal_types or ["gene_family"],
        },
    }


def incoming_row():
    return {
        "anchor_id": "cent_litchi",
        "anchor_name": "Litchi chinensis",
        "assertion_id": "cass_spl_litchi",
        "source": "SPL gene family",
        "source_type": "gene_family",
        "source_canonical_id": "cent_spl",
        "relation": "18_homologues_identified",
        "relation_category": "identified",
        "target": "Litchi chinensis",
        "target_type": "species",
        "target_canonical_id": "cent_litchi",
        "evidence": "Eighteen SPL homologues were identified.",
        "paper_id": "10.1080_14620316.2018.1543557",
        "title": "Genome-wide identification of litchi SPL genes",
        "pmid": "",
        "pmcid": "",
        "doi": "10.1080/14620316.2018.1543557",
        "species": "Litchi chinensis",
        "direction": "incoming",
        "support_count": 1,
        "paper_count": 1,
        "best_score": 0.9,
    }


def test_incoming_navigation_preserves_original_assertion_direction():
    retriever = AgenticGraphRetriever(driver=FakeDriver([incoming_row()]))

    evidence, chains, audit = retriever.retrieve(plan())

    assert len(evidence) == 1
    assert evidence[0].source == "SPL gene family"
    assert evidence[0].relation == "18_homologues_identified"
    assert evidence[0].target == "Litchi chinensis"
    assert evidence[0].raw["navigation_direction"] == "incoming"
    assert evidence[0].raw["answer_entity"] == "SPL gene family"
    assert evidence[0].raw["answer_entity_type"] == "gene_family"
    assert evidence[0].raw["assertion_id"] == "cass_spl_litchi"
    assert evidence[0].doi == "10.1080/14620316.2018.1543557"

    assert chains[0]["start"] == "Litchi chinensis"
    assert chains[0]["end"] == "SPL gene family"
    assert chains[0]["navigation_direction"] == "incoming"
    assert chains[0]["hops"][0]["source"] == "SPL gene family"
    assert chains[0]["hops"][0]["target"] == "Litchi chinensis"
    assert audit["returned_assertion_count"] == 1


def test_terminal_type_filter_rejects_non_gene_family_neighbors():
    row = incoming_row()
    row["source"] = "ABA"
    row["source_type"] = "hormone"
    retriever = AgenticGraphRetriever(driver=FakeDriver([row]))

    evidence, chains, audit = retriever.retrieve(plan())

    assert evidence == []
    assert chains == []
    assert audit["filtered_reasons"] == {"terminal_type_mismatch": 1}


def test_direction_plan_filters_rows_even_if_driver_returns_both():
    outgoing = incoming_row()
    outgoing.update(
        {
            "assertion_id": "cass_litchi_lccco",
            "source": "Litchi chinensis",
            "source_type": "species",
            "source_canonical_id": "cent_litchi",
            "relation": "has_gene_family",
            "target": "LcCCO",
            "target_type": "gene_family",
            "target_canonical_id": "cent_lccco",
            "direction": "outgoing",
        }
    )
    retriever = AgenticGraphRetriever(driver=FakeDriver([incoming_row(), outgoing]))

    evidence, chains, audit = retriever.retrieve(plan(incoming=False, outgoing=True))

    assert [item.raw["assertion_id"] for item in evidence] == ["cass_litchi_lccco"]
    assert chains[0]["answer_entity"] == "LcCCO"
    assert audit["filtered_reasons"] == {"direction_disabled": 1}


def test_duplicate_assertions_are_returned_once():
    retriever = AgenticGraphRetriever(
        driver=FakeDriver([incoming_row(), incoming_row()])
    )

    evidence, chains, audit = retriever.retrieve(plan())

    assert len(evidence) == 1
    assert len(chains) == 1
    assert audit["filtered_reasons"] == {"duplicate_assertion": 1}


def test_legacy_schema_uses_existing_entity_relation_anchor_path_and_aliases(
    monkeypatch,
):
    monkeypatch.setenv("QA_V1_GRAPH_SCHEMA", "legacy")
    row = incoming_row()
    row.update(
        {
            "paper_id": "PMC12115174_plants-14-01453",
            "pmcid": "",
            "doi": "",
        }
    )
    driver = FakeDriver([row])
    retriever = AgenticGraphRetriever(driver=driver)

    evidence, chains, audit = retriever.retrieve(plan())

    statement, kwargs = driver.calls[0]
    parameters = kwargs["parameters_"]
    assert "MATCH (anchor:Entity)" in statement
    assert "[r:RELATION]" in statement
    assert "SRT2CanonicalAssertion" not in statement
    assert "r.relation_category" not in statement
    assert set(parameters["anchor_names"]) == {
        "litchi chinensis",
        "l. chinensis",
        "litchi",
        "lychee",
        "荔枝",
    }
    assert parameters["limit"] == 100
    assert audit["query_mode"] == "legacy_anchor_neighborhood"
    assert evidence[0].evidence == "Eighteen SPL homologues were identified."
    assert evidence[0].pmcid == "PMC12115174"
    assert chains[0]["hops"][0]["pmcid"] == "PMC12115174"
    assert chains[0]["navigation_direction"] == "incoming"