'commit'

2025-09-07 08:56:35 +08:00
parent bb94553638
commit 25eb08b1db
7 changed files with 626 additions and 20 deletions
--- a/dsLightRag/Volcengine/T3_Chat.py
+++ b/dsLightRag/Volcengine/T3_Chat.py
@@ -0,0 +1,107 @@
+import json
+import logging
+import time
+from Volcengine.VikingDBMemoryService import initialize_services, MEMORY_COLLECTION_NAME
+from volcenginesdkarkruntime import Ark
+from Config.Config import VOLC_API_KEY
+
+# 配置日志
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger('MemoryChatSimulator')
+
+def simulate_chat_session(memory_service, llm_client, user_id, session_topic):
+    """模拟完整对话会话：创建对话→存储记忆→搜索记忆"""
+    # 初始化对话历史
+    conversation_history = []
+    logger.info(f"\n=== 开始新对话会话: {session_topic} ===")
+
+    # 模拟两轮对话
+    user_messages = [
+        f"你好，我是{user_id}，{session_topic}",
+        "我最近感觉压力很大，不知道该怎么办"
+    ]
+
+    for message in user_messages:
+        # 处理对话轮次
+        assistant_reply = memory_service.handle_conversation_turn(
+            llm_client=llm_client,
+            user_id=user_id,
+            user_message=message,
+            conversation_history=conversation_history
+        )
+        time.sleep(1)  # 模拟思考延迟
+
+    # 归档对话记忆
+    logger.info("\n=== 归档对话记忆 ===")
+    archive_success = memory_service.archive_conversation(
+        user_id=user_id,
+        assistant_id="simulated_assistant",
+        conversation_history=conversation_history,
+        topic_name=session_topic
+    )
+
+    # 搜索相关记忆
+    if archive_success:
+        logger.info("\n=== 搜索相关记忆 ===")
+        # 等待索引更新
+        logger.info("等待记忆库索引更新...")
+        time.sleep(5)  # 实际环境可能需要更长时间
+
+        # 搜索与当前主题相关的记忆
+        search_query = f"{session_topic}相关的问题"
+        relevant_memories = memory_service.search_relevant_memories(
+            collection_name=MEMORY_COLLECTION_NAME,
+            user_id=user_id,
+            query=search_query,
+            limit=3
+        )
+
+        logger.info(f"\n搜索到 '{search_query}' 的相关记忆 ({len(relevant_memories)}条):")
+        for i, memory in enumerate(relevant_memories, 1):
+            logger.info(f"\n记忆 {i} (相关度: {memory['score']:.3f}):")
+            logger.info(f"内容: {memory['memory_info']}")
+
+    return archive_success
+
+def main():
+    """主函数：初始化服务并执行模拟测试"""
+    logger.info("===== 记忆库存储与搜索模拟测试 ====")
+
+    try:
+        # 初始化服务
+        logger.info("初始化记忆库服务和LLM客户端...")
+        memory_service, llm_client = initialize_services()
+
+        # 确保集合就绪
+        logger.info(f"检查集合 '{MEMORY_COLLECTION_NAME}' 是否就绪...")
+        
+        # 先验证集合是否存在
+        try:
+            collection_info = memory_service.get_collection(MEMORY_COLLECTION_NAME)
+            logger.info(f"集合 '{MEMORY_COLLECTION_NAME}' 存在，详细信息: {json.dumps(collection_info, ensure_ascii=False)}...")
+        except Exception as e:
+            logger.error(f"集合 '{MEMORY_COLLECTION_NAME}' 不存在: {str(e)}")
+            logger.error("请先运行T2_CreateIndex.py创建集合")
+            return
+        
+        if not memory_service.wait_for_collection_ready(timeout=120):
+            logger.error("集合未就绪，无法继续测试")
+            return
+
+        # 模拟两个不同用户的对话会话
+        user_sessions = [
+            ("student_001", "学习压力问题"),
+            ("teacher_001", "教学方法讨论")
+        ]
+
+        for user_id, topic in user_sessions:
+            simulate_chat_session(memory_service, llm_client, user_id, topic)
+            time.sleep(2)  # 会话间隔
+
+        logger.info("\n===== 模拟测试完成 ====")
+
+    except Exception as e:
+        logger.error(f"测试过程中发生错误: {str(e)}", exc_info=True)
+
+if __name__ == "__main__":
+    main()