dsProject/dsLightRag/Volcengine/T3_ChatWithMemory.py

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"
        # 助手ID:助手
        assistant_id = "assistant"

        # 告知大模型用户信息
        logger.info("告知大模型用户信息...")
        user_info = "李明，男生，15岁，家住长春"

        # 记录信息到记忆体
        logger.info("记录用户信息到记忆体...")

        # 准备对话历史
        conversation_history = []

        # 使用正确的handle_conversation_turn方法参数
        response = handle_conversation_turn(
            memory_service=memory_service,  # 内存记忆服务
            llm_client=llm_client,  # 大模型客户端
            collection_name=collection_name,  # 集合名称
            user_id=user_id,  # 用户ID
            assistant_id=assistant_id,  # 助手ID
            user_message=f"请记住以下用户信息：{user_info}",  # 要记忆的信息
            conversation_history=conversation_history  # 对话历史
        )

        logger.info(f"模型回复: {response}")

        # 归档对话
        archive_conversation(
            memory_service=memory_service,  # 内存记忆服务
            collection_name=collection_name,  # 集合名称
            user_id=user_id,  # 用户ID
            assistant_id=assistant_id,  # 助手ID
            conversation_history=conversation_history,  # 对话历史
            topic_name="user_info"  # 会话主题
        )

        # 等待一段时间确保索引更新
        logger.info("等待索引更新...")
        time.sleep(5)

        # 验证大模型是否记住个人信息
        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, # 用户ID
            assistant_id=assistant_id,  # 助手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()