436 lines
18 KiB
Python
436 lines
18 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
Cloudflare Memory Bank 实现测试
|
||
"""
|
||
|
||
import unittest
|
||
import asyncio
|
||
import os
|
||
import sys
|
||
from unittest.mock import patch, MagicMock, AsyncMock
|
||
from datetime import datetime
|
||
|
||
# 添加项目根目录到Python路径
|
||
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
||
|
||
from src.jixia.memory.cloudflare_memory_bank import CloudflareMemoryBank, MemoryEntry
|
||
|
||
|
||
class TestCloudflareMemoryBank(unittest.TestCase):
|
||
"""测试CloudflareMemoryBank类"""
|
||
|
||
def setUp(self):
|
||
"""测试前的设置"""
|
||
# Mock掉 aiohttp.ClientSession 以避免实际网络请求
|
||
self.patcher = patch('src.jixia.memory.cloudflare_memory_bank.aiohttp.ClientSession')
|
||
self.mock_session_class = self.patcher.start()
|
||
self.mock_session = AsyncMock()
|
||
self.mock_session_class.return_value = self.mock_session
|
||
|
||
# Mock掉 get_cloudflare_config 以避免实际读取配置
|
||
self.config_patcher = patch('src.jixia.memory.cloudflare_memory_bank.get_cloudflare_config')
|
||
self.mock_get_config = self.config_patcher.start()
|
||
self.mock_get_config.return_value = {
|
||
'account_id': 'test-account',
|
||
'api_token': 'test-token',
|
||
'vectorize_index': 'test-index',
|
||
'embed_model': '@cf/baai/bge-m3',
|
||
'autorag_domain': 'test.example.com'
|
||
}
|
||
|
||
# 创建CloudflareMemoryBank实例
|
||
self.memory_bank = CloudflareMemoryBank()
|
||
|
||
# 重置一些内部状态
|
||
self.memory_bank.config = self.mock_get_config.return_value
|
||
self.memory_bank.account_id = 'test-account'
|
||
self.memory_bank.api_token = 'test-token'
|
||
self.memory_bank.vectorize_index = 'test-index'
|
||
self.memory_bank.embed_model = '@cf/baai/bge-m3'
|
||
self.memory_bank.autorag_domain = 'test.example.com'
|
||
|
||
def tearDown(self):
|
||
"""测试后的清理"""
|
||
self.patcher.stop()
|
||
self.config_patcher.stop()
|
||
|
||
def test_init(self):
|
||
"""测试初始化"""
|
||
self.assertEqual(self.memory_bank.account_id, "test-account")
|
||
self.assertEqual(self.memory_bank.api_token, "test-token")
|
||
self.assertEqual(self.memory_bank.vectorize_index, "test-index")
|
||
self.assertEqual(self.memory_bank.embed_model, "@cf/baai/bge-m3")
|
||
self.assertEqual(self.memory_bank.autorag_domain, "test.example.com")
|
||
|
||
async def test_create_memory_bank(self):
|
||
"""测试创建记忆空间"""
|
||
memory_bank_id = await self.memory_bank.create_memory_bank("tieguaili")
|
||
|
||
# 验证返回的ID格式
|
||
self.assertEqual(memory_bank_id, "cf_memory_tieguaili")
|
||
|
||
async def test_create_memory_bank_with_display_name(self):
|
||
"""测试创建记忆空间时指定显示名称"""
|
||
memory_bank_id = await self.memory_bank.create_memory_bank(
|
||
"tieguaili",
|
||
"铁拐李的专属记忆银行"
|
||
)
|
||
|
||
# 验证返回的ID格式
|
||
self.assertEqual(memory_bank_id, "cf_memory_tieguaili")
|
||
|
||
async def test_generate_embedding(self):
|
||
"""测试生成嵌入向量"""
|
||
# Mock响应
|
||
mock_response = AsyncMock()
|
||
mock_response.status = 200
|
||
mock_response.json = AsyncMock(return_value={
|
||
"result": {
|
||
"data": [
|
||
{
|
||
"embedding": [0.1, 0.2, 0.3, 0.4, 0.5]
|
||
}
|
||
]
|
||
}
|
||
})
|
||
|
||
# Mock session.post
|
||
self.mock_session.post.return_value.__aenter__.return_value = mock_response
|
||
|
||
# 调用方法
|
||
embedding = await self.memory_bank._generate_embedding("测试文本")
|
||
|
||
# 验证结果
|
||
self.assertEqual(embedding, [0.1, 0.2, 0.3, 0.4, 0.5])
|
||
|
||
# 验证调用了正确的URL和参数
|
||
expected_url = "https://api.cloudflare.com/client/v4/accounts/test-account/ai/run/@cf/baai/bge-m3"
|
||
self.mock_session.post.assert_called_once()
|
||
call_args = self.mock_session.post.call_args
|
||
self.assertEqual(call_args[0][0], expected_url)
|
||
self.assertEqual(call_args[1]['json'], {"text": ["测试文本"]})
|
||
|
||
async def test_generate_embedding_api_error(self):
|
||
"""测试生成嵌入向量时API错误"""
|
||
# Mock响应
|
||
mock_response = AsyncMock()
|
||
mock_response.status = 500
|
||
mock_response.text = AsyncMock(return_value="Internal Server Error")
|
||
|
||
# Mock session.post
|
||
self.mock_session.post.return_value.__aenter__.return_value = mock_response
|
||
|
||
# 验证抛出异常
|
||
with self.assertRaises(Exception) as context:
|
||
await self.memory_bank._generate_embedding("测试文本")
|
||
|
||
self.assertIn("Failed to generate embedding", str(context.exception))
|
||
|
||
async def test_add_memory(self):
|
||
"""测试添加记忆"""
|
||
# Mock _generate_embedding 方法
|
||
with patch.object(self.memory_bank, '_generate_embedding', new=AsyncMock(return_value=[0.1, 0.2, 0.3])) as mock_embed:
|
||
# Mock upsert 响应
|
||
mock_response = AsyncMock()
|
||
mock_response.status = 200
|
||
mock_response.json = AsyncMock(return_value={"result": {"upserted": 1}})
|
||
|
||
# Mock session.post
|
||
self.mock_session.post.return_value.__aenter__.return_value = mock_response
|
||
|
||
# 添加记忆
|
||
memory_id = await self.memory_bank.add_memory(
|
||
agent_name="tieguaili",
|
||
content="在讨论NVIDIA股票时，我倾向于逆向思维，关注潜在风险。",
|
||
memory_type="preference",
|
||
debate_topic="NVIDIA投资分析",
|
||
metadata={"source": "manual"}
|
||
)
|
||
|
||
# 验证返回的ID格式 (以mem_开头)
|
||
self.assertTrue(memory_id.startswith("mem_tieguaili_"))
|
||
|
||
# 验证调用了生成嵌入的方法
|
||
mock_embed.assert_called_once_with("在讨论NVIDIA股票时，我倾向于逆向思维，关注潜在风险。")
|
||
|
||
# 验证调用了upsert API
|
||
self.mock_session.post.assert_called()
|
||
|
||
# 验证upsert调用的参数
|
||
upsert_call = None
|
||
for call in self.mock_session.post.call_args_list:
|
||
if 'vectorize/indexes/test-index/upsert' in call[0][0]:
|
||
upsert_call = call
|
||
break
|
||
|
||
self.assertIsNotNone(upsert_call)
|
||
call_args, call_kwargs = upsert_call
|
||
self.assertIn("vectorize/indexes/test-index/upsert", call_args[0])
|
||
self.assertIn("vectors", call_kwargs['json'])
|
||
|
||
async def test_add_memory_api_error(self):
|
||
"""测试添加记忆时API错误"""
|
||
# Mock _generate_embedding 方法
|
||
with patch.object(self.memory_bank, '_generate_embedding', new=AsyncMock(return_value=[0.1, 0.2, 0.3])):
|
||
# Mock upsert 响应
|
||
mock_response = AsyncMock()
|
||
mock_response.status = 500
|
||
mock_response.text = AsyncMock(return_value="Internal Server Error")
|
||
|
||
# Mock session.post
|
||
self.mock_session.post.return_value.__aenter__.return_value = mock_response
|
||
|
||
# 验证抛出异常
|
||
with self.assertRaises(Exception) as context:
|
||
await self.memory_bank.add_memory(
|
||
agent_name="tieguaili",
|
||
content="测试内容"
|
||
)
|
||
|
||
self.assertIn("Failed to upsert memory", str(context.exception))
|
||
|
||
async def test_search_memories(self):
|
||
"""测试搜索记忆"""
|
||
# Mock _generate_embedding 方法
|
||
with patch.object(self.memory_bank, '_generate_embedding', new=AsyncMock(return_value=[0.1, 0.2, 0.3])) as mock_embed:
|
||
# Mock query 响应
|
||
mock_response = AsyncMock()
|
||
mock_response.status = 200
|
||
mock_response.json = AsyncMock(return_value={
|
||
"result": {
|
||
"matches": [
|
||
{
|
||
"metadata": {
|
||
"content": "在讨论NVIDIA股票时，我倾向于逆向思维，关注潜在风险。",
|
||
"memory_type": "preference",
|
||
"agent_name": "tieguaili"
|
||
},
|
||
"score": 0.95
|
||
}
|
||
]
|
||
}
|
||
})
|
||
|
||
# Mock session.post
|
||
self.mock_session.post.return_value.__aenter__.return_value = mock_response
|
||
|
||
# 搜索记忆
|
||
results = await self.memory_bank.search_memories(
|
||
agent_name="tieguaili",
|
||
query="NVIDIA",
|
||
memory_type="preference",
|
||
limit=5
|
||
)
|
||
|
||
# 验证结果
|
||
self.assertEqual(len(results), 1)
|
||
self.assertEqual(results[0]["content"], "在讨论NVIDIA股票时，我倾向于逆向思维，关注潜在风险。")
|
||
self.assertEqual(results[0]["relevance_score"], 0.95)
|
||
|
||
# 验证调用了生成嵌入的方法
|
||
mock_embed.assert_called_once_with("NVIDIA")
|
||
|
||
# 验证调用了query API
|
||
self.mock_session.post.assert_called()
|
||
|
||
# 验证query调用的参数
|
||
query_call = None
|
||
for call in self.mock_session.post.call_args_list:
|
||
if 'vectorize/indexes/test-index/query' in call[0][0]:
|
||
query_call = call
|
||
break
|
||
|
||
self.assertIsNotNone(query_call)
|
||
call_args, call_kwargs = query_call
|
||
self.assertIn("vectorize/indexes/test-index/query", call_args[0])
|
||
self.assertIn("vector", call_kwargs['json'])
|
||
self.assertIn("filter", call_kwargs['json'])
|
||
self.assertEqual(call_kwargs['json']['filter'], {"agent_name": "tieguaili", "memory_type": "preference"})
|
||
|
||
async def test_search_memories_api_error(self):
|
||
"""测试搜索记忆时API错误"""
|
||
# Mock _generate_embedding 方法
|
||
with patch.object(self.memory_bank, '_generate_embedding', new=AsyncMock(return_value=[0.1, 0.2, 0.3])):
|
||
# Mock query 响应
|
||
mock_response = AsyncMock()
|
||
mock_response.status = 500
|
||
mock_response.text = AsyncMock(return_value="Internal Server Error")
|
||
|
||
# Mock session.post
|
||
self.mock_session.post.return_value.__aenter__.return_value = mock_response
|
||
|
||
# 验证返回空列表而不是抛出异常
|
||
results = await self.memory_bank.search_memories(
|
||
agent_name="tieguaili",
|
||
query="NVIDIA"
|
||
)
|
||
|
||
self.assertEqual(results, [])
|
||
|
||
async def test_get_agent_context(self):
|
||
"""测试获取智能体上下文"""
|
||
# Mock search_memories 方法
|
||
with patch.object(self.memory_bank, 'search_memories', new=AsyncMock()) as mock_search:
|
||
# 设置mock返回值
|
||
mock_search.side_effect = [
|
||
[ # conversation memories
|
||
{"content": "NVIDIA的估值过高，存在泡沫风险。", "relevance_score": 0.9}
|
||
],
|
||
[ # preference memories
|
||
{"content": "倾向于逆向思维，关注潜在风险。", "relevance_score": 0.8}
|
||
],
|
||
[ # strategy memories
|
||
{"content": "使用技术分析策略。", "relevance_score": 0.7}
|
||
]
|
||
]
|
||
|
||
# 获取上下文
|
||
context = await self.memory_bank.get_agent_context("tieguaili", "NVIDIA投资分析")
|
||
|
||
# 验证上下文包含预期内容
|
||
self.assertIn("# 铁拐李的记忆上下文", context)
|
||
self.assertIn("## 历史对话记忆", context)
|
||
self.assertIn("## 偏好记忆", context)
|
||
self.assertIn("## 策略记忆", context)
|
||
self.assertIn("NVIDIA的估值过高，存在泡沫风险。", context)
|
||
self.assertIn("倾向于逆向思维，关注潜在风险。", context)
|
||
self.assertIn("使用技术分析策略。", context)
|
||
|
||
async def test_get_agent_context_no_memories(self):
|
||
"""测试获取智能体上下文但无相关记忆"""
|
||
# Mock search_memories 方法
|
||
with patch.object(self.memory_bank, 'search_memories', new=AsyncMock(return_value=[])):
|
||
# 获取上下文
|
||
context = await self.memory_bank.get_agent_context("tieguaili", "NVIDIA投资分析")
|
||
|
||
# 验证上下文包含暂无相关记忆的提示
|
||
self.assertIn("# 铁拐李的记忆上下文", context)
|
||
self.assertIn("暂无相关记忆。", context)
|
||
|
||
async def test_save_debate_session(self):
|
||
"""测试保存辩论会话"""
|
||
# Mock add_memory 方法
|
||
with patch.object(self.memory_bank, 'add_memory', new=AsyncMock()) as mock_add:
|
||
conversation_history = [
|
||
{"agent": "tieguaili", "content": "NVIDIA的估值过高，存在泡沫风险。"},
|
||
{"agent": "lvdongbin", "content": "NVIDIA在AI领域的领先地位不可忽视。"},
|
||
{"agent": "tieguaili", "content": "但我们需要考虑竞争加剧和增长放缓的可能性。"}
|
||
]
|
||
|
||
outcomes = {
|
||
"winner": "lvdongbin",
|
||
"insights": {
|
||
"tieguaili": "铁拐李的风险意识值得肯定，但在AI趋势的判断上略显保守。"
|
||
}
|
||
}
|
||
|
||
# 保存辩论会话
|
||
await self.memory_bank.save_debate_session(
|
||
debate_topic="NVIDIA投资分析",
|
||
participants=["tieguaili", "lvdongbin"],
|
||
conversation_history=conversation_history,
|
||
outcomes=outcomes
|
||
)
|
||
|
||
# 验证调用了add_memory两次（对话总结和策略洞察）
|
||
self.assertEqual(mock_add.call_count, 2)
|
||
|
||
# 验证第一次调用是对话总结
|
||
call_args1 = mock_add.call_args_list[0][1]
|
||
self.assertEqual(call_args1['agent_name'], 'tieguaili')
|
||
self.assertEqual(call_args1['memory_type'], 'conversation')
|
||
self.assertIn('铁拐李在本次辩论中的主要观点', call_args1['content'])
|
||
|
||
# 验证第二次调用是策略洞察
|
||
call_args2 = mock_add.call_args_list[1][1]
|
||
self.assertEqual(call_args2['agent_name'], 'tieguaili')
|
||
self.assertEqual(call_args2['memory_type'], 'strategy')
|
||
self.assertIn('铁拐李的风险意识值得肯定', call_args2['content'])
|
||
|
||
def test_summarize_conversation(self):
|
||
"""测试对话总结"""
|
||
conversation_history = [
|
||
{"agent": "tieguaili", "content": "第一点看法：NVIDIA的估值过高，存在泡沫风险。"},
|
||
{"agent": "lvdongbin", "content": "NVIDIA在AI领域的领先地位不可忽视。"},
|
||
{"agent": "tieguaili", "content": "第二点看法：我们需要考虑竞争加剧和增长放缓的可能性。"},
|
||
{"agent": "tieguaili", "content": "第三点看法：从技术分析角度看，股价已出现超买信号。"}
|
||
]
|
||
|
||
summary = self.memory_bank._summarize_conversation(conversation_history, "tieguaili")
|
||
|
||
# 验证总结包含预期内容
|
||
self.assertIn("铁拐李在本次辩论中的主要观点", summary)
|
||
self.assertIn("第一点看法：NVIDIA的估值过高，存在泡沫风险。", summary)
|
||
self.assertIn("第二点看法：我们需要考虑竞争加剧和增长放缓的可能性。", summary)
|
||
self.assertIn("第三点看法：从技术分析角度看，股价已出现超买信号。", summary)
|
||
|
||
def test_extract_strategy_insight_winner(self):
|
||
"""测试提取策略洞察 - 获胜者"""
|
||
outcomes = {
|
||
"winner": "tieguaili",
|
||
"insights": {}
|
||
}
|
||
|
||
insight = self.memory_bank._extract_strategy_insight(outcomes, "tieguaili")
|
||
|
||
self.assertIn("铁拐李在本次辩论中获胜", insight)
|
||
|
||
def test_extract_strategy_insight_from_insights(self):
|
||
"""测试从洞察中提取策略洞察"""
|
||
outcomes = {
|
||
"winner": "lvdongbin",
|
||
"insights": {
|
||
"tieguaili": "铁拐李的风险意识值得肯定，但在AI趋势的判断上略显保守。"
|
||
}
|
||
}
|
||
|
||
insight = self.memory_bank._extract_strategy_insight(outcomes, "tieguaili")
|
||
|
||
self.assertEqual(insight, "铁拐李的风险意识值得肯定，但在AI趋势的判断上略显保守。")
|
||
|
||
|
||
if __name__ == '__main__':
|
||
# 创建一个异步测试运行器
|
||
def run_async_test(test_case):
|
||
"""运行异步测试用例"""
|
||
loop = asyncio.new_event_loop()
|
||
asyncio.set_event_loop(loop)
|
||
try:
|
||
return loop.run_until_complete(test_case)
|
||
finally:
|
||
loop.close()
|
||
|
||
# 获取所有以test_开头的异步方法并运行它们
|
||
suite = unittest.TestSuite()
|
||
test_instance = TestCloudflareMemoryBank()
|
||
test_instance.setUp()
|
||
test_instance.addCleanup(test_instance.tearDown)
|
||
|
||
# 添加同步测试
|
||
suite.addTest(TestCloudflareMemoryBank('test_init'))
|
||
suite.addTest(TestCloudflareMemoryBank('test_summarize_conversation'))
|
||
suite.addTest(TestCloudflareMemoryBank('test_extract_strategy_insight_winner'))
|
||
suite.addTest(TestCloudflareMemoryBank('test_extract_strategy_insight_from_insights'))
|
||
|
||
# 添加异步测试
|
||
async_tests = [
|
||
'test_create_memory_bank',
|
||
'test_create_memory_bank_with_display_name',
|
||
'test_generate_embedding',
|
||
'test_generate_embedding_api_error',
|
||
'test_add_memory',
|
||
'test_add_memory_api_error',
|
||
'test_search_memories',
|
||
'test_search_memories_api_error',
|
||
'test_get_agent_context',
|
||
'test_get_agent_context_no_memories',
|
||
'test_save_debate_session'
|
||
]
|
||
|
||
for test_name in async_tests:
|
||
test_method = getattr(test_instance, test_name)
|
||
suite.addTest(unittest.FunctionTestCase(lambda tm=test_method: run_async_test(tm())))
|
||
|
||
# 运行测试
|
||
runner = unittest.TextTestRunner(verbosity=2)
|
||
runner.run(suite) |