'commit'

dsLightRag/Volcengine/VikingDBMemoryService.py

@@ -15,17 +15,19 @@ from volcengine.auth.SignerV4 import SignerV4
 from volcengine.base.Service import Service
 from volcenginesdkarkruntime import Ark
 
+from Config.Config import VOLC_SECRETKEY, VOLC_ACCESSKEY, VOLC_API_KEY
+
 
 def initialize_services(ak=None, sk=None, ark_api_key=None):
     """初始化记忆数据库服务和LLM客户端"""
     load_dotenv()
     # 如果参数未提供，尝试从环境变量获取
     if not ak:
-        ak = os.getenv("VOLC_ACCESSKEY")
+        ak = VOLC_ACCESSKEY
     if not sk:
-        sk = os.getenv("VOLC_SECRETKEY")
+        sk = VOLC_SECRETKEY
     if not ark_api_key:
-        ark_api_key = os.getenv("VOLC_API_KEY")
+        ark_api_key = VOLC_API_KEY
 
     if not all([ak, sk, ark_api_key]):
         raise ValueError("必须提供 VOLC_ACCESSKEY, VOLC_SECRETKEY, 和 VOLC_API_KEY。")

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"
+        # 调用封装的记忆体创建函数
+        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
-        ensure_collection_exists(memory_service, collection_name)
     except Exception as e:
         print(f"初始化失败: {e}")
         return
@@ -114,4 +138,5 @@ def main():
     print("\n端到端记忆测试完成！")
 
 if __name__ == "__main__":
-    main()
+    setup_memory_collection()
+#    main()