PlantAI-rag/tests/test_qa_v1_1_73_answer_composer.py

from src.rag.qa_v1_1_73.answer_composer import AnswerComposer
from src.rag.qa_v1_1_73.schemas import EvidenceItem, RetrievalHypothesis


class CapturingLLM:
    def __init__(self) -> None:
        self.user_prompt = ""

    def complete(self, system: str, user: str, **kwargs) -> str:
        self.user_prompt = user
        return "ok"


def test_answer_composer_includes_optional_correction_note_only_on_retry() -> None:
    llm = CapturingLLM()
    composer = AnswerComposer(llm)

    composer.compose("问题", RetrievalHypothesis(), [], {}, [], "", "")
    assert "Citation correction required" not in llm.user_prompt

    composer.compose(
        "问题",
        RetrievalHypothesis(),
        [],
        {},
        [],
        "",
        "Use the same evidence and remove internal labels.",
    )
    assert "Citation correction required" in llm.user_prompt
    assert "Use the same evidence and remove internal labels." in llm.user_prompt


def test_answer_prompt_prioritizes_complete_path_terminal_not_known_bridge() -> None:
    llm = CapturingLLM()
    composer = AnswerComposer(llm)
    composer.compose(
        "A reaches B, and B reaches C. What does C contain?",
        RetrievalHypothesis(),
        [],
        {},
        [],
        "G1: A -> B -> C -> D",
    )
    prompt = llm.user_prompt.lower()
    assert "continuous" in prompt
    assert "terminal" in prompt
    assert "known bridge" in prompt
    assert "do not stop" in prompt


def test_answer_prompt_requests_evidence_rich_process_explanations() -> None:
    llm = CapturingLLM()
    composer = AnswerComposer(llm)
    composer.compose(
        "WRKY基因家族的功能有哪些？",
        RetrievalHypothesis(question_type="function"),
        [],
        {},
        [],
        "G1.H1: WRKY family -> seed dormancy\nG1.H2: seed dormancy -> germination",
    )
    prompt = (composer.SYSTEM + "\n" + llm.user_prompt).lower()

    assert "overview" in prompt
    assert "intermediate biological process" in prompt
    assert "outcome" in prompt
    assert "all supplied evidence" in prompt
    assert "wholly unrelated" in prompt


def test_answer_prompt_does_not_require_specific_genes_or_overclaim_graph_chains() -> None:
    llm = CapturingLLM()
    composer = AnswerComposer(llm)
    composer.compose(
        "WRKY基因家族的功能有哪些？",
        RetrievalHypothesis(question_type="function"),
        [],
        {},
        [],
        "G1.H1: WRKY family -> leaf senescence",
    )
    prompt = (composer.SYSTEM + "\n" + llm.user_prompt).lower()

    assert "not required" in prompt
    assert "specific gene" in prompt
    assert "process-level" in prompt
    assert "direct molecular mechanism" in prompt


def test_answer_prompt_uses_public_citations_and_conservative_species_context() -> None:
    llm = CapturingLLM()
    composer = AnswerComposer(llm)
    item = EvidenceItem(
        evidence_id="S1",
        kb="V16/V15",
        source="AtRALF1",
        relation="inhibits",
        target="root growth",
        evidence="In Arabidopsis thaliana, AtRALF1 inhibits root growth.",
        title="RALF signalling in Arabidopsis thaliana",
        pmid="34573426",
        species="Arabidopsis thaliana",
    )

    composer.compose(
        "RALF基因家族有哪些功能？",
        RetrievalHypothesis(question_type="function"),
        [item],
        {},
        [],
    )
    prompt = composer.SYSTEM + "\n" + llm.user_prompt

    assert "[PMID: 34573426]" in prompt
    assert "[Data: Sources" not in composer.SYSTEM
    assert "Do not write paper titles in the answer" in prompt
    assert "species + gene or family member + supported function or process" in prompt
    assert "Do not promote paper species to claim species" in prompt
    assert "Answer in the same language as the user's question" in prompt


def test_answer_prompt_marks_uncitable_source_as_unavailable() -> None:
    llm = CapturingLLM()
    composer = AnswerComposer(llm)
    item = EvidenceItem(
        evidence_id="S7",
        kb="V16/V15",
        source="WRKY family",
        relation="participates_in",
        target="stress response",
        evidence="WRKY proteins participate in stress responses.",
    )

    composer.compose("WRKY有哪些功能？", RetrievalHypothesis(), [item], {}, [])

    assert "[S7]" in llm.user_prompt
    assert "Preferred citation: unavailable" in llm.user_prompt
    assert "Omit factual claims whose Preferred citation is unavailable" in (
        composer.SYSTEM + llm.user_prompt
    )


def test_evidence_pack_does_not_expose_uncorroborated_raw_species() -> None:
    llm = CapturingLLM()
    composer = AnswerComposer(llm)
    item = EvidenceItem(
        evidence_id="S1",
        kb="V16/V15",
        source="WRKY family",
        relation="participates_in",
        target="stress response",
        evidence="WRKY proteins participate in stress responses.",
        title="WRKY transcription factors in Juglans regia",
        species="Carya illinoinensis",
        pmid="34573426",
    )

    composer.compose("WRKY有哪些功能？", RetrievalHypothesis(), [item], {}, [])

    assert "Claim species: -" in llm.user_prompt
    assert "Paper species: -" in llm.user_prompt
    assert "Carya illinoinensis" not in llm.user_prompt


def test_answer_prompt_enforces_evidence_strength_and_developmental_scope() -> None:
    prompt = AnswerComposer.SYSTEM.lower()

    assert "expression" in prompt
    assert "association" in prompt
    assert "prediction" in prompt
    assert "major" in prompt
    assert "causal" in prompt
    assert "spike" in prompt
    assert "spikelet development" in prompt
    assert "flowering-time" in prompt
    assert "functional perturbation" in prompt


def test_answer_prompt_handles_major_gene_questions_with_only_indirect_evidence() -> None:
    prompt = AnswerComposer.SYSTEM.lower()

    assert "does not establish a major regulator" in prompt
    assert "never call expression evidence direct functional evidence" in prompt
    assert "heterologous" in prompt
    assert "native-species" in prompt


def test_answer_prompt_receives_agent_final_handoff_constraints() -> None:
    llm = CapturingLLM()
    composer = AnswerComposer(llm)

    composer.compose(
        "荔枝中有哪些基因家族已经被研究？",
        RetrievalHypothesis(),
        [],
        {},
        [],
        agent_handoff={
            "allowed": True,
            "evidence_ids": ["S1"],
            "chain_ids": ["AG1"],
            "answer_constraints": ["Only report Litchi-connected gene families"],
        },
    )

    assert "Agent-approved final handoff" in llm.user_prompt
    assert "Only report Litchi-connected gene families" in llm.user_prompt
    assert "rejected candidates" in llm.user_prompt