'commit'

dsLightRag/Volcengine/chat.py

@@ -1,5 +1,10 @@
+import os
 import time
 
+from volcenginesdkarkruntime import Ark
+
+from Config.Config import VOLC_API_KEY
+
 """
 在记忆库准备好后，我们先模拟一段包含两轮的完整对话。
 对话结束后，把这段对话历史消息写入记忆库。然后再开启一个新话题，提出和刚才相关的问题，
@@ -7,7 +12,9 @@ AI 就能用刚写入的记忆来回答。
 注意：首次写入需要 3–5 分钟建立索引，这段时间内检索会报错。
 """
 import json
+import time
 from VikingDBMemoryService import initialize_services, ensure_collection_exists, search_relevant_memories
+
 def handle_conversation_turn(memory_service, llm_client, collection_name, user_id, user_message, conversation_history):
     """处理一轮对话，包括记忆搜索和LLM响应。"""
     print("\n" + "=" * 60)
@@ -71,15 +78,32 @@ def archive_conversation(memory_service, collection_name, user_id, assistant_id,
         print(f"归档对话失败: {e}")
         return False
 
+def setup_memory_collection(collection_name="emotional_support"):
+    """独立封装记忆体创建逻辑，返回memory_service供测试使用"""
+    try:
+        memory_service, _ = initialize_services()
+        ensure_collection_exists(memory_service, collection_name)
+        print(f"记忆体 '{collection_name}' 创建/验证成功")
+        return memory_service
+    except Exception as e:
+        print(f"记忆体创建失败: {e}")
+        return None
+
 def main():
     print("开始端到端记忆测试...")
+    collection_name="emotional_support"
 
     try:
-        memory_service, llm_client = initialize_services()
-        collection_name = "emotional_support"
-        user_id = "xiaoming" # 用户ID:小明
-        assistant_id = "assistant1" # 助手ID:助手1
-        ensure_collection_exists(memory_service, collection_name)
+        # 调用封装的记忆体创建函数
+        memory_service = setup_memory_collection()
+        if not memory_service:
+            return
+        llm_client = Ark(
+            base_url="https://ark.cn-beijing.volces.com/api/v3",
+            api_key=VOLC_API_KEY
+        )
+        user_id = "xiaoming"  # 用户ID:小明
+        assistant_id = "assistant1"  # 助手ID:助手1
     except Exception as e:
         print(f"初始化失败: {e}")
         return
@@ -114,4 +138,5 @@ def main():
     print("\n端到端记忆测试完成！")
 
 if __name__ == "__main__":
-    main()
+    setup_memory_collection()
+#    main()