PlantAI-rag/tests/test_qa_v1_1_73_direct_flow_metadata.py

from __future__ import annotations

import json
from threading import Barrier, Lock

from src.rag.qa_v1_1_73.direct_flow_retriever import (
    DirectFlowRetriever,
    load_paper_metadata_registry,
)
from src.rag.qa_v1_1_73.schemas import EvidenceItem, RetrievalHypothesis


METADATA = {
    "title": "Genome-Wide Identification of WRKY",
    "doi": "10.3390/genes12091444",
    "pmid": "34573426",
    "pmcid": "PMC8466090",
}


def test_item_from_payload_preserves_citation_metadata() -> None:
    retriever = DirectFlowRetriever.__new__(DirectFlowRetriever)
    item = retriever.item_from_payload(
        {
            "assertion_id": "cass_test",
            "source": "WRKY",
            "relation": "involved_in",
            "target": "stress response",
            "evidence": "WRKY participates in stress response.",
            "paper_id": "PMC8466090_genes-12-01444",
            **METADATA,
        },
        "qdrant_vector",
        "WRKY",
    )

    assert item.title == METADATA["title"]
    assert item.doi == METADATA["doi"]
    assert item.pmid == METADATA["pmid"]
    assert item.pmcid == METADATA["pmcid"]


def test_neo4j_seed_search_requests_and_maps_citation_metadata() -> None:
    class FakeSession:
        def __enter__(self):
            return self

        def __exit__(self, *_args):
            return False

        def run(self, statement, **_parameters):
            for field in ("title", "doi", "pmid", "pmcid"):
                assert f"a.{field} AS {field}" in statement
            return [
                {
                    "assertion_id": "cass_test",
                    "source": "WRKY",
                    "relation": "involved_in",
                    "target": "stress response",
                    "evidence": "evidence",
                    "paper_id": "PMC8466090_genes-12-01444",
                    "score": 0.9,
                    **METADATA,
                }
            ]

    class FakeDriver:
        def session(self):
            return FakeSession()

    retriever = DirectFlowRetriever.__new__(DirectFlowRetriever)
    retriever.driver = FakeDriver()

    items = retriever.neo4j_seed_search(["cent_wrky"], [], 10, method="neo4j_seed_exact")

    assert len(items) == 1
    assert items[0].doi == METADATA["doi"]
    assert items[0].pmcid == METADATA["pmcid"]


def test_load_registry_and_join_by_paper_id(tmp_path) -> None:
    registry = tmp_path / "papers.jsonl"
    registry.write_text(
        json.dumps(
            {"paper_id": "PMC8466090_genes-12-01444", **METADATA},
            ensure_ascii=False,
        )
        + "\n",
        encoding="utf-8",
    )
    retriever = DirectFlowRetriever.__new__(DirectFlowRetriever)
    retriever.paper_metadata = load_paper_metadata_registry(registry)

    item = retriever.item_from_payload(
        {
            "assertion_id": "cass_test",
            "source": "WRKY",
            "relation": "involved_in",
            "target": "stress response",
            "paper_id": "PMC8466090_genes-12-01444",
        },
        "neo4j_seed_exact",
        "WRKY",
    )

    assert item.title == METADATA["title"]
    assert item.doi == METADATA["doi"]
    assert item.pmid == METADATA["pmid"]
    assert item.pmcid == METADATA["pmcid"]


def test_payload_metadata_takes_precedence_over_registry() -> None:
    retriever = DirectFlowRetriever.__new__(DirectFlowRetriever)
    retriever.paper_metadata = {
        "p1": {"paper_id": "p1", "title": "Registry title", "doi": "10.1/old"}
    }

    item = retriever.item_from_payload(
        {"paper_id": "p1", "title": "Payload title", "doi": "10.1/new"},
        "qdrant_vector",
        "query",
    )

    assert item.title == "Payload title"
    assert item.doi == "10.1/new"


def test_legacy_direct_flow_never_creates_or_calls_canonical_driver(monkeypatch) -> None:
    def fail_driver(*_args, **_kwargs):
        raise AssertionError("legacy profile must not create a Neo4j driver")

    monkeypatch.setattr(
        "src.rag.qa_v1_1_73.direct_flow_retriever.GraphDatabase.driver",
        fail_driver,
    )
    retriever = DirectFlowRetriever(graph_schema="legacy", vector_query_workers=1)

    def fail_search(*_args, **_kwargs):
        raise AssertionError("legacy profile must not call canonical Neo4j search")

    retriever.neo4j_seed_search = fail_search
    items, audit = retriever.retrieve(
        "ABRE",
        {},
        RetrievalHypothesis(),
        use_qdrant=False,
        use_neo4j=True,
        use_graph_expansion=True,
    )

    assert retriever.driver is None
    assert items == []
    assert audit["options"]["effective_use_neo4j"] is False
    assert audit["options"]["effective_use_graph_expansion"] is False
    assert audit["neo4j_status"] == "disabled_for_legacy_schema"


def test_vector_queries_run_two_at_a_time_but_merge_in_query_order() -> None:
    retriever = DirectFlowRetriever.__new__(DirectFlowRetriever)
    retriever.graph_schema = "legacy"
    retriever.vector_query_workers = 2
    barrier = Barrier(2, timeout=1.0)
    lock = Lock()
    active = 0
    max_active = 0

    def qdrant_search(query, limit):
        nonlocal active, max_active
        with lock:
            active += 1
            max_active = max(max_active, active)
        try:
            barrier.wait()
            return [
                EvidenceItem(
                    evidence_id="",
                    kb="test",
                    source=query,
                    relation="mentions",
                    target="target",
                    evidence="evidence",
                    paper_id=f"paper-{query}",
                    relation_id=f"relation-{query}",
                    retrieval_queries=[query],
                )
            ]
        finally:
            with lock:
                active -= 1

    retriever.qdrant_search = qdrant_search
    items, audit = retriever.retrieve(
        "Q0",
        {"question_lane": ["Q1"]},
        RetrievalHypothesis(),
        qdrant_top_k=20,
        use_qdrant=True,
        use_neo4j=False,
        use_graph_expansion=False,
    )

    assert max_active == 2
    assert [row["query"] for row in audit["qdrant"]] == ["Q0", "Q1"]
    assert [item.retrieval_queries[0] for item in items] == ["Q0", "Q1"]