363 lines
13 KiB
Python
363 lines
13 KiB
Python
"""U4 测试: RetrieveKnowledgeTool - RAG 管线内置工具
|
|
|
|
测试 retrieve_knowledge 工具的创建、执行、自动注册和集成。
|
|
"""
|
|
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from agentkit.memory.base import Memory, MemoryItem
|
|
from agentkit.memory.retriever import MemoryRetriever, RetrieveKnowledgeTool
|
|
from agentkit.tools.base import Tool
|
|
|
|
|
|
# ── In-Memory Memory 实现（用于测试） ────────────────────
|
|
|
|
|
|
class InMemoryMemory(Memory):
|
|
"""基于内存的 Memory 实现，用于测试"""
|
|
|
|
def __init__(self):
|
|
self._store: dict[str, MemoryItem] = {}
|
|
|
|
async def store(self, key: str, value, metadata=None) -> None:
|
|
self._store[key] = MemoryItem(
|
|
key=key, value=value, metadata=metadata or {}, score=1.0
|
|
)
|
|
|
|
async def retrieve(self, key: str) -> MemoryItem | None:
|
|
return self._store.get(key)
|
|
|
|
async def search(self, query: str, top_k: int = 5, filters=None) -> list[MemoryItem]:
|
|
results = []
|
|
for item in self._store.values():
|
|
if query.lower() in str(item.value).lower() or query.lower() in item.key.lower():
|
|
results.append(item)
|
|
return results[:top_k]
|
|
|
|
async def delete(self, key: str) -> bool:
|
|
return self._store.pop(key, None) is not None
|
|
|
|
|
|
# ── TestRetrieveKnowledgeToolCreation ──────────────────────
|
|
|
|
|
|
class TestRetrieveKnowledgeToolCreation:
|
|
"""RetrieveKnowledgeTool 创建测试"""
|
|
|
|
def test_create_retrieve_tool_returns_tool_when_semantic_configured(self):
|
|
"""有 semantic memory 时 create_retrieve_tool() 返回 Tool"""
|
|
semantic = InMemoryMemory()
|
|
retriever = MemoryRetriever(semantic_memory=semantic)
|
|
|
|
tool = retriever.create_retrieve_tool()
|
|
|
|
assert tool is not None
|
|
assert isinstance(tool, Tool)
|
|
|
|
def test_create_retrieve_tool_returns_none_when_no_semantic(self):
|
|
"""无 semantic memory 时 create_retrieve_tool() 返回 None"""
|
|
retriever = MemoryRetriever()
|
|
|
|
tool = retriever.create_retrieve_tool()
|
|
|
|
assert tool is None
|
|
|
|
def test_create_retrieve_tool_with_working_only_returns_none(self):
|
|
"""仅有 working memory 时返回 None"""
|
|
working = InMemoryMemory()
|
|
retriever = MemoryRetriever(working_memory=working)
|
|
|
|
tool = retriever.create_retrieve_tool()
|
|
|
|
assert tool is None
|
|
|
|
def test_tool_has_correct_name(self):
|
|
"""工具名称为 retrieve_knowledge"""
|
|
semantic = InMemoryMemory()
|
|
retriever = MemoryRetriever(semantic_memory=semantic)
|
|
|
|
tool = retriever.create_retrieve_tool()
|
|
|
|
assert tool.name == "retrieve_knowledge"
|
|
|
|
def test_tool_has_description(self):
|
|
"""工具包含描述"""
|
|
semantic = InMemoryMemory()
|
|
retriever = MemoryRetriever(semantic_memory=semantic)
|
|
|
|
tool = retriever.create_retrieve_tool()
|
|
|
|
assert isinstance(tool.description, str)
|
|
assert len(tool.description) > 0
|
|
|
|
def test_tool_has_input_schema(self):
|
|
"""工具包含 input_schema"""
|
|
semantic = InMemoryMemory()
|
|
retriever = MemoryRetriever(semantic_memory=semantic)
|
|
|
|
tool = retriever.create_retrieve_tool()
|
|
|
|
assert tool.input_schema is not None
|
|
assert tool.input_schema["type"] == "object"
|
|
assert "query" in tool.input_schema["properties"]
|
|
assert "query" in tool.input_schema["required"]
|
|
|
|
def test_tool_is_retrieve_knowledge_tool_instance(self):
|
|
"""工具是 RetrieveKnowledgeTool 实例"""
|
|
semantic = InMemoryMemory()
|
|
retriever = MemoryRetriever(semantic_memory=semantic)
|
|
|
|
tool = retriever.create_retrieve_tool()
|
|
|
|
assert isinstance(tool, RetrieveKnowledgeTool)
|
|
|
|
|
|
# ── TestRetrieveKnowledgeToolExecution ─────────────────────
|
|
|
|
|
|
class TestRetrieveKnowledgeToolExecution:
|
|
"""RetrieveKnowledgeTool 执行测试"""
|
|
|
|
async def test_execute_calls_retriever_retrieve(self):
|
|
"""execute() 调用 MemoryRetriever.retrieve()"""
|
|
semantic = InMemoryMemory()
|
|
await semantic.store("s1", "AI趋势报告", metadata={"source": "report.pdf"})
|
|
retriever = MemoryRetriever(semantic_memory=semantic)
|
|
tool = retriever.create_retrieve_tool()
|
|
|
|
result = await tool.execute(query="AI趋势")
|
|
|
|
assert "results" in result
|
|
assert len(result["results"]) >= 1
|
|
|
|
async def test_execute_results_formatted_correctly(self):
|
|
"""结果包含 content, score, source, document_title"""
|
|
semantic = InMemoryMemory()
|
|
await semantic.store(
|
|
"s1",
|
|
"AI趋势报告内容",
|
|
metadata={"source": "report.pdf", "document_title": "2024 AI Report"},
|
|
)
|
|
retriever = MemoryRetriever(semantic_memory=semantic)
|
|
tool = retriever.create_retrieve_tool()
|
|
|
|
result = await tool.execute(query="AI趋势")
|
|
|
|
assert "results" in result
|
|
for item in result["results"]:
|
|
assert "content" in item
|
|
assert "score" in item
|
|
assert "source" in item
|
|
assert "document_title" in item
|
|
|
|
async def test_execute_empty_query_returns_error(self):
|
|
"""空 query 返回错误"""
|
|
semantic = InMemoryMemory()
|
|
retriever = MemoryRetriever(semantic_memory=semantic)
|
|
tool = retriever.create_retrieve_tool()
|
|
|
|
result = await tool.execute(query="")
|
|
|
|
assert "error" in result
|
|
assert result["results"] == []
|
|
|
|
async def test_execute_max_calls_limit(self):
|
|
"""超过 max_calls 限制后返回错误"""
|
|
semantic = InMemoryMemory()
|
|
await semantic.store("s1", "Some content")
|
|
retriever = MemoryRetriever(semantic_memory=semantic)
|
|
tool = retriever.create_retrieve_tool(max_calls=3)
|
|
|
|
# 前 3 次调用应该成功
|
|
for i in range(3):
|
|
result = await tool.execute(query="content")
|
|
assert "error" not in result or result.get("call_count") == i + 1
|
|
|
|
# 第 4 次调用应该返回错误
|
|
result = await tool.execute(query="content")
|
|
assert "error" in result
|
|
assert "Maximum retrieval calls" in result["error"]
|
|
assert result["results"] == []
|
|
|
|
async def test_execute_call_count_tracking(self):
|
|
"""call_count 在响应中正确跟踪"""
|
|
semantic = InMemoryMemory()
|
|
await semantic.store("s1", "Some content")
|
|
retriever = MemoryRetriever(semantic_memory=semantic)
|
|
tool = retriever.create_retrieve_tool(max_calls=5)
|
|
|
|
for i in range(1, 4):
|
|
result = await tool.execute(query="content")
|
|
assert result["call_count"] == i
|
|
|
|
async def test_execute_exception_handling(self):
|
|
"""retriever 抛出异常时返回错误响应"""
|
|
retriever = MemoryRetriever(semantic_memory=InMemoryMemory())
|
|
tool = retriever.create_retrieve_tool()
|
|
|
|
# Mock retriever.retrieve to raise exception
|
|
tool._retriever.retrieve = AsyncMock(side_effect=Exception("Service unavailable"))
|
|
|
|
result = await tool.execute(query="test")
|
|
|
|
assert "error" in result
|
|
assert "Service unavailable" in result["error"]
|
|
assert result["results"] == []
|
|
|
|
async def test_execute_returns_query_in_response(self):
|
|
"""响应中包含原始查询"""
|
|
semantic = InMemoryMemory()
|
|
await semantic.store("s1", "Some content")
|
|
retriever = MemoryRetriever(semantic_memory=semantic)
|
|
tool = retriever.create_retrieve_tool()
|
|
|
|
result = await tool.execute(query="AI趋势")
|
|
|
|
assert result["query"] == "AI趋势"
|
|
|
|
|
|
# ── TestRetrieveKnowledgeToolAutoRegistration ──────────────
|
|
|
|
|
|
class TestRetrieveKnowledgeToolAutoRegistration:
|
|
"""RetrieveKnowledgeTool 自动注册测试"""
|
|
|
|
def test_agent_with_semantic_memory_has_tool(self):
|
|
"""ConfigDrivenAgent 配置了 semantic memory 时自动注册 retrieve_knowledge"""
|
|
from agentkit.core.config_driven import AgentConfig, ConfigDrivenAgent
|
|
|
|
config = AgentConfig.from_dict({
|
|
"name": "test_agent",
|
|
"agent_type": "test",
|
|
"task_mode": "llm_generate",
|
|
"prompt": {
|
|
"identity": "Test agent",
|
|
"instructions": "Test",
|
|
},
|
|
"memory": {
|
|
"semantic": {
|
|
"enabled": True,
|
|
"base_url": "http://localhost:8080",
|
|
"knowledge_base_ids": ["kb1"],
|
|
},
|
|
},
|
|
})
|
|
|
|
# Patch imports inside the try block of ConfigDrivenAgent.__init__
|
|
with patch("agentkit.memory.http_rag.HttpRAGService") as mock_rag, \
|
|
patch("agentkit.memory.semantic.SemanticMemory") as mock_sem:
|
|
mock_sem.return_value = InMemoryMemory()
|
|
agent = ConfigDrivenAgent(config=config)
|
|
|
|
tool_names = [t.name for t in agent._tools]
|
|
assert "retrieve_knowledge" in tool_names
|
|
|
|
def test_agent_without_semantic_memory_does_not_have_tool(self):
|
|
"""ConfigDrivenAgent 未配置 semantic memory 时不注册 retrieve_knowledge"""
|
|
from agentkit.core.config_driven import AgentConfig, ConfigDrivenAgent
|
|
|
|
config = AgentConfig.from_dict({
|
|
"name": "test_agent",
|
|
"agent_type": "test",
|
|
"task_mode": "llm_generate",
|
|
"prompt": {
|
|
"identity": "Test agent",
|
|
"instructions": "Test",
|
|
},
|
|
})
|
|
|
|
agent = ConfigDrivenAgent(config=config)
|
|
|
|
tool_names = [t.name for t in agent._tools]
|
|
assert "retrieve_knowledge" not in tool_names
|
|
|
|
def test_auto_registered_tool_is_retrieve_knowledge_instance(self):
|
|
"""自动注册的工具是 RetrieveKnowledgeTool 实例"""
|
|
from agentkit.core.config_driven import AgentConfig, ConfigDrivenAgent
|
|
|
|
config = AgentConfig.from_dict({
|
|
"name": "test_agent",
|
|
"agent_type": "test",
|
|
"task_mode": "llm_generate",
|
|
"prompt": {
|
|
"identity": "Test agent",
|
|
"instructions": "Test",
|
|
},
|
|
"memory": {
|
|
"semantic": {
|
|
"enabled": True,
|
|
"base_url": "http://localhost:8080",
|
|
"knowledge_base_ids": ["kb1"],
|
|
},
|
|
},
|
|
})
|
|
|
|
with patch("agentkit.memory.http_rag.HttpRAGService"), \
|
|
patch("agentkit.memory.semantic.SemanticMemory") as mock_sem:
|
|
mock_sem.return_value = InMemoryMemory()
|
|
agent = ConfigDrivenAgent(config=config)
|
|
|
|
retrieve_tools = [t for t in agent._tools if t.name == "retrieve_knowledge"]
|
|
assert len(retrieve_tools) == 1
|
|
assert isinstance(retrieve_tools[0], RetrieveKnowledgeTool)
|
|
|
|
|
|
# ── TestRetrieveKnowledgeToolIntegration ───────────────────
|
|
|
|
|
|
class TestRetrieveKnowledgeToolIntegration:
|
|
"""RetrieveKnowledgeTool 集成测试"""
|
|
|
|
async def test_tool_works_with_query_transformer(self):
|
|
"""工具配合 query transformer 工作"""
|
|
from agentkit.memory.query_transformer import QueryTransformerBase, TransformedQuery
|
|
|
|
class SimpleTransformer(QueryTransformerBase):
|
|
async def transform(self, query: str) -> TransformedQuery:
|
|
return TransformedQuery(
|
|
main_query=f"enhanced: {query}",
|
|
sub_queries=[],
|
|
)
|
|
|
|
semantic = InMemoryMemory()
|
|
await semantic.store("s1", "enhanced: AI trends data")
|
|
retriever = MemoryRetriever(
|
|
semantic_memory=semantic,
|
|
query_transformer=SimpleTransformer(),
|
|
)
|
|
tool = retriever.create_retrieve_tool()
|
|
|
|
result = await tool.execute(query="AI")
|
|
|
|
assert "results" in result
|
|
|
|
async def test_tool_returns_structured_results_for_llm(self):
|
|
"""工具返回 LLM 可用的结构化结果"""
|
|
semantic = InMemoryMemory()
|
|
await semantic.store(
|
|
"s1",
|
|
"GEO optimization improves brand visibility",
|
|
metadata={"source": "guide.md", "document_title": "GEO Guide"},
|
|
)
|
|
await semantic.store(
|
|
"s2",
|
|
"Another relevant document about SEO",
|
|
metadata={"source": "seo.md", "document_title": "SEO Basics"},
|
|
)
|
|
retriever = MemoryRetriever(semantic_memory=semantic)
|
|
tool = retriever.create_retrieve_tool()
|
|
|
|
result = await tool.execute(query="optimization")
|
|
|
|
assert isinstance(result, dict)
|
|
assert "query" in result
|
|
assert "results" in result
|
|
assert "call_count" in result
|
|
assert isinstance(result["results"], list)
|
|
for item in result["results"]:
|
|
assert isinstance(item, dict)
|
|
assert "content" in item
|
|
assert "score" in item
|