sladro
61026cd1b6
Co-authored-by: factory-droid[bot] <138933559+factory-droid[bot]@users.noreply.github.com>
107 lines
3.5 KiB
Python
107 lines
3.5 KiB
Python
"""Tests for multi-source RAG selection and aggregation."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any, Dict
|
|
|
|
from bidmaster.nodes.base import NodeContext
|
|
from bidmaster.nodes.content.generate_content import GenerateContentNode
|
|
from bidmaster.utils.document_context import DocumentContextMatch
|
|
|
|
|
|
class DummyTenderSearcher:
|
|
def __init__(self, label: str = "Tender snippet"):
|
|
self.label = label
|
|
self.last_query: str | None = None
|
|
|
|
def search(self, query: str, top_k: int): # pragma: no cover - simple stub
|
|
self.last_query = query
|
|
return [
|
|
DocumentContextMatch(
|
|
text=self.label,
|
|
section="章节一",
|
|
score=0.9,
|
|
source_type="paragraph",
|
|
metadata={},
|
|
)
|
|
]
|
|
|
|
|
|
class DummyRagTool:
|
|
def __init__(self, kb_content: str = "KB snippet"):
|
|
self.last_context: Dict[str, Any] | None = None
|
|
self.kb_content = kb_content
|
|
self.search_queries: list[str] = []
|
|
|
|
def search(self, query: str, k: int): # pragma: no cover - simple stub
|
|
self.search_queries.append(query)
|
|
return [{"content": self.kb_content, "metadata": {"source": "kb.docx"}}]
|
|
|
|
def generate_content(self, task_id: str, context: Dict[str, Any]):
|
|
self.last_context = context
|
|
return "generated"
|
|
|
|
|
|
def _base_state(chapter: Dict[str, Any], rag_tool: DummyRagTool) -> Dict[str, Any]:
|
|
return {
|
|
"current_chapter": chapter,
|
|
"chapter_queue": [chapter],
|
|
"chapter_configs": {},
|
|
"chapter_children_map": {},
|
|
"generated_contents": {},
|
|
"rag_tool": rag_tool,
|
|
}
|
|
|
|
|
|
def test_generate_content_auto_uses_all_available_sources():
|
|
chapter = {"id": "chapter_1", "title": "概述", "level": 1}
|
|
rag_tool = DummyRagTool()
|
|
state: Dict[str, Any] = _base_state(chapter, rag_tool)
|
|
state["available_rag_sources"] = {
|
|
"tender_doc": {"available": True},
|
|
"global_kb": {"available": True},
|
|
}
|
|
state["tender_doc_searcher"] = DummyTenderSearcher()
|
|
|
|
node = GenerateContentNode()
|
|
node.execute(state, NodeContext())
|
|
|
|
assert state["generated_contents"]["chapter_1"] == "generated"
|
|
assert rag_tool.last_context is not None
|
|
rag_context = rag_tool.last_context.get("rag_context")
|
|
assert "【招标文件】" in rag_context
|
|
assert "【知识库】" in rag_context
|
|
|
|
|
|
def test_generate_content_falls_back_when_tender_doc_missing():
|
|
chapter = {"id": "chapter_1", "title": "概述", "level": 1}
|
|
rag_tool = DummyRagTool(kb_content="Only KB")
|
|
state: Dict[str, Any] = _base_state(chapter, rag_tool)
|
|
state["available_rag_sources"] = {
|
|
"tender_doc": {"available": False},
|
|
"global_kb": {"available": True},
|
|
}
|
|
|
|
node = GenerateContentNode()
|
|
node.execute(state, NodeContext())
|
|
|
|
rag_context = rag_tool.last_context.get("rag_context")
|
|
assert "Only KB" in rag_context
|
|
assert "【招标文件】" not in rag_context
|
|
|
|
|
|
def test_generate_content_fallbacks_to_global_when_sources_unavailable():
|
|
chapter = {"id": "chapter_1", "title": "概述", "level": 1}
|
|
rag_tool = DummyRagTool(kb_content="fallback")
|
|
state: Dict[str, Any] = _base_state(chapter, rag_tool)
|
|
state["available_rag_sources"] = {
|
|
"tender_doc": {"available": False},
|
|
"global_kb": {"available": False},
|
|
}
|
|
|
|
node = GenerateContentNode()
|
|
node.execute(state, NodeContext())
|
|
|
|
# 即使标记不可用，也会退回全局知识库，保证填充不中断
|
|
assert "fallback" in rag_tool.last_context.get("rag_context")
|