PlantAI-rag/tests/test_qa_v1_1_73_wheat_spl_stability.py

from pathlib import Path

from src.rag.qa_v1_1_73.domain_config import DomainConfig
from src.rag.qa_v1_1_73.pipeline import QAV1Pipeline
from src.rag.qa_v1_1_73.query_builder import QueryBuilder
from src.rag.qa_v1_1_73.schemas import EvidenceItem, RetrievalHypothesis


QUESTION = "小麦参与开花中哪些SPL基因发挥主要作用"


def domain() -> DomainConfig:
    return DomainConfig(
        Path("configs/rag/entity_alias.yaml"),
        Path("configs/rag/relation_taxonomy.yaml"),
    )


def evidence(
    index: int,
    *,
    pmid: str = "",
    source: str = "candidate",
    target: str = "unrelated",
) -> EvidenceItem:
    return EvidenceItem(
        evidence_id="",
        kb="V16/V15" if pmid else "verify_fix3",
        source=source,
        relation="regulates",
        target=target,
        evidence=f"fixture evidence {index}",
        pmid=pmid,
        vector_score=1.0 - index * 0.001,
    )


def test_wheat_spl_flowering_question_lane_is_deterministic_when_llm_misses_entities() -> None:
    pipeline = QAV1Pipeline.__new__(QAV1Pipeline)
    pipeline.domain = domain()
    hypothesis = RetrievalHypothesis()

    pipeline._complete_direct_question_entities(QUESTION, hypothesis)
    queries = QueryBuilder(pipeline.domain).build(QUESTION, "", hypothesis)

    assert queries["question_lane"] == [
        QUESTION,
        "小麦 SPL 开花",
        "Triticum aestivum SPL gene family flowering",
        "wheat SPL genes ovary anthesis expression",
    ]


def test_relation_expansion_does_not_contaminate_evidence_facet_query() -> None:
    pipeline = QAV1Pipeline.__new__(QAV1Pipeline)
    pipeline.domain = domain()
    hypothesis = RetrievalHypothesis(relation_intents=["function_effect"])

    pipeline._complete_direct_question_entities(QUESTION, hypothesis)
    queries = QueryBuilder(pipeline.domain).build(QUESTION, "", hypothesis)

    assert "wheat SPL genes ovary anthesis expression" in queries["question_lane"]
    assert all(
        term not in "wheat SPL genes ovary anthesis expression"
        for term in ("promotes", "inhibits", "regulates")
    )


def test_question_lane_preserves_canonical_query_when_variant_limit_is_tight() -> None:
    pipeline = QAV1Pipeline.__new__(QAV1Pipeline)
    pipeline.domain = domain()
    hypothesis = RetrievalHypothesis()

    pipeline._complete_direct_question_entities(QUESTION, hypothesis)

    assert QueryBuilder(pipeline.domain, max_variants=2).build(
        QUESTION, "", hypothesis
    )["question_lane"] == [
        QUESTION,
        "Triticum aestivum SPL gene family flowering",
    ]
    assert QueryBuilder(pipeline.domain, max_variants=1).build(
        QUESTION, "", hypothesis
    )["question_lane"] == [QUESTION]


def test_ascii_aliases_do_not_match_inside_larger_tokens() -> None:
    entities = domain().direct_entities_from_question("buckwheat deflowering SPLASH")

    assert entities == []


def test_ascii_aliases_still_match_complete_scientific_tokens() -> None:
    entities = domain().direct_entities_from_question("wheat flowering SPL genes")

    assert [entity.canonical_candidate for entity in entities] == [
        "Triticum aestivum",
        "SPL gene family",
        "flowering",
    ]


def test_litchi_chinese_alias_uses_existing_species_graph_seed_pipeline() -> None:
    entities = domain().direct_entities_from_question(
        "荔枝中有哪些基因家族已经被研究"
    )

    assert len(entities) == 1
    assert entities[0].canonical_candidate == "Litchi chinensis"
    assert entities[0].entity_type == "species"
    assert entities[0].verification_status == "verified_alias"
    assert entities[0].use_as_graph_seed is True


def test_flowering_evidence_terms_are_configuration_driven() -> None:
    config = domain()
    entities = config.direct_entities_from_question("wheat SPL genes flowering")

    assert config.evidence_query_terms(entities) == [
        "wheat",
        "SPL genes",
        "ovary",
        "anthesis",
        "expression",
    ]


def test_unconfigured_entities_do_not_create_evidence_terms() -> None:
    config = domain()
    entities = config.direct_entities_from_question("Arabidopsis GST genes")

    assert config.evidence_query_terms(entities) == []


def test_unrelated_question_does_not_gain_flowering_evidence_facet() -> None:
    pipeline = QAV1Pipeline.__new__(QAV1Pipeline)
    pipeline.domain = domain()
    hypothesis = RetrievalHypothesis()

    pipeline._complete_direct_question_entities(
        "拟南芥GST基因家族有哪些功能", hypothesis
    )
    queries = QueryBuilder(pipeline.domain).build(
        "拟南芥GST基因家族有哪些功能", "", hypothesis
    )

    assert "wheat SPL genes ovary anthesis expression" not in queries["question_lane"]


def test_chat_filter_prevents_high_score_uncitable_candidates_displacing_taspl038_evidence() -> None:
    uncitable = [evidence(index) for index in range(30)]
    taspl038 = evidence(
        100,
        pmid="32912142",
        source="TaSPL038",
        target="ovary_anthesis",
    )
    raw_items = uncitable + [taspl038]

    unfiltered_top30 = QAV1Pipeline._pre_rerank(raw_items, 30)
    assert taspl038 not in unfiltered_top30

    chat_items, chat_audit = QAV1Pipeline._filter_public_answer_candidates(
        raw_items, retrieval_only=False
    )
    chat_top30 = QAV1Pipeline._pre_rerank(chat_items, 30)
    assert taspl038 in chat_top30
    assert chat_audit == {
        "applied": True,
        "input_count": 31,
        "kept_count": 1,
        "filtered_uncitable_count": 30,
    }

    search_items, search_audit = QAV1Pipeline._filter_public_answer_candidates(
        raw_items, retrieval_only=True
    )
    assert search_items == raw_items
    assert search_audit == {
        "applied": False,
        "input_count": 31,
        "kept_count": 31,
        "filtered_uncitable_count": 0,
    }