257 lines
8.8 KiB
Python
257 lines
8.8 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
Vertex AI Memory Bank 测试脚本
|
||
验证稷下学宫记忆银行功能
|
||
"""
|
||
|
||
import asyncio
|
||
import sys
|
||
import os
|
||
|
||
# 添加项目根目录到路径
|
||
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||
|
||
from src.jixia.memory.vertex_memory_bank import VertexMemoryBank, initialize_baxian_memory_banks
|
||
from src.jixia.agents.memory_enhanced_agent import MemoryEnhancedAgent, create_memory_enhanced_council
|
||
from config.doppler_config import get_google_genai_config, validate_config
|
||
|
||
|
||
async def test_memory_bank_basic():
|
||
"""测试Memory Bank基础功能"""
|
||
print("🧪 测试 Memory Bank 基础功能...")
|
||
|
||
try:
|
||
# 验证配置
|
||
if not validate_config("google_adk"):
|
||
print("❌ Google ADK 配置验证失败")
|
||
return False
|
||
|
||
config = get_google_genai_config()
|
||
if not config.get('project_id'):
|
||
print("❌ Google Cloud Project ID 未配置")
|
||
print("请设置环境变量: GOOGLE_CLOUD_PROJECT_ID")
|
||
return False
|
||
|
||
# 创建Memory Bank实例
|
||
memory_bank = VertexMemoryBank.from_config()
|
||
print(f"✅ Memory Bank 实例创建成功")
|
||
print(f" 项目ID: {config['project_id']}")
|
||
print(f" 区域: {config['location']}")
|
||
|
||
# 测试创建记忆银行
|
||
bank_id = await memory_bank.create_memory_bank(
|
||
agent_name="tieguaili",
|
||
display_name="铁拐李测试记忆银行"
|
||
)
|
||
print(f"✅ 创建记忆银行成功: {bank_id}")
|
||
|
||
# 测试添加记忆
|
||
memory_id = await memory_bank.add_memory(
|
||
agent_name="tieguaili",
|
||
content="测试记忆：在分析NVIDIA时，我倾向于关注潜在的市场风险和估值泡沫。",
|
||
memory_type="preference",
|
||
debate_topic="NVIDIA投资分析",
|
||
metadata={"test": True, "priority": "high"}
|
||
)
|
||
print(f"✅ 添加记忆成功: {memory_id}")
|
||
|
||
# 测试搜索记忆
|
||
results = await memory_bank.search_memories(
|
||
agent_name="tieguaili",
|
||
query="NVIDIA 风险",
|
||
limit=5
|
||
)
|
||
print(f"✅ 搜索记忆成功，找到 {len(results)} 条结果")
|
||
|
||
for i, result in enumerate(results, 1):
|
||
print(f" {i}. {result['content'][:50]}... (相关度: {result.get('relevance_score', 'N/A')})")
|
||
|
||
# 测试获取上下文
|
||
context = await memory_bank.get_agent_context("tieguaili", "NVIDIA投资分析")
|
||
print(f"✅ 获取上下文成功，长度: {len(context)} 字符")
|
||
|
||
return True
|
||
|
||
except Exception as e:
|
||
print(f"❌ Memory Bank 基础测试失败: {e}")
|
||
return False
|
||
|
||
|
||
async def test_memory_enhanced_agent():
|
||
"""测试记忆增强智能体"""
|
||
print("\n🧪 测试记忆增强智能体...")
|
||
|
||
try:
|
||
# 创建记忆银行
|
||
memory_bank = VertexMemoryBank.from_config()
|
||
|
||
# 创建记忆增强智能体
|
||
agent = MemoryEnhancedAgent("tieguaili", memory_bank)
|
||
print(f"✅ 创建记忆增强智能体: {agent.personality.chinese_name}")
|
||
|
||
# 测试基于记忆的响应
|
||
response = await agent.respond_with_memory(
|
||
message="你对NVIDIA的最新财报有什么看法？",
|
||
topic="NVIDIA投资分析"
|
||
)
|
||
print(f"✅ 智能体响应成功")
|
||
print(f" 响应长度: {len(response)} 字符")
|
||
print(f" 响应预览: {response[:100]}...")
|
||
|
||
# 测试学习偏好
|
||
await agent.learn_preference(
|
||
preference="用户偏好保守的投资策略，关注风险控制",
|
||
topic="投资偏好"
|
||
)
|
||
print("✅ 学习用户偏好成功")
|
||
|
||
# 测试保存策略洞察
|
||
await agent.save_strategy_insight(
|
||
insight="在高估值环境下，应该更加关注基本面分析和风险管理",
|
||
topic="投资策略"
|
||
)
|
||
print("✅ 保存策略洞察成功")
|
||
|
||
return True
|
||
|
||
except Exception as e:
|
||
print(f"❌ 记忆增强智能体测试失败: {e}")
|
||
return False
|
||
|
||
|
||
async def test_baxian_memory_council():
|
||
"""测试八仙记忆议会"""
|
||
print("\n🧪 测试八仙记忆议会...")
|
||
|
||
try:
|
||
# 创建记忆增强议会
|
||
council = await create_memory_enhanced_council()
|
||
print(f"✅ 创建八仙记忆议会成功，智能体数量: {len(council.agents)}")
|
||
|
||
# 列出所有智能体
|
||
for agent_name, agent in council.agents.items():
|
||
print(f" - {agent.personality.chinese_name} ({agent_name})")
|
||
|
||
# 进行简短的记忆辩论测试
|
||
print("\n🏛️ 开始记忆增强辩论测试...")
|
||
result = await council.conduct_memory_debate(
|
||
topic="比特币投资价值分析",
|
||
participants=["tieguaili", "lvdongbin"], # 只选择两个智能体进行快速测试
|
||
rounds=1
|
||
)
|
||
|
||
print(f"✅ 辩论完成")
|
||
print(f" 主题: {result['topic']}")
|
||
print(f" 参与者: {len(result['participants'])} 位")
|
||
print(f" 发言次数: {result['total_exchanges']}")
|
||
|
||
# 显示辩论内容
|
||
for exchange in result['conversation_history']:
|
||
print(f" {exchange['chinese_name']}: {exchange['content'][:80]}...")
|
||
|
||
# 获取集体记忆摘要
|
||
summary = await council.get_collective_memory_summary("比特币投资价值分析")
|
||
print(f"\n📚 集体记忆摘要长度: {len(summary)} 字符")
|
||
|
||
return True
|
||
|
||
except Exception as e:
|
||
print(f"❌ 八仙记忆议会测试失败: {e}")
|
||
return False
|
||
|
||
|
||
async def test_memory_bank_initialization():
|
||
"""测试Memory Bank初始化"""
|
||
print("\n🧪 测试 Memory Bank 初始化...")
|
||
|
||
try:
|
||
config = get_google_genai_config()
|
||
project_id = config.get('project_id')
|
||
location = config.get('location', 'us-central1')
|
||
|
||
if not project_id:
|
||
print("❌ 项目ID未配置，跳过初始化测试")
|
||
return False
|
||
|
||
# 初始化所有八仙记忆银行
|
||
memory_bank = await initialize_baxian_memory_banks(project_id, location)
|
||
print(f"✅ 八仙记忆银行初始化成功")
|
||
print(f" 记忆银行数量: {len(memory_bank.memory_banks)}")
|
||
|
||
for agent_name, bank_name in memory_bank.memory_banks.items():
|
||
chinese_name = memory_bank.baxian_agents.get(agent_name, agent_name)
|
||
print(f" - {chinese_name}: {bank_name}")
|
||
|
||
return True
|
||
|
||
except Exception as e:
|
||
print(f"❌ Memory Bank 初始化测试失败: {e}")
|
||
return False
|
||
|
||
|
||
async def main():
|
||
"""主测试函数"""
|
||
print("🏛️ 稷下学宫 Vertex AI Memory Bank 测试")
|
||
print("=" * 50)
|
||
|
||
# 检查配置
|
||
print("🔧 检查配置...")
|
||
config = get_google_genai_config()
|
||
|
||
print(f"Google API Key: {'已配置' if config.get('api_key') else '未配置'}")
|
||
print(f"Project ID: {config.get('project_id', '未配置')}")
|
||
print(f"Location: {config.get('location', 'us-central1')}")
|
||
print(f"Memory Bank: {'启用' if config.get('memory_bank_enabled', 'TRUE') == 'TRUE' else '禁用'}")
|
||
|
||
if not config.get('project_id'):
|
||
print("\n❌ 测试需要 Google Cloud Project ID")
|
||
print("请设置环境变量: GOOGLE_CLOUD_PROJECT_ID=your-project-id")
|
||
return
|
||
|
||
# 运行测试
|
||
tests = [
|
||
("Memory Bank 基础功能", test_memory_bank_basic),
|
||
("记忆增强智能体", test_memory_enhanced_agent),
|
||
("八仙记忆议会", test_baxian_memory_council),
|
||
("Memory Bank 初始化", test_memory_bank_initialization)
|
||
]
|
||
|
||
results = []
|
||
|
||
for test_name, test_func in tests:
|
||
print(f"\n{'='*20}")
|
||
print(f"测试: {test_name}")
|
||
print(f"{'='*20}")
|
||
|
||
try:
|
||
result = await test_func()
|
||
results.append((test_name, result))
|
||
except Exception as e:
|
||
print(f"❌ 测试 {test_name} 出现异常: {e}")
|
||
results.append((test_name, False))
|
||
|
||
# 显示测试结果摘要
|
||
print(f"\n{'='*50}")
|
||
print("🏛️ 测试结果摘要")
|
||
print(f"{'='*50}")
|
||
|
||
passed = 0
|
||
total = len(results)
|
||
|
||
for test_name, result in results:
|
||
status = "✅ 通过" if result else "❌ 失败"
|
||
print(f"{status} {test_name}")
|
||
if result:
|
||
passed += 1
|
||
|
||
print(f"\n📊 总体结果: {passed}/{total} 测试通过")
|
||
|
||
if passed == total:
|
||
print("🎉 所有测试通过！Vertex AI Memory Bank 集成成功！")
|
||
else:
|
||
print("⚠️ 部分测试失败，请检查配置和网络连接")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
# 运行测试
|
||
asyncio.run(main()) |