Files
PlantAI-rag/tests/test_qa_v1_1_73_graph_chain_citations.py

218 lines
7.5 KiB
Python

from __future__ import annotations
from types import SimpleNamespace
from src.rag.qa_v1_1_73.graph_path_chain import GraphHop, NativeGraphPathRetriever
from src.rag.qa_v1_1_73.knowledge_chain import _edge_dict
from src.rag.qa_v1_1_73.pipeline import QAV1Pipeline
from src.rag.qa_v1_1_73.schemas import EvidenceItem, RetrievalHypothesis
def test_graph_hop_preserves_publication_metadata() -> None:
hop = GraphHop(
hop_index=1,
from_id="cent_wrky",
from_name="WRKY",
to_id="cent_senescence",
to_name="leaf senescence",
source="WRKY",
relation="regulates",
target="leaf senescence",
paper_id="PMC8466090_genes-12-01444",
title="WRKY transcription factors in plants",
pmid="34573426",
pmcid="PMC8466090",
doi="10.3390/genes12091444",
species="Arabidopsis thaliana",
evidence="WRKY factors regulate leaf senescence.",
)
value = hop.to_dict()
assert value["title"] == "WRKY transcription factors in plants"
assert value["pmid"] == "34573426"
assert value["pmcid"] == "PMC8466090"
assert value["doi"] == "10.3390/genes12091444"
assert value["species"] == "Arabidopsis thaliana"
def test_local_graph_edge_uses_the_source_citation_contract() -> None:
edge = _edge_dict(
EvidenceItem(
evidence_id="G1.H1",
kb="V16/V15",
source="WRKY",
relation="regulates",
target="leaf senescence",
evidence="WRKY factors regulate leaf senescence.",
paper_id="PMC8466090_genes-12-01444",
title="WRKY transcription factors in plants",
pmid="34573426",
pmcid="PMC8466090",
doi="10.3390/genes12091444",
species="Arabidopsis thaliana",
),
rank=1,
total=1,
)
assert edge["title"] == "WRKY transcription factors in plants"
assert edge["doi"] == "10.3390/genes12091444"
assert edge["species"] == "Arabidopsis thaliana"
assert edge["preferred_citation"] == "[PMID: 34573426]"
context = QAV1Pipeline._format_graph_chain_context(
{"chains": [{"chain_id": "G1", "hop_count": 1, "chain_score": 1.0, "edges": [edge]}]}
)
assert "Paper: WRKY transcription factors in plants" in context
assert "DOI: 10.3390/genes12091444" in context
assert "Species: Arabidopsis thaliana" in context
assert "Preferred citation: [PMID: 34573426]" in context
def test_native_graph_context_uses_real_metadata_without_paper_id_guessing() -> None:
context = QAV1Pipeline._format_native_graph_chain_context(
[
{
"chain_id": "G1",
"hop_count": 2,
"path_score": 0.9,
"hops": [
{
"assertion_id": "a1",
"paper_id": "PMC8466090_genes-12-01444",
"title": "WRKY transcription factors in plants",
"pmid": "34573426",
"pmcid": "PMC8466090",
"doi": "10.3390/genes12091444",
"species": "Arabidopsis thaliana",
"source": "WRKY",
"relation": "regulates",
"target": "leaf senescence",
"evidence": "WRKY factors regulate leaf senescence.",
},
{
"assertion_id": "a2",
"paper_id": "PMC9999999_compound-id",
"source": "leaf senescence",
"relation": "causes",
"target": "leaf yellowing",
"evidence": "Senescence causes leaf yellowing.",
},
],
}
]
)
assert "Title: WRKY transcription factors in plants" in context
assert "Species: Arabidopsis thaliana" in context
assert "Preferred citation: [PMID: 34573426]" in context
second_record = context.split("Record G1.H2", 1)[1]
assert "PMID: -" in second_record
assert "PMCID: -" in second_record
assert "DOI: -" in second_record
assert "Preferred citation: unavailable" in second_record
assert "PMCID: PMC9999999" not in second_record
def test_native_retrieval_maps_assertion_publication_metadata() -> None:
retriever = NativeGraphPathRetriever.__new__(NativeGraphPathRetriever)
retriever._resolve_start_nodes = lambda _ids, _names, limit: [
{"canonical_id": "cent_wrky", "name": "WRKY"}
]
retriever._question_template_paths = lambda _question, max_hops: []
retriever._question_chain_spec = lambda _question: ([], None)
retriever._fetch_neighbors = lambda *_args, **_kwargs: [
{
"from_id": "cent_wrky",
"from_name": "WRKY",
"to_id": "cent_senescence",
"to_name": "leaf senescence",
"source": "WRKY",
"relation": "regulates",
"target": "leaf senescence",
"source_canonical_id": "cent_wrky",
"target_canonical_id": "cent_senescence",
"assertion_id": "a1",
"paper_id": "PMC8466090_genes-12-01444",
"title": "WRKY transcription factors in plants",
"pmid": "34573426",
"pmcid": "PMC8466090",
"doi": "10.3390/genes12091444",
"species": "Arabidopsis thaliana",
"evidence": "WRKY factors regulate leaf senescence.",
"best_score": 1.0,
}
]
retriever._score_edge = lambda *_args: (1.0, {})
retriever._score_path = lambda *_args: (1.0, {})
chains, _audit = retriever.retrieve_chains(
"WRKY如何调控叶片衰老?",
RetrievalHypothesis(),
seed_ids=["cent_wrky"],
seed_names=["WRKY"],
max_hops=1,
top_k_chains=1,
)
hop = chains[0]["hops"][0]
assert hop["title"] == "WRKY transcription factors in plants"
assert hop["pmid"] == "34573426"
assert hop["pmcid"] == "PMC8466090"
assert hop["doi"] == "10.3390/genes12091444"
assert hop["species"] == "Arabidopsis thaliana"
def test_neighbor_query_requests_assertion_publication_metadata() -> None:
statements = []
class FakeSession:
def __enter__(self):
return self
def __exit__(self, *_args):
return False
def run(self, statement, **_parameters):
statements.append(statement)
return []
class FakeDriver:
def session(self):
return FakeSession()
retriever = NativeGraphPathRetriever.__new__(NativeGraphPathRetriever)
retriever.driver = FakeDriver()
assert retriever._fetch_neighbors(
["cent_wrky"],
set(),
relation_clues=[],
priority_relation_clues=[],
entity_clues=[],
answer_type_groups=[],
per_node_expand=10,
) == []
for field in ("title", "pmid", "pmcid", "doi", "species"):
assert f"a.{field} AS {field}" in statements[0]
def test_answer_validator_accepts_exact_graph_chain_preferred_pmid() -> None:
pipeline = QAV1Pipeline.__new__(QAV1Pipeline)
pipeline.answer_composer = SimpleNamespace(
compose=lambda *_args, **_kwargs: "WRKY调控叶片衰老。[PMID: 34573426]"
)
answer, audit = pipeline._compose_validated_answer(
"WRKY如何调控叶片衰老?",
RetrievalHypothesis(),
[],
{},
[],
"Preferred citation: [PMID: 34573426]",
)
assert answer.endswith("[PMID: 34573426]")
assert audit["valid"] is True