import logging
import sys
import time

from Volcengine.Kit.VikingDBMemoryService import MEMORY_COLLECTION_NAME, initialize_services, \
    handle_conversation_turn, archive_conversation

# 控制日志输出
logger = logging.getLogger('ChatWithMemory')
logger.setLevel(logging.INFO)

# 只添加一次处理器，避免重复日志
if not logger.handlers:
    handler = logging.StreamHandler()
    handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s'))
    logger.addHandler(handler)



def main():
    logger.info("开始测试大模型记忆功能...")
    
    try:
        # 使用initialize_services函数初始化服务和LLM客户端
        memory_service, llm_client = initialize_services()
        
        collection_name = MEMORY_COLLECTION_NAME
        user_id = "liming" # 用户李明
        assistant_id = "assistant" # 助手ID:助手
        
        # 验证大模型是否记住个人信息
        logger.info("验证大模型是否记住个人信息...")
        test_conversation_history = []
        
        test_response = handle_conversation_turn(
            memory_service=memory_service,
            llm_client=llm_client,
            collection_name=collection_name,
            user_id=user_id,
            user_message="请告诉我李明的个人信息",
            conversation_history=test_conversation_history
        )
        
        logger.info(f"测试回复: {test_response}")
        
        # 检查回复中是否包含关键信息
        keywords = ["李明", "男", "15", "长春"]
        found_keywords = [kw for kw in keywords if kw in test_response]
        
        if len(found_keywords) == len(keywords):
            logger.info("✅ 大模型成功记住了用户信息！")
        else:
            logger.info(f"❌ 大模型可能没有完全记住用户信息。找到的关键词: {found_keywords}")
            
            # 尝试直接搜索记忆
            logger.info("尝试直接搜索记忆...")
            filter_params = {
                "user_id": [user_id],
                "assistant_id": assistant_id,  # 添加这一行
                "memory_type": ["sys_event_v1", "sys_profile_v1"]
            }
            search_result = memory_service.search_memory(
                collection_name=collection_name,
                query="李明 15岁 长春 男生",
                filter=filter_params,
                limit=5
            )
            
            logger.info(f"搜索结果: {search_result}")
            
    except Exception as e:
        logger.error(f"操作失败: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()