Files
PlantAI-rag/tests/test_qa_v1_1_73_subquestion_pipeline.py

611 lines
20 KiB
Python

from __future__ import annotations
from concurrent.futures import ThreadPoolExecutor
import inspect
from pathlib import Path
from threading import Event, Lock
import time
from src.rag.qa_v1_1_73.config import QAConfig
from src.rag.qa_v1_1_73.pipeline import QAV1Pipeline
from src.rag.qa_v1_1_73.schemas import EvidenceItem, HypothesisEntity, RetrievalHypothesis
from src.rag.qa_v1_1_73.subquestion_orchestrator import SubquestionRetrievalOrchestrator
def _config() -> QAConfig:
return QAConfig(
project_root=Path("/tmp"),
gfanalysis_root=Path("/tmp"),
alias_path=Path("/tmp/alias.yaml"),
taxonomy_path=Path("/tmp/taxonomy.yaml"),
reranker_model=Path("/tmp/reranker"),
initial_recall_per_query=50,
max_query_variants=6,
max_evidence=40,
min_evidence=8,
max_rounds=2,
graph_hops=2,
include_candidate=False,
enable_provisional_draft=True,
enable_second_round=True,
enable_bge_rerank=True,
graph_schema="legacy",
enable_agentic_controller=True,
agentic_max_replans=1,
enable_question_decomposition=True,
)
def test_worker_configuration_is_bounded(monkeypatch) -> None:
monkeypatch.setenv("QA_V1_SUBQUESTION_WORKERS", "99")
monkeypatch.setenv("QA_V1_VECTOR_QUERY_WORKERS", "0")
config = QAConfig.from_env()
assert config.subquestion_workers == 3
assert config.vector_query_workers == 1
def test_worker_configuration_preserves_safe_values(monkeypatch) -> None:
monkeypatch.setenv("QA_V1_SUBQUESTION_WORKERS", "2")
monkeypatch.setenv("QA_V1_VECTOR_QUERY_WORKERS", "2")
config = QAConfig.from_env()
assert config.subquestion_workers == 2
assert config.vector_query_workers == 2
def test_shared_reranker_is_serialized_by_pipeline_lock() -> None:
class TrackingReranker:
def __init__(self):
self.lock = Lock()
self.active = 0
self.max_active = 0
def rerank(self, question, items):
with self.lock:
self.active += 1
self.max_active = max(self.max_active, self.active)
try:
time.sleep(0.05)
return list(items), []
finally:
with self.lock:
self.active -= 1
pipeline = QAV1Pipeline.__new__(QAV1Pipeline)
pipeline.reranker = TrackingReranker()
pipeline._reranker_lock = Lock()
with ThreadPoolExecutor(max_workers=2) as executor:
futures = [
executor.submit(pipeline._rerank, f"Q{index}", [])
for index in range(2)
]
[future.result() for future in futures]
assert pipeline.reranker.max_active == 1
def _parent_hypothesis() -> RetrievalHypothesis:
return RetrievalHypothesis(
question_type="mechanism",
direct_question_entities=[
HypothesisEntity(
mention="HSP基因家族",
canonical_candidate="HSP gene family",
entity_type="gene_family",
origin="question",
confidence=0.9,
verification_status="question_direct_unverified",
),
HypothesisEntity(
mention="高温",
canonical_candidate="high temperature",
entity_type="condition",
origin="question",
confidence=0.9,
verification_status="question_direct_unverified",
),
],
relation_intents=["heat stress response", "chaperone function"],
concept_queries=["parent query must not leak"],
answer_slots=["parent slot must not leak"],
)
def _task(qid: str, question: str) -> dict:
return {
"subquestion_id": qid,
"question": question,
"role": "required",
"answer_slot": "molecular action",
"anchors": [
{
"mention": "HSP基因家族",
"canonical_name": "HSP gene family",
"canonical_id": "",
"entity_type": "gene_family",
"role": "start_anchor",
}
],
"target_entity_types": ["process", "protein", "trait"],
"relation_intents": ["chaperone function"],
"constraints": ["heat stress"],
"retrieval": {
"vector_top_k": 30,
"retrieve_native_chains": True,
"max_hops": 3,
"allow_incoming": True,
"allow_outgoing": True,
"max_attempts": 2,
},
"acceptance_requirements": {
"requires_anchor_connection": False,
"requires_terminal_type_match": True,
"requires_structured_assertion": True,
"requires_paper_provenance": True,
},
}
def _evidence(evidence_id: str, relation: str = "prevents") -> EvidenceItem:
return EvidenceItem(
evidence_id=evidence_id,
kb="test",
source="HSPs",
relation=relation,
target="protein aggregation",
evidence="HSPs prevent protein aggregation.",
paper_id="paper-hsp",
doi="10.1000/hsp",
relation_id=evidence_id,
raw={"source_type": "protein", "target_type": "process"},
)
def test_anchor_only_agent_plan_keeps_vector_support_out_of_directflow_neo4j():
pipeline = QAV1Pipeline.__new__(QAV1Pipeline)
switches = pipeline._direct_flow_switches(
{
"anchors": [{"canonical_name": "Litchi chinensis"}],
"vector_plan": {"enabled": True},
"graph_plan": {
"enabled": True,
"anchor_only_start": True,
"exclude_unrelated_vector_start_seeds": True,
},
},
use_qdrant=True,
use_neo4j=True,
use_graph_expansion=True,
graph_schema="canonical",
)
assert switches == {
"use_qdrant": True,
"use_neo4j": False,
"use_graph_expansion": False,
}
def test_legacy_schema_disables_direct_and_native_canonical_graph_paths():
pipeline = QAV1Pipeline.__new__(QAV1Pipeline)
switches = pipeline._direct_flow_switches(
{
"vector_plan": {"enabled": True},
"graph_plan": {"enabled": True},
},
use_qdrant=True,
use_neo4j=True,
use_graph_expansion=True,
graph_schema="legacy",
)
assert switches == {
"use_qdrant": True,
"use_neo4j": False,
"use_graph_expansion": False,
}
assert "disabled_for_legacy_schema" in inspect.getsource(QAV1Pipeline.answer)
assert "disabled_for_legacy_schema" in inspect.getsource(
QAV1Pipeline._retrieve_subquestion_once
)
def test_anchor_only_plan_uses_same_switch_to_skip_duplicate_native_graph():
answer_source = inspect.getsource(QAV1Pipeline.answer)
child_source = inspect.getsource(QAV1Pipeline._retrieve_subquestion_once)
assert (
'if native_chain_enabled and direct_flow_switches["use_neo4j"]'
in answer_source
)
assert 'if bool(options.get("native_chain_enabled", True)) and bool(' in child_source
assert 'direct_flow_switches["use_neo4j"]' in child_source
def test_agentic_anchor_candidate_helper_uses_factory_and_closes_retriever():
evidence = _evidence("anchor-1")
chain = {"chain_id": "AG1", "hops": [{"assertion_id": "anchor-1"}]}
instances = []
class TrackingAnchorRetriever:
def __init__(self):
self.plan = None
self.closed = False
instances.append(self)
def retrieve(self, plan):
self.plan = plan
return [evidence], [chain], {
"query_mode": "canonical_anchor_neighborhood",
"status": "ok",
}
def close(self):
self.closed = True
pipeline = QAV1Pipeline.__new__(QAV1Pipeline)
pipeline.agentic_graph_retriever_factory = TrackingAnchorRetriever
plan = {"anchors": [{"canonical_name": "Litchi chinensis"}]}
returned = pipeline._retrieve_agentic_anchor_candidates(plan)
assert returned[0] == [evidence]
assert returned[1] == [chain]
assert returned[2]["query_mode"] == "canonical_anchor_neighborhood"
assert instances[0].plan == plan
assert instances[0].closed is True
def test_child_hypothesis_is_isolated_to_subquestion_answer_slot_and_queries():
pipeline = QAV1Pipeline.__new__(QAV1Pipeline)
child = pipeline._hypothesis_for_subquestion(
_parent_hypothesis(),
_task("Q1", "HSP如何防止高温造成的蛋白质损伤?"),
{"vector_queries": ["HSP protein refolding heat damage"]},
)
assert child.question_type == "mechanism"
assert [item.canonical_candidate for item in child.direct_question_entities] == [
"HSP gene family",
"high temperature",
]
assert child.relation_intents[0] == "chaperone function"
assert "prevent_protein_aggregation" in child.relation_intents
assert "protein_refolding" in child.relation_intents
assert child.answer_slots == ["molecular action"]
assert child.concept_queries[:2] == [
"HSP如何防止高温造成的蛋白质损伤?",
"HSP protein refolding heat damage",
]
assert "parent query must not leak" not in child.concept_queries
def test_mechanism_subquestion_runtime_contract_caps_expensive_graph_search():
pipeline = QAV1Pipeline.__new__(QAV1Pipeline)
limits = pipeline._subquestion_runtime_limits(
_task("Q0", "HSP基因家族如何帮助植物应对高温?"),
_parent_hypothesis(),
{},
native_chain_max_hops=3,
native_chain_beam_width=60,
native_chain_per_node_expand=12,
)
assert limits == {
"vector_top_k": 30,
"max_hops": 1,
"beam_width": 12,
"per_node_expand": 4,
}
def test_q0_runtime_uses_decomposition_type_when_parent_is_function():
pipeline = QAV1Pipeline.__new__(QAV1Pipeline)
parent = _parent_hypothesis()
parent.question_type = "function"
task = _task("Q0", "GST基因家族在植物胁迫响应中有什么作用?")
task["question_type"] = "mechanism"
limits = pipeline._subquestion_runtime_limits(
task,
parent,
{},
native_chain_max_hops=3,
native_chain_beam_width=90,
native_chain_per_node_expand=24,
)
assert limits == {
"vector_top_k": 30,
"max_hops": 1,
"beam_width": 12,
"per_node_expand": 4,
}
def test_question_decomposition_starts_in_background_before_q0_retrieval_finishes():
started = Event()
release = Event()
class BlockingDecomposer:
def decompose(self, question, **kwargs):
started.set()
assert release.wait(timeout=1)
return {"original_question": question, "subquestions": []}
pipeline = QAV1Pipeline.__new__(QAV1Pipeline)
pipeline.question_decomposer = BlockingDecomposer()
future = pipeline._start_question_decomposition(
"GST基因家族在植物胁迫响应中有什么作用?",
verified_entities=[],
relation_intents=[],
answer_slots=[],
)
assert started.wait(timeout=1)
assert future.done() is False
release.set()
decomposition, elapsed = future.result(timeout=1)
assert decomposition["original_question"].startswith("GST")
assert elapsed >= 0
def test_answer_overlaps_decomposition_with_q0_top30_then_uses_q0_graph_limits():
source = inspect.getsource(QAV1Pipeline.answer)
start = source.index("_start_question_decomposition")
direct_recall = source.index("retriever.retrieve(")
await_decomposition = source.index("decomposition_future.result()")
native_chains = source.index("_retrieve_native_chains(")
assert start < direct_recall < await_decomposition < native_chains
assert 'max_hops=q0_runtime_limits["max_hops"]' in source
assert 'beam_width=q0_runtime_limits["beam_width"]' in source
assert 'per_node_expand=q0_runtime_limits["per_node_expand"]' in source
def test_decomposition_keeps_real_agent_plan_for_q0():
source = inspect.getsource(QAV1Pipeline.answer)
start = source.index("agentic_plan: Dict[str, Any]")
end = source.index('timings["agentic_retrieval_plan"]', start)
planning_block = source[start:end]
assert "self.agentic_controller.plan(question, hypothesis)" in planning_block
assert "self.agentic_controller._fallback_plan" not in planning_block
def test_child_retrieval_uses_agent_plan_and_merges_anchor_candidates():
controller_calls = []
direct_calls = []
anchor_calls = []
class PlanningController:
def plan(self, question, hypothesis):
controller_calls.append((question, hypothesis))
return {
"anchors": [{"canonical_name": "Litchi chinensis"}],
"vector_plan": {"enabled": True},
"graph_plan": {
"enabled": True,
"anchor_only_start": True,
"exclude_unrelated_vector_start_seeds": True,
"required_terminal_types": ["gene_family"],
},
}
class TrackingDirectRetriever:
def retrieve(self, question, queries, hypothesis, **kwargs):
direct_calls.append(kwargs)
return [_evidence("direct-1")], {"seed_extraction": {}}
def close(self):
return None
class TrackingAnchorRetriever:
def retrieve(self, plan):
anchor_calls.append(plan)
return (
[_evidence("anchor-1", "identified_in")],
[{"chain_id": "AG1", "hops": [{"assertion_id": "anchor-1"}]}],
{
"query_mode": "canonical_anchor_neighborhood",
"status": "ok",
},
)
def close(self):
return None
class QueryBuilder:
@staticmethod
def build(question, provisional_answer, hypothesis):
return {"question_lane": [question], "hypothesis_lane": []}
class Reranker:
@staticmethod
def rerank(question, items):
return list(items), []
pipeline = QAV1Pipeline.__new__(QAV1Pipeline)
pipeline.config = _config()
pipeline.agentic_controller = PlanningController()
pipeline.direct_retriever_factory = TrackingDirectRetriever
pipeline.agentic_graph_retriever_factory = TrackingAnchorRetriever
pipeline.query_builder = QueryBuilder()
pipeline.reranker = Reranker()
pipeline._reranker_lock = Lock()
pipeline._filter_public_answer_candidates = lambda items, **kwargs: (list(items), {})
pipeline._pre_rerank = lambda items, limit: list(items)
pipeline._select_graph_balanced = lambda items, **kwargs: (list(items), {})
result = pipeline._retrieve_subquestion_once(
_task("Q1", "荔枝中有哪些基因家族已经做过研究?"),
_parent_hypothesis(),
attempt=1,
refinement={},
options={
"qdrant_top_k": 120,
"graph_top_k": 80,
"rerank_limit": 160,
"use_qdrant": True,
"use_neo4j": True,
"use_graph_expansion": True,
"use_bge_rerank": True,
"native_chain_enabled": False,
"native_chain_max_hops": 3,
"native_chain_beam_width": 60,
"native_chain_per_node_expand": 12,
},
)
assert controller_calls[0][0].startswith("荔枝")
assert direct_calls[0]["use_qdrant"] is True
assert direct_calls[0]["use_neo4j"] is False
assert direct_calls[0]["use_graph_expansion"] is False
assert anchor_calls == [result["retrieval_audit"]["agentic_plan"]]
assert {item.relation_id for item in result["evidence_top30"]} == {
"direct-1",
"anchor-1",
}
assert result["chains"][0]["chain_id"].startswith("Q1.A1.")
assert (
result["retrieval_audit"]["agentic_anchor_audit"]["query_mode"]
== "canonical_anchor_neighborhood"
)
class _AcceptingController:
def review_subquestion_groups(self, question, groups, *, attempt):
return {
"groups": [
{
"subquestion_id": group["subquestion_id"],
"status": "SATISFIED",
"accepted_evidence_ids": [group["evidence"][0].evidence_id],
"accepted_chain_ids": [],
"refinement": {},
}
for group in groups
]
}
class _PipelineWithFakeSubquestionRetrieval(QAV1Pipeline):
def _retrieve_subquestion_once(self, task, parent_hypothesis, *, attempt, refinement, options):
self.retrieval_calls.append((task["subquestion_id"], attempt, options["top_k_for_llm"]))
return {
"evidence_top30": [_evidence(f'{task["subquestion_id"]}.A{attempt}.E1')],
"chains": [],
"retrieval_audit": {"vector_top_k": options["top_k_for_llm"]},
}
def test_pipeline_orchestration_reuses_q0_top30_and_keeps_rejected_candidates_in_audit():
pipeline = _PipelineWithFakeSubquestionRetrieval.__new__(
_PipelineWithFakeSubquestionRetrieval
)
pipeline.subquestion_orchestrator = SubquestionRetrievalOrchestrator()
pipeline.agentic_controller = _AcceptingController()
pipeline.retrieval_calls = []
q0_top30 = [_evidence("Q0.A1.E1"), _evidence("Q0.A1.E2", "upregulated_in")]
decomposition = {
"original_question": "HSP基因家族如何帮助植物应对高温?",
"subquestions": [
_task("Q0", "HSP基因家族如何帮助植物应对高温?"),
_task("Q1", "HSP如何防止蛋白质损伤?"),
],
"synthesis_order": ["Q0", "Q1"],
}
result = pipeline._run_subquestion_orchestration(
"HSP基因家族如何帮助植物应对高温?",
_parent_hypothesis(),
decomposition,
q0_bundle={
"evidence_top30": q0_top30,
"chains": [],
"retrieval_audit": {"vector_top_k": 30},
},
options={"top_k_for_llm": 30},
)
assert pipeline.retrieval_calls == [("Q1", 1, 30)]
assert [item.evidence_id for item in result["accepted_groups"][0]["evidence"]] == [
"Q0.A1.E1"
]
assert [
item.evidence_id
for item in result["audit_groups"][0]["attempts"][0]["evidence_top30"]
] == ["Q0.A1.E1", "Q0.A1.E2"]
def test_config_enables_question_decomposition_by_default(monkeypatch):
monkeypatch.delenv("QA_ENABLE_QUESTION_DECOMPOSITION", raising=False)
assert QAConfig.from_env().enable_question_decomposition is True
def test_q0_without_decomposition_preserves_pre_1114_native_graph_limits():
pipeline = QAV1Pipeline.__new__(QAV1Pipeline)
limits = pipeline._subquestion_runtime_limits(
_task("Q0", "HSP基因家族如何帮助植物应对高温?"),
_parent_hypothesis(),
{},
native_chain_max_hops=3,
native_chain_beam_width=60,
native_chain_per_node_expand=12,
apply_subquestion_constraints=False,
)
assert limits == {
"vector_top_k": 30,
"max_hops": 3,
"beam_width": 60,
"per_node_expand": 12,
}
class _EmptyDomain:
@staticmethod
def direct_entities_from_question(question):
return []
def test_pipeline_recognizes_literal_uppercase_gene_family_when_llm_returns_no_entities():
pipeline = QAV1Pipeline.__new__(QAV1Pipeline)
pipeline.domain = _EmptyDomain()
hypothesis = RetrievalHypothesis(question_type="mechanism")
pipeline._complete_direct_question_entities(
"HSP基因家族如何帮助植物应对高温?", hypothesis
)
assert [item.mention for item in hypothesis.direct_question_entities] == [
"HSP基因家族"
]
assert hypothesis.direct_question_entities[0].canonical_candidate == "HSP gene family"
assert hypothesis.direct_question_entities[0].origin == "question"
def test_mechanism_child_adds_stable_action_queries_when_planner_relations_are_empty():
pipeline = QAV1Pipeline.__new__(QAV1Pipeline)
task = _task("Q2", "What mechanisms does the HSP gene family use under heat stress?")
task["relation_intents"] = []
child = pipeline._hypothesis_for_subquestion(_parent_hypothesis(), task, {})
assert "prevent_protein_aggregation" in child.relation_intents
assert "protein_refolding" in child.relation_intents
assert "protein_homeostasis" in child.relation_intents
assert any("prevent protein aggregation" in value for value in child.concept_queries)
assert any("refolding heat-damaged proteins" in value for value in child.concept_queries)