384 lines
16 KiB
Python
384 lines
16 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
Vertex 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.vertex_memory_bank import VertexMemoryBank, MemoryEntry
|
||
|
||
|
||
class TestVertexMemoryBank(unittest.TestCase):
|
||
"""测试VertexMemoryBank类"""
|
||
|
||
def setUp(self):
|
||
"""测试前的设置"""
|
||
# Mock掉aiplatform.init以避免实际初始化
|
||
patcher = patch('src.jixia.memory.vertex_memory_bank.aiplatform.init')
|
||
self.mock_init = patcher.start()
|
||
self.addCleanup(patcher.stop)
|
||
|
||
# 创建VertexMemoryBank实例
|
||
self.memory_bank = VertexMemoryBank(
|
||
project_id="test-project",
|
||
location="us-central1"
|
||
)
|
||
|
||
# 重置本地存储
|
||
self.memory_bank.local_memories = {}
|
||
self.memory_bank.memory_banks = {}
|
||
|
||
def test_init(self):
|
||
"""测试初始化"""
|
||
self.assertEqual(self.memory_bank.project_id, "test-project")
|
||
self.assertEqual(self.memory_bank.location, "us-central1")
|
||
self.assertEqual(self.memory_bank.local_memories, {})
|
||
self.assertEqual(self.memory_bank.memory_banks, {})
|
||
|
||
# 验证调用了aiplatform.init
|
||
self.mock_init.assert_called_once_with(project="test-project", location="us-central1")
|
||
|
||
def test_from_config(self):
|
||
"""测试从配置创建实例"""
|
||
with patch('src.jixia.memory.vertex_memory_bank.get_google_genai_config') as mock_config:
|
||
mock_config.return_value = {
|
||
'project_id': 'config-project',
|
||
'location': 'europe-west1'
|
||
}
|
||
|
||
memory_bank = VertexMemoryBank.from_config()
|
||
|
||
self.assertEqual(memory_bank.project_id, "config-project")
|
||
self.assertEqual(memory_bank.location, "europe-west1")
|
||
|
||
def test_from_config_missing_project_id(self):
|
||
"""测试从配置创建实例时缺少project_id"""
|
||
with patch('src.jixia.memory.vertex_memory_bank.get_google_genai_config') as mock_config:
|
||
mock_config.return_value = {
|
||
'project_id': None,
|
||
'location': 'europe-west1'
|
||
}
|
||
|
||
with self.assertRaises(ValueError) as context:
|
||
VertexMemoryBank.from_config()
|
||
|
||
self.assertIn("Google Cloud Project ID 未配置", str(context.exception))
|
||
|
||
async def test_create_memory_bank(self):
|
||
"""测试创建记忆银行"""
|
||
memory_bank_id = await self.memory_bank.create_memory_bank("tieguaili")
|
||
|
||
# 验证返回的ID格式
|
||
self.assertEqual(memory_bank_id, "memory_bank_tieguaili_test-project")
|
||
|
||
# 验证内部状态
|
||
self.assertIn("tieguaili", self.memory_bank.memory_banks)
|
||
self.assertEqual(self.memory_bank.memory_banks["tieguaili"], memory_bank_id)
|
||
self.assertIn("tieguaili", self.memory_bank.local_memories)
|
||
self.assertEqual(self.memory_bank.local_memories["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, "memory_bank_tieguaili_test-project")
|
||
|
||
async def test_add_memory(self):
|
||
"""测试添加记忆"""
|
||
# 先创建记忆银行
|
||
await self.memory_bank.create_memory_bank("tieguaili")
|
||
|
||
# 添加记忆
|
||
memory_id = await self.memory_bank.add_memory(
|
||
agent_name="tieguaili",
|
||
content="在讨论NVIDIA股票时，我倾向于逆向思维，关注潜在风险。",
|
||
memory_type="preference",
|
||
debate_topic="NVIDIA投资分析",
|
||
metadata={"source": "manual"}
|
||
)
|
||
|
||
# 验证返回的ID格式
|
||
self.assertEqual(memory_id, "memory_tieguaili_0")
|
||
|
||
# 验证记忆已存储
|
||
self.assertEqual(len(self.memory_bank.local_memories["tieguaili"]), 1)
|
||
|
||
stored_memory = self.memory_bank.local_memories["tieguaili"][0]
|
||
self.assertEqual(stored_memory["id"], memory_id)
|
||
self.assertEqual(stored_memory["content"], "在讨论NVIDIA股票时，我倾向于逆向思维，关注潜在风险。")
|
||
self.assertEqual(stored_memory["memory_type"], "preference")
|
||
self.assertEqual(stored_memory["debate_topic"], "NVIDIA投资分析")
|
||
self.assertIn("source", stored_memory["metadata"])
|
||
self.assertEqual(stored_memory["metadata"]["source"], "manual")
|
||
self.assertIn("agent_name", stored_memory["metadata"])
|
||
self.assertEqual(stored_memory["metadata"]["agent_name"], "tieguaili")
|
||
|
||
async def test_add_memory_creates_bank_if_not_exists(self):
|
||
"""测试添加记忆时自动创建记忆银行"""
|
||
# 不先创建记忆银行，直接添加记忆
|
||
memory_id = await self.memory_bank.add_memory(
|
||
agent_name="tieguaili",
|
||
content="测试内容"
|
||
)
|
||
|
||
# 验证记忆银行已被自动创建
|
||
self.assertIn("tieguaili", self.memory_bank.memory_banks)
|
||
self.assertIn("tieguaili", self.memory_bank.local_memories)
|
||
|
||
# 验证记忆已存储
|
||
self.assertEqual(len(self.memory_bank.local_memories["tieguaili"]), 1)
|
||
|
||
async def test_search_memories(self):
|
||
"""测试搜索记忆"""
|
||
# 先创建记忆银行并添加一些记忆
|
||
await self.memory_bank.create_memory_bank("tieguaili")
|
||
|
||
await self.memory_bank.add_memory(
|
||
agent_name="tieguaili",
|
||
content="在讨论NVIDIA股票时，我倾向于逆向思维，关注潜在风险。",
|
||
memory_type="preference",
|
||
debate_topic="NVIDIA投资分析"
|
||
)
|
||
|
||
await self.memory_bank.add_memory(
|
||
agent_name="tieguaili",
|
||
content="我喜欢关注苹果公司的创新产品发布会。",
|
||
memory_type="preference",
|
||
debate_topic="AAPL投资分析"
|
||
)
|
||
|
||
# 搜索NVIDIA相关记忆
|
||
results = await self.memory_bank.search_memories(
|
||
agent_name="tieguaili",
|
||
query="NVIDIA"
|
||
)
|
||
|
||
self.assertEqual(len(results), 1)
|
||
self.assertEqual(results[0]["content"], "在讨论NVIDIA股票时，我倾向于逆向思维，关注潜在风险。")
|
||
self.assertIn("relevance_score", results[0])
|
||
|
||
async def test_search_memories_with_type_filter(self):
|
||
"""测试带类型过滤的搜索记忆"""
|
||
# 先创建记忆银行并添加不同类型的记忆
|
||
await self.memory_bank.create_memory_bank("tieguaili")
|
||
|
||
await self.memory_bank.add_memory(
|
||
agent_name="tieguaili",
|
||
content="在讨论NVIDIA股票时，我倾向于逆向思维，关注潜在风险。",
|
||
memory_type="preference",
|
||
debate_topic="NVIDIA投资分析"
|
||
)
|
||
|
||
await self.memory_bank.add_memory(
|
||
agent_name="tieguaili",
|
||
content="在NVIDIA的辩论中，我使用了技术分析策略。",
|
||
memory_type="strategy",
|
||
debate_topic="NVIDIA投资分析"
|
||
)
|
||
|
||
# 搜索NVIDIA相关记忆，只返回preference类型
|
||
results = await self.memory_bank.search_memories(
|
||
agent_name="tieguaili",
|
||
query="NVIDIA",
|
||
memory_type="preference"
|
||
)
|
||
|
||
self.assertEqual(len(results), 1)
|
||
self.assertEqual(results[0]["metadata"]["memory_type"], "preference")
|
||
|
||
async def test_search_memories_no_results(self):
|
||
"""测试搜索无结果的情况"""
|
||
# 搜索不存在的记忆银行
|
||
results = await self.memory_bank.search_memories(
|
||
agent_name="nonexistent",
|
||
query="test"
|
||
)
|
||
|
||
self.assertEqual(results, [])
|
||
|
||
# 搜索空的记忆银行
|
||
await self.memory_bank.create_memory_bank("tieguaili")
|
||
results = await self.memory_bank.search_memories(
|
||
agent_name="tieguaili",
|
||
query="test"
|
||
)
|
||
|
||
self.assertEqual(results, [])
|
||
|
||
async def test_get_agent_context(self):
|
||
"""测试获取智能体上下文"""
|
||
# 先创建记忆银行并添加一些记忆
|
||
await self.memory_bank.create_memory_bank("tieguaili")
|
||
|
||
await self.memory_bank.add_memory(
|
||
agent_name="tieguaili",
|
||
content="在讨论NVIDIA股票时，我倾向于逆向思维，关注潜在风险。",
|
||
memory_type="preference",
|
||
debate_topic="NVIDIA投资分析"
|
||
)
|
||
|
||
await self.memory_bank.add_memory(
|
||
agent_name="tieguaili",
|
||
content="在NVIDIA的辩论中，我使用了技术分析策略。",
|
||
memory_type="strategy",
|
||
debate_topic="NVIDIA投资分析"
|
||
)
|
||
|
||
# 获取上下文
|
||
context = await self.memory_bank.get_agent_context("tieguaili", "NVIDIA投资分析")
|
||
|
||
# 验证上下文包含预期内容
|
||
self.assertIn("# 铁拐李的记忆上下文", context)
|
||
self.assertIn("## 偏好记忆", context)
|
||
self.assertIn("## 策略记忆", context)
|
||
self.assertIn("在讨论NVIDIA股票时，我倾向于逆向思维，关注潜在风险。", context)
|
||
self.assertIn("在NVIDIA的辩论中，我使用了技术分析策略。", context)
|
||
|
||
async def test_get_agent_context_no_memories(self):
|
||
"""测试获取智能体上下文但无相关记忆"""
|
||
# 先创建记忆银行
|
||
await self.memory_bank.create_memory_bank("tieguaili")
|
||
|
||
# 获取上下文
|
||
context = await self.memory_bank.get_agent_context("tieguaili", "NVIDIA投资分析")
|
||
|
||
# 验证上下文包含暂无相关记忆的提示
|
||
self.assertIn("# 铁拐李的记忆上下文", context)
|
||
self.assertIn("暂无相关记忆。", context)
|
||
|
||
async def test_save_debate_session(self):
|
||
"""测试保存辩论会话"""
|
||
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
|
||
)
|
||
|
||
# 验证铁拐李的记忆已保存
|
||
self.assertIn("tieguaili", self.memory_bank.local_memories)
|
||
self.assertEqual(len(self.memory_bank.local_memories["tieguaili"]), 2)
|
||
|
||
# 验证第一条记忆是对话总结
|
||
summary_memory = self.memory_bank.local_memories["tieguaili"][0]
|
||
self.assertEqual(summary_memory["memory_type"], "conversation")
|
||
self.assertIn("铁拐李在本次辩论中的主要观点", summary_memory["content"])
|
||
|
||
# 验证第二条记忆是策略洞察
|
||
strategy_memory = self.memory_bank.local_memories["tieguaili"][1]
|
||
self.assertEqual(strategy_memory["memory_type"], "strategy")
|
||
self.assertIn("铁拐李的风险意识值得肯定", strategy_memory["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 = TestVertexMemoryBank()
|
||
test_instance.setUp()
|
||
|
||
# 添加同步测试
|
||
suite.addTest(TestVertexMemoryBank('test_init'))
|
||
suite.addTest(TestVertexMemoryBank('test_from_config'))
|
||
suite.addTest(TestVertexMemoryBank('test_from_config_missing_project_id'))
|
||
suite.addTest(TestVertexMemoryBank('test_summarize_conversation'))
|
||
suite.addTest(TestVertexMemoryBank('test_extract_strategy_insight_winner'))
|
||
suite.addTest(TestVertexMemoryBank('test_extract_strategy_insight_from_insights'))
|
||
|
||
# 添加异步测试
|
||
async_tests = [
|
||
'test_create_memory_bank',
|
||
'test_create_memory_bank_with_display_name',
|
||
'test_add_memory',
|
||
'test_add_memory_creates_bank_if_not_exists',
|
||
'test_search_memories',
|
||
'test_search_memories_with_type_filter',
|
||
'test_search_memories_no_results',
|
||
'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) |