305 lines
11 KiB
Python
305 lines
11 KiB
Python
from __future__ import annotations
|
||
|
||
from src.rag.qa_v1_1_73.question_decomposition import (
|
||
QUESTION_DECOMPOSITION_PROMPT,
|
||
QuestionDecompositionPlanner,
|
||
)
|
||
|
||
|
||
class FakeLLM:
|
||
available = True
|
||
|
||
def __init__(self, response):
|
||
self.response = response
|
||
self.calls = []
|
||
|
||
def complete_json(self, system, user, *, max_tokens=2500):
|
||
self.calls.append((system, user, max_tokens))
|
||
if isinstance(self.response, Exception):
|
||
raise self.response
|
||
return self.response
|
||
|
||
|
||
def _subquestion(index: int) -> dict:
|
||
return {
|
||
"subquestion_id": f"Q{index}",
|
||
"question": f"子问题{index}",
|
||
"role": "required" if index == 0 else "supporting",
|
||
"answer_slot": f"槽位{index}",
|
||
"anchors": [],
|
||
"target_entity_types": ["unknown"],
|
||
"relation_intents": ["lookup"],
|
||
"constraints": [],
|
||
"retrieval": {
|
||
"vector_top_k": 99,
|
||
"retrieve_native_chains": True,
|
||
"max_hops": 9,
|
||
"allow_incoming": True,
|
||
"allow_outgoing": True,
|
||
"max_attempts": 3,
|
||
},
|
||
"acceptance_requirements": {
|
||
"requires_anchor_connection": False,
|
||
"requires_terminal_type_match": True,
|
||
"requires_structured_assertion": True,
|
||
"requires_paper_provenance": True,
|
||
},
|
||
"contribution_to_final_answer": f"回答槽位{index}",
|
||
}
|
||
|
||
|
||
def test_prompt_contains_rules_schema_and_two_domain_examples():
|
||
assert "question_decomposition_v1" in QUESTION_DECOMPOSITION_PROMPT
|
||
assert "Always create Q0" in QUESTION_DECOMPOSITION_PROMPT
|
||
assert "vector_top_k=30" in QUESTION_DECOMPOSITION_PROMPT
|
||
assert "Attempt 3 is forbidden" in QUESTION_DECOMPOSITION_PROMPT
|
||
assert "R2R3-MYB" in QUESTION_DECOMPOSITION_PROMPT
|
||
assert "Litchi chinensis" in QUESTION_DECOMPOSITION_PROMPT
|
||
assert '"subquestions"' in QUESTION_DECOMPOSITION_PROMPT
|
||
|
||
|
||
def test_planner_enforces_fixed_retrieval_contract_and_six_question_cap():
|
||
response = {
|
||
"mode": "DECOMPOSE",
|
||
"prompt_version": "invented",
|
||
"original_question": "changed",
|
||
"question_type": "naming",
|
||
"answer_goal": "解释命名",
|
||
"global_constraints": {},
|
||
"subquestions": [_subquestion(index) for index in range(7)],
|
||
"synthesis_order": [f"Q{index}" for index in range(7)],
|
||
"decomposition_summary": "test",
|
||
}
|
||
planner = QuestionDecompositionPlanner(FakeLLM(response))
|
||
|
||
plan = planner.decompose(
|
||
"R2R3-MYB蛋白为什么被称为R2R3-MYB？",
|
||
verified_entities=[],
|
||
relation_intents=["naming_origin"],
|
||
answer_slots=["名称来源"],
|
||
)
|
||
|
||
assert plan["prompt_version"] == "question_decomposition_v1"
|
||
assert plan["original_question"] == "R2R3-MYB蛋白为什么被称为R2R3-MYB？"
|
||
assert len(plan["subquestions"]) == 6
|
||
assert plan["subquestions"][0]["subquestion_id"] == "Q0"
|
||
assert plan["subquestions"][0]["question"] == "R2R3-MYB蛋白为什么被称为R2R3-MYB？"
|
||
assert all(item["retrieval"]["vector_top_k"] == 30 for item in plan["subquestions"])
|
||
assert all(item["retrieval"]["max_attempts"] == 2 for item in plan["subquestions"])
|
||
assert all(item["retrieval"]["max_hops"] <= 3 for item in plan["subquestions"])
|
||
assert plan["synthesis_order"] == [f"Q{index}" for index in range(6)]
|
||
assert plan["plan_source"] == "llm_agent"
|
||
|
||
|
||
def test_planner_falls_back_to_atomic_q0_when_agent_output_fails():
|
||
planner = QuestionDecompositionPlanner(FakeLLM(ValueError("bad json")))
|
||
|
||
plan = planner.decompose(
|
||
"WRKY转录因子的典型结构特征是什么？",
|
||
verified_entities=[{"mention": "WRKY", "canonical_name": "WRKY"}],
|
||
relation_intents=["structure"],
|
||
answer_slots=["典型结构"],
|
||
)
|
||
|
||
assert plan["plan_source"] == "deterministic_fallback"
|
||
assert len(plan["subquestions"]) == 1
|
||
assert plan["subquestions"][0]["subquestion_id"] == "Q0"
|
||
assert plan["subquestions"][0]["question"] == "WRKY转录因子的典型结构特征是什么？"
|
||
assert plan["subquestions"][0]["retrieval"]["vector_top_k"] == 30
|
||
assert plan["subquestions"][0]["retrieval"]["max_attempts"] == 2
|
||
|
||
|
||
def test_planner_rejects_unverified_anchors_and_relation_intents():
|
||
response = {
|
||
"subquestions": [
|
||
{
|
||
**_subquestion(0),
|
||
"anchors": [
|
||
{
|
||
"mention": "WRKY",
|
||
"canonical_name": "WRKY",
|
||
"canonical_id": "cent_wrky",
|
||
"entity_type": "gene_family",
|
||
},
|
||
{
|
||
"mention": "InventedFamily",
|
||
"canonical_name": "InventedFamily",
|
||
"canonical_id": "cent_fake",
|
||
"entity_type": "gene_family",
|
||
},
|
||
],
|
||
"relation_intents": ["structure", "invented_relation"],
|
||
}
|
||
]
|
||
}
|
||
planner = QuestionDecompositionPlanner(FakeLLM(response))
|
||
|
||
plan = planner.decompose(
|
||
"WRKY转录因子的典型结构特征是什么？",
|
||
verified_entities=[
|
||
{
|
||
"mention": "WRKY",
|
||
"canonical_name": "WRKY",
|
||
"canonical_id": "cent_wrky",
|
||
"entity_type": "gene_family",
|
||
}
|
||
],
|
||
relation_intents=["structure"],
|
||
answer_slots=["典型结构"],
|
||
)
|
||
|
||
q0 = plan["subquestions"][0]
|
||
assert [item["canonical_id"] for item in q0["anchors"]] == ["cent_wrky"]
|
||
assert q0["relation_intents"] == ["structure"]
|
||
|
||
|
||
def test_planner_restores_verified_hsp_anchor_and_normalizes_mechanism_target_types():
|
||
response = {
|
||
"question_type": "mechanism",
|
||
"subquestions": [
|
||
{
|
||
**_subquestion(0),
|
||
"question": "changed",
|
||
"anchors": [],
|
||
"target_entity_types": [
|
||
"molecular_mechanism",
|
||
"protein_function",
|
||
"cellular_process",
|
||
],
|
||
"acceptance_requirements": {
|
||
"requires_anchor_connection": True,
|
||
"requires_terminal_type_match": True,
|
||
"requires_structured_assertion": True,
|
||
"requires_paper_provenance": True,
|
||
},
|
||
},
|
||
{
|
||
**_subquestion(1),
|
||
"question": "What mechanisms does the HSP gene family use under heat stress?",
|
||
"anchors": [],
|
||
"target_entity_types": ["molecular_mechanism"],
|
||
},
|
||
],
|
||
}
|
||
planner = QuestionDecompositionPlanner(FakeLLM(response))
|
||
|
||
plan = planner.decompose(
|
||
"HSP基因家族如何帮助植物应对高温？",
|
||
verified_entities=[
|
||
{
|
||
"mention": "HSP基因家族",
|
||
"canonical_name": "HSP gene family",
|
||
"canonical_id": "",
|
||
"entity_type": "gene_family",
|
||
}
|
||
],
|
||
relation_intents=[],
|
||
answer_slots=["molecular mechanism"],
|
||
)
|
||
|
||
assert all(
|
||
[anchor["canonical_name"] for anchor in item["anchors"]]
|
||
== ["HSP gene family"]
|
||
for item in plan["subquestions"]
|
||
)
|
||
assert plan["subquestions"][0]["target_entity_types"] == [
|
||
"process",
|
||
"biological_process",
|
||
]
|
||
assert plan["subquestions"][1]["target_entity_types"] == ["process"]
|
||
|
||
|
||
def test_mechanism_planner_drops_supporting_questions_that_invent_new_relations_or_entities():
|
||
response = {
|
||
"question_type": "mechanism",
|
||
"subquestions": [
|
||
{
|
||
**_subquestion(0),
|
||
"relation_intents": ["protein_homeostasis"],
|
||
},
|
||
{
|
||
**_subquestion(1),
|
||
"question": "HSP如何通过分子伴侣作用防止蛋白质聚集？",
|
||
"relation_intents": ["protein_homeostasis"],
|
||
},
|
||
{
|
||
**_subquestion(2),
|
||
"question": "HSP如何与蛋白酶体协作清除错误折叠蛋白？",
|
||
"relation_intents": ["proteolysis_cooperation", "protein_homeostasis"],
|
||
},
|
||
{
|
||
**_subquestion(3),
|
||
"question": "HSP90如何调节HSF信号？",
|
||
"relation_intents": ["protein_homeostasis"],
|
||
},
|
||
{
|
||
**_subquestion(4),
|
||
"question": "不同植物中有哪些具体HSP基因参与高温响应？",
|
||
"relation_intents": ["protein_homeostasis"],
|
||
},
|
||
],
|
||
}
|
||
planner = QuestionDecompositionPlanner(FakeLLM(response))
|
||
|
||
plan = planner.decompose(
|
||
"HSP基因家族如何帮助植物应对高温？",
|
||
verified_entities=[
|
||
{
|
||
"mention": "HSP基因家族",
|
||
"canonical_name": "HSP gene family",
|
||
"canonical_id": "",
|
||
"entity_type": "gene_family",
|
||
}
|
||
],
|
||
relation_intents=["protein_homeostasis"],
|
||
answer_slots=["molecular mechanism"],
|
||
)
|
||
|
||
assert [item["question"] for item in plan["subquestions"]] == [
|
||
"HSP基因家族如何帮助植物应对高温？",
|
||
"HSP如何通过分子伴侣作用防止蛋白质聚集？",
|
||
]
|
||
assert plan["synthesis_order"] == ["Q0", "Q1"]
|
||
|
||
|
||
def test_mechanism_planner_adds_safe_broad_support_when_all_agent_support_is_rejected():
|
||
response = {
|
||
"question_type": "mechanism",
|
||
"subquestions": [
|
||
{**_subquestion(0), "relation_intents": ["protein_homeostasis"]},
|
||
{
|
||
**_subquestion(1),
|
||
"question": "HSP90如何通过HSF-HSE信号通路调控热响应？",
|
||
"relation_intents": ["invented_signaling"],
|
||
},
|
||
],
|
||
}
|
||
planner = QuestionDecompositionPlanner(FakeLLM(response))
|
||
|
||
plan = planner.decompose(
|
||
"HSP基因家族如何帮助植物应对高温？",
|
||
verified_entities=[
|
||
{
|
||
"mention": "HSP基因家族",
|
||
"canonical_name": "HSP gene family",
|
||
"entity_type": "gene_family",
|
||
},
|
||
{
|
||
"mention": "高温",
|
||
"canonical_name": "high temperature",
|
||
"entity_type": "condition",
|
||
},
|
||
],
|
||
relation_intents=["protein_homeostasis"],
|
||
answer_slots=["molecular mechanism"],
|
||
)
|
||
|
||
assert len(plan["subquestions"]) == 2
|
||
q1 = plan["subquestions"][1]
|
||
assert q1["subquestion_id"] == "Q1"
|
||
assert "直接分子作用" in q1["question"]
|
||
assert "被保护对象" in q1["question"]
|
||
assert "HSP90" not in q1["question"]
|
||
assert "HSF" not in q1["question"]
|
||
assert q1["retrieval"]["vector_top_k"] == 30
|