269 lines
8.3 KiB
Python
269 lines
8.3 KiB
Python
from __future__ import annotations
|
|
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
from src.rag.qa_v1_1_73.pipeline import QAAnswerError, QAV1Pipeline
|
|
from src.rag.qa_v1_1_73.schemas import EvidenceItem, RetrievalHypothesis
|
|
|
|
|
|
def item(
|
|
evidence_id: str,
|
|
*,
|
|
kb: str = "V16/V15",
|
|
pmid: str = "",
|
|
pmcid: str = "",
|
|
doi: str = "",
|
|
relation: str = "regulates",
|
|
) -> EvidenceItem:
|
|
return EvidenceItem(
|
|
evidence_id=evidence_id,
|
|
kb=kb,
|
|
source="WRKY",
|
|
relation=relation,
|
|
target="leaf senescence",
|
|
evidence="WRKY is associated with leaf senescence.",
|
|
paper_id="paper-" + evidence_id,
|
|
title="WRKY study " + evidence_id,
|
|
pmid=pmid,
|
|
pmcid=pmcid,
|
|
doi=doi,
|
|
)
|
|
|
|
|
|
class SequenceComposer:
|
|
def __init__(self, answers):
|
|
self.answers = iter(answers)
|
|
self.calls = []
|
|
|
|
def compose(self, *args):
|
|
self.calls.append(args)
|
|
return next(self.answers)
|
|
|
|
|
|
def test_chat_composer_excludes_uncitable_candidate_kbs_but_keeps_citable_evidence() -> None:
|
|
composer = SequenceComposer(["WRKY与叶片衰老相关。[PMID: 34573426]"])
|
|
pipeline = QAV1Pipeline.__new__(QAV1Pipeline)
|
|
pipeline.answer_composer = composer
|
|
evidence = [
|
|
item("S1", pmid="34573426"),
|
|
item("S2", kb="verify_fix3"),
|
|
item("S3", kb="pecan_llm_v2"),
|
|
]
|
|
|
|
pipeline._compose_validated_answer(
|
|
"WRKY有什么功能？", RetrievalHypothesis(), evidence, {}, [], ""
|
|
)
|
|
|
|
assert composer.calls[0][2] == [evidence[0]]
|
|
|
|
|
|
def test_internal_labels_trigger_one_llm_rewrite_with_identical_evidence() -> None:
|
|
composer = SequenceComposer(
|
|
[
|
|
"WRKY参与衰老。[S1] [Data: Sources (S1)] G1.H1",
|
|
"WRKY与叶片衰老相关。[PMCID: PMC8466090]",
|
|
]
|
|
)
|
|
pipeline = QAV1Pipeline.__new__(QAV1Pipeline)
|
|
pipeline.answer_composer = composer
|
|
evidence = [item("S1", pmcid="PMC8466090")]
|
|
|
|
answer, audit = pipeline._compose_validated_answer(
|
|
"WRKY有什么功能？", RetrievalHypothesis(), evidence, {}, [], ""
|
|
)
|
|
|
|
assert answer.endswith("[PMCID: PMC8466090]")
|
|
assert len(composer.calls) == 2
|
|
assert composer.calls[0][2] == composer.calls[1][2] == evidence
|
|
assert composer.calls[0][5] == composer.calls[1][5] == ""
|
|
assert composer.calls[0][6] == ""
|
|
assert "same evidence" in composer.calls[1][6].lower()
|
|
assert "internal" in composer.calls[1][6].lower()
|
|
assert "malformed" in composer.calls[1][6].lower()
|
|
assert "unknown" in composer.calls[1][6].lower()
|
|
assert audit["rewrite_attempted"] is True
|
|
|
|
|
|
def test_full_width_public_citation_brackets_trigger_one_rewrite() -> None:
|
|
composer = SequenceComposer(
|
|
[
|
|
"WRKY与叶片衰老相关。【PMID: 34573426】",
|
|
"WRKY与叶片衰老相关。[PMID: 34573426]",
|
|
]
|
|
)
|
|
pipeline = QAV1Pipeline.__new__(QAV1Pipeline)
|
|
pipeline.answer_composer = composer
|
|
evidence = [item("S1", pmid="34573426")]
|
|
|
|
answer, audit = pipeline._compose_validated_answer(
|
|
"WRKY有什么功能？", RetrievalHypothesis(), evidence, {}, [], ""
|
|
)
|
|
|
|
assert answer.endswith("[PMID: 34573426]")
|
|
assert len(composer.calls) == 2
|
|
assert composer.calls[0][2] == composer.calls[1][2] == evidence
|
|
assert audit["cited_pmids"] == ["34573426"]
|
|
assert audit["rewrite_attempted"] is True
|
|
|
|
|
|
def test_second_internal_label_violation_fails_closed_without_cleaning() -> None:
|
|
composer = SequenceComposer(["结论。[S1]", "结论。[G1.H1]"])
|
|
pipeline = QAV1Pipeline.__new__(QAV1Pipeline)
|
|
pipeline.answer_composer = composer
|
|
|
|
with pytest.raises(QAAnswerError) as exc:
|
|
pipeline._compose_validated_answer(
|
|
"WRKY有什么功能？",
|
|
RetrievalHypothesis(),
|
|
[item("S1", doi="10.3390/genes12091444")],
|
|
{},
|
|
[],
|
|
"",
|
|
)
|
|
|
|
assert len(composer.calls) == 2
|
|
assert exc.value.error_code == "invalid_llm_citation"
|
|
assert exc.value.answer == ""
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"answer",
|
|
[
|
|
"结论。[PMID: 99999999]",
|
|
"结论。[PMCID: PMC9999999]",
|
|
"结论。[DOI: 10.9999/not-allowed]",
|
|
"结论。[DOI: not-a-doi]",
|
|
],
|
|
)
|
|
def test_all_visible_public_identifiers_must_belong_to_allowed_evidence(answer: str) -> None:
|
|
composer = SequenceComposer([answer, answer])
|
|
pipeline = QAV1Pipeline.__new__(QAV1Pipeline)
|
|
pipeline.answer_composer = composer
|
|
|
|
with pytest.raises(QAAnswerError) as exc:
|
|
pipeline._compose_validated_answer(
|
|
"WRKY有什么功能？",
|
|
RetrievalHypothesis(),
|
|
[item("S1", pmcid="PMC8466090", doi="10.3390/genes12091444")],
|
|
{},
|
|
[],
|
|
"Preferred citation: [PMID: 34573426]",
|
|
)
|
|
|
|
assert len(composer.calls) == 2
|
|
assert exc.value.answer == ""
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"answer",
|
|
[
|
|
"结论。[PMID: 34573426",
|
|
"结论。[PMID 34573426]",
|
|
"结论。[PMCID: PMC8466090",
|
|
"结论。[PMCID PMC8466090]",
|
|
"结论。[DOI: 10.3390/genes12091444",
|
|
"结论。[DOI 10.3390/genes12091444]",
|
|
],
|
|
)
|
|
def test_malformed_citation_like_labels_fail_closed(answer: str) -> None:
|
|
composer = SequenceComposer([answer, answer])
|
|
pipeline = QAV1Pipeline.__new__(QAV1Pipeline)
|
|
pipeline.answer_composer = composer
|
|
|
|
with pytest.raises(QAAnswerError):
|
|
pipeline._compose_validated_answer(
|
|
"WRKY有什么功能？",
|
|
RetrievalHypothesis(),
|
|
[
|
|
item(
|
|
"S1",
|
|
pmid="34573426",
|
|
pmcid="PMC8466090",
|
|
doi="10.3390/genes12091444",
|
|
)
|
|
],
|
|
{},
|
|
[],
|
|
"",
|
|
)
|
|
|
|
assert len(composer.calls) == 2
|
|
|
|
|
|
def test_plain_identifier_words_without_citation_shape_are_not_rejected() -> None:
|
|
answer = "PMID、PMCID和DOI是文献标识符。"
|
|
composer = SequenceComposer([answer])
|
|
pipeline = QAV1Pipeline.__new__(QAV1Pipeline)
|
|
pipeline.answer_composer = composer
|
|
|
|
result, audit = pipeline._compose_validated_answer(
|
|
"这些缩写是什么？", RetrievalHypothesis(), [], {}, [], ""
|
|
)
|
|
|
|
assert result == answer
|
|
assert audit["valid"] is True
|
|
|
|
|
|
def test_chat_sources_are_cited_publication_cards_deduped_across_srt_and_graph_hops() -> None:
|
|
evidence = [
|
|
item("S1", pmid="34573426"),
|
|
item("S2", pmid="34573426", relation="associated_with"),
|
|
item("S3", pmid="11111111"),
|
|
item("S4", kb="verify_fix3"),
|
|
]
|
|
graph_chains = [
|
|
{
|
|
"chain_id": "G1",
|
|
"hops": [
|
|
{
|
|
"assertion_id": "a1",
|
|
"paper_id": "paper-pmc",
|
|
"title": "Graph PMCID study",
|
|
"pmcid": "PMC9000001",
|
|
"species": "Arabidopsis thaliana",
|
|
"source": "A",
|
|
"relation": "regulates",
|
|
"target": "B",
|
|
},
|
|
{
|
|
"assertion_id": "a2",
|
|
"paper_id": "paper-doi",
|
|
"title": "Graph DOI study",
|
|
"doi": "10.1000/graph-doi",
|
|
"source": "B",
|
|
"relation": "regulates",
|
|
"target": "C",
|
|
},
|
|
],
|
|
}
|
|
]
|
|
answer = (
|
|
"结论。[PMID: 34573426] [PMCID: PMC9000001] "
|
|
"[DOI: 10.1000/graph-doi]"
|
|
)
|
|
|
|
sources = QAV1Pipeline._response_sources(
|
|
answer, evidence, graph_chains, retrieval_only=False, include_raw=False
|
|
)
|
|
|
|
assert [source["preferred_citation"] for source in sources] == [
|
|
"[PMID: 34573426]",
|
|
"[PMCID: PMC9000001]",
|
|
"[DOI: 10.1000/graph-doi]",
|
|
]
|
|
assert all(source.get("evidence_id") != "S3" for source in sources)
|
|
assert sources[1]["species"] == "Arabidopsis thaliana"
|
|
assert len(sources) == 3
|
|
|
|
|
|
def test_retrieval_only_sources_retain_uncitable_candidates_for_audit() -> None:
|
|
evidence = [item("S1", pmid="34573426"), item("S2", kb="verify_fix3")]
|
|
|
|
sources = QAV1Pipeline._response_sources(
|
|
"", evidence, [], retrieval_only=True, include_raw=False
|
|
)
|
|
|
|
assert [source["evidence_id"] for source in sources] == ["S1", "S2"]
|