218 lines
7.5 KiB
Python
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
|