'commit'

dsLightRag/Volcengine/chat.py

@@ -1,3 +1,4 @@
+import logging
 import os
 import time
 
@@ -14,11 +15,16 @@ AI 就能用刚写入的记忆来回答。
 import json
 import time
 from VikingDBMemoryService import initialize_services, ensure_collection_exists, search_relevant_memories
-
+# 控制日志输出
+logger = logging.getLogger('CollectionMemory')
+logger.setLevel(logging.INFO)
+handler = logging.StreamHandler()
+handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s'))
+logger.addHandler(handler)
 def handle_conversation_turn(memory_service, llm_client, collection_name, user_id, user_message, conversation_history):
     """处理一轮对话，包括记忆搜索和LLM响应。"""
-    print("\n" + "=" * 60)
+    logger.info("\n" + "=" * 60)
-    print(f"用户: {user_message}")
+    logger.info(f"用户: {user_message}")
 
     relevant_memories = search_relevant_memories(memory_service, collection_name, user_id, user_message)
 
@@ -28,7 +34,7 @@ def handle_conversation_turn(memory_service, llm_client, collection_name, user_i
             [f"- {json.dumps(mem['memory_info'], ensure_ascii=False)}" for mem in relevant_memories])
         system_prompt += f"\n\n这是我们过去的一些对话记忆，请参考：\n{memory_context}\n\n请利用这些信息来更好地理解和回应用户。"
 
-    print("AI正在思考...")
+    logger.info("AI正在思考...")
 
     try:
         messages = [{"role": "system", "content": system_prompt}] + conversation_history + [
@@ -39,10 +45,10 @@ def handle_conversation_turn(memory_service, llm_client, collection_name, user_i
         )
         assistant_reply = completion.choices[0].message.content
     except Exception as e:
-        print(f"LLM调用失败: {e}")
+        logger.info(f"LLM调用失败: {e}")
         assistant_reply = "抱歉，我现在有点混乱，无法回应。我们可以稍后再聊吗？"
 
-    print(f"伙伴: {assistant_reply}")
+    logger.info(f"伙伴: {assistant_reply}")
 
     conversation_history.extend([
         {"role": "user", "content": user_message},
@@ -53,10 +59,10 @@ def handle_conversation_turn(memory_service, llm_client, collection_name, user_i
 def archive_conversation(memory_service, collection_name, user_id, assistant_id, conversation_history, topic_name):
     """将对话历史归档到记忆数据库。"""
     if not conversation_history:
-        print("没有对话可以归档。")
+        logger.info("没有对话可以归档。")
         return False
 
-    print(f"\n正在归档关于 '{topic_name}' 的对话...")
+    logger.info(f"\n正在归档关于 '{topic_name}' 的对话...")
     session_id = f"{topic_name}_{int(time.time())}"
     metadata = {
         "default_user_id": user_id,
@@ -71,11 +77,11 @@ def archive_conversation(memory_service, collection_name, user_id, assistant_id,
             messages=conversation_history,
             metadata=metadata
         )
-        print(f"对话已成功归档，会话ID: {session_id}")
+        logger.info(f"对话已成功归档，会话ID: {session_id}")
-        print("正在等待记忆索引更新...")
+        logger.info("正在等待记忆索引更新...")
         return True
     except Exception as e:
-        print(f"归档对话失败: {e}")
+        logger.info(f"归档对话失败: {e}")
         return False
 
 def setup_memory_collection(collection_name="emotional_support"):
@@ -83,18 +89,18 @@ def setup_memory_collection(collection_name="emotional_support"):
     try:
         memory_service, _ = initialize_services()
         ensure_collection_exists(memory_service, collection_name)
-        print(f"记忆体 '{collection_name}' 创建/验证成功")
+        logger.info(f"记忆体 '{collection_name}' 创建/验证成功")
         
         # 添加集合就绪等待
-        print("等待集合准备就绪...")
+        logger.info("等待集合准备就绪...")
         if wait_for_collection_ready(memory_service, collection_name):
-            print(f"集合 '{collection_name}' 已就绪")
+            logger.info(f"集合 '{collection_name}' 已就绪")
             return memory_service
         else:
-            print(f"集合 '{collection_name}' 未能就绪")
+            logger.info(f"集合 '{collection_name}' 未能就绪")
             return None
     except Exception as e:
-        print(f"记忆体创建失败: {e}")
+        logger.info(f"记忆体创建失败: {e}")
         return None
 
 def wait_for_collection_ready(memory_service, collection_name, timeout=300, interval=10):
@@ -112,18 +118,18 @@ def wait_for_collection_ready(memory_service, collection_name, timeout=300, inte
             collection_info = memory_service.get_collection(collection_name)
             # 根据Volcengine API文档，状态可能在Status字段中，值可能为"READY"、"CREATING"等
             status = collection_info.get("Status", "UNKNOWN")
-            print(f"集合 '{collection_name}' 当前状态: {status}")
+            logger.info(f"集合 '{collection_name}' 当前状态: {status}")
             if status == "READY":
                 return True
             time.sleep(interval)
         except Exception as e:
-            print(f"检查集合状态失败: {e}")
+            logger.info(f"检查集合状态失败: {e}")
             time.sleep(interval)
-    print(f"集合 '{collection_name}' 在{timeout}秒内未就绪")
+    logger.info(f"集合 '{collection_name}' 在{timeout}秒内未就绪")
     return False
 
 def main():
-    print("开始端到端记忆测试...")
+    logger.info("开始端到端记忆测试...")
     collection_name="emotional_support"
 
     try:
@@ -138,10 +144,10 @@ def main():
         user_id = "xiaoming"  # 用户ID:小明
         assistant_id = "assistant1"  # 助手ID:助手1
     except Exception as e:
-        print(f"初始化失败: {e}")
+        logger.info(f"初始化失败: {e}")
         return
 
-    print("\n--- 阶段 1: 初始对话 ---")
+    logger.info("\n--- 阶段 1: 初始对话 ---")
     initial_conversation_history = []
     handle_conversation_turn(
         memory_service, llm_client, collection_name, user_id,
@@ -154,13 +160,13 @@ def main():
         initial_conversation_history
     )
 
-    print("\n--- 阶段 2: 归档记忆 ---")
+    logger.info("\n--- 阶段 2: 归档记忆 ---")
     archive_conversation(
         memory_service, collection_name, user_id, assistant_id,
         initial_conversation_history, "study_stress_discussion"
     )
 
-    print("\n--- 阶段 3: 验证记忆 ---")
+    logger.info("\n--- 阶段 3: 验证记忆 ---")
     verification_conversation_history = []
     handle_conversation_turn(
         memory_service, llm_client, collection_name, user_id,
@@ -168,7 +174,7 @@ def main():
         verification_conversation_history
     )
 
-    print("\n端到端记忆测试完成！")
+    logger.info("\n端到端记忆测试完成！")
 
 if __name__ == "__main__":
     setup_memory_collection()