#!/usr/bin/env python3
"""
测试 Vertex AI Memory Bank 功能
"""

import asyncio
import sys
import os

# 添加项目根目录到Python路径
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '.')))

from src.jixia.memory.factory import get_memory_backend

async def test_vertex_memory_bank():
    """测试 Vertex Memory Bank 功能"""
    print("🧪 Vertex AI Memory Bank 功能测试\n")
    
    try:
        # 获取 Vertex Memory Bank 后端
        print("🔍 正在获取 Vertex Memory Bank 后端...")
        memory_bank = get_memory_backend(prefer='vertex')
        print("✅ 成功获取 Vertex Memory Bank 后端\n")
        
        # 测试创建记忆银行
        print("🔍 正在为吕洞宾创建记忆银行...")
        bank_id = await memory_bank.create_memory_bank("lvdongbin", "吕洞宾的记忆银行")
        print(f"✅ 成功为吕洞宾创建记忆银行: {bank_id}\n")
        
        # 测试添加记忆
        print("🔍 正在为吕洞宾添加记忆...")
        memory_id = await memory_bank.add_memory(
            agent_name="lvdongbin",
            content="在讨论NVIDIA股票时，我倾向于使用DCF模型评估其内在价值，并关注其在AI领域的竞争优势。",
            memory_type="preference",
            debate_topic="NVIDIA投资分析",
            metadata={"confidence": "high"}
        )
        print(f"✅ 成功为吕洞宾添加记忆: {memory_id}\n")
        
        # 测试搜索记忆
        print("🔍 正在搜索吕洞宾关于NVIDIA的记忆...")
        results = await memory_bank.search_memories(
            agent_name="lvdongbin",
            query="NVIDIA",
            memory_type="preference"
        )
        print(f"✅ 搜索完成，找到 {len(results)} 条相关记忆\n")
        
        if results:
            print("🔍 搜索结果:")
            for i, result in enumerate(results, 1):
                print(f"  {i}. {result['content']}")
                print(f"     相关性评分: {result['relevance_score']:.4f}\n")
        
        # 测试获取上下文
        print("🔍 正在获取吕洞宾关于NVIDIA投资分析的上下文...")
        context = await memory_bank.get_agent_context("lvdongbin", "NVIDIA投资分析")
        print("✅ 成功获取上下文\n")
        print("🔍 上下文内容:")
        print(context)
        print("\n")
        
        print("🎉 Vertex AI Memory Bank 功能测试完成!")
        
    except Exception as e:
        print(f"❌ 测试过程中发生错误: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    asyncio.run(test_vertex_memory_bank())