'commit'

dsLightRag/Volcengine/DropIndex.py

@@ -1,56 +0,0 @@
-import json
-import requests
-from volcengine.base.Request import Request
-from volcengine.Credentials import Credentials
-from volcengine.auth.SignerV4 import SignerV4
-
-from Config.Config import VOLC_ACCESSKEY, VOLC_SECRETKEY
-
-AK = VOLC_ACCESSKEY
-SK = VOLC_SECRETKEY
-Domain = "api-knowledgebase.mlp.cn-beijing.volces.com"
-
-def prepare_request(method, path, ak, sk, data=None):
-    r = Request()
-    r.set_shema("http")   # 注意这里是 http
-    r.set_method(method)
-    r.set_host(Domain)
-    r.set_path(path)
-
-    if data is not None:
-        r.set_body(json.dumps(data))
-
-    # 使用 air 服务和 cn-north-1 区域
-    credentials = Credentials(ak, sk, 'air', 'cn-north-1')
-    SignerV4.sign(r, credentials)
-    return r
-
-def internal_request(method, api, payload, params=None):
-    req = prepare_request(
-        method=method,
-        path=api,
-        ak=AK,
-        sk=SK,
-        data=payload
-    )
-
-    r = requests.request(
-        method=req.method,
-        url="{}://{}{}".format(req.schema, req.host, req.path),
-        headers=req.headers,
-        data=req.body,
-        params=params,
-    )
-    return r
-
-# 删除记忆库
-path = "/api/memory/collection/delete"
-payload = {
-    "CollectionName": "emotional_support"
-}
-rsp = internal_request("POST", path, payload)
-print(rsp.json())
-if rsp.status_code == 200:
-    print("记忆库删除成功:", rsp.json())
-else:
-    print("错误:", rsp.status_code, rsp.text)

dsLightRag/Volcengine/T1_DropIndex.py

@@ -0,0 +1,38 @@
+import json
+from VikingDBMemoryService import VikingDBMemoryService
+from Config.Config import VOLC_ACCESSKEY, VOLC_SECRETKEY
+
+def drop_existing_collection(collection_name):
+    # 初始化记忆库服务
+    memory_service = VikingDBMemoryService(
+        ak=VOLC_ACCESSKEY,
+        sk=VOLC_SECRETKEY,
+        host="api-knowledgebase.mlp.cn-beijing.volces.com",
+        region="cn-beijing"
+    )
+
+    try:
+        # 检查集合是否存在
+        print(f"正在检查集合 '{collection_name}'...")
+        memory_service.get_collection(collection_name)
+        print(f"集合 '{collection_name}' 已存在，准备删除...")
+
+        # 删除集合
+        response = memory_service.drop_collection(collection_name)
+        print(f"删除响应: {json.dumps(response, ensure_ascii=False, indent=2)}")
+        print(f"集合 '{collection_name}' 删除成功")
+        return True
+
+    except Exception as e:
+        error_msg = str(e)
+        if "collection not exist" in error_msg:
+            print(f"集合 '{collection_name}' 不存在，无需删除")
+            return False
+        else:
+            print(f"操作失败: {error_msg}")
+            raise
+
+if __name__ == "__main__":
+    # 可修改为目标集合名称
+    TARGET_COLLECTION = "emotional_support"
+    drop_existing_collection(TARGET_COLLECTION)

dsLightRag/Volcengine/T2_CreateIndex.py

@@ -0,0 +1,57 @@
+import json
+
+from Config.Config import VOLC_ACCESSKEY, VOLC_SECRETKEY
+from VikingDBMemoryService import VikingDBMemoryService
+from Volcengine.chat import wait_for_collection_ready
+
+
+def create_memory_collection(collection_name, description="情感陪伴记忆库"):
+    # 初始化记忆库服务
+    memory_service = VikingDBMemoryService(
+        ak=VOLC_ACCESSKEY,
+        sk=VOLC_SECRETKEY,
+        host="api-knowledgebase.mlp.cn-beijing.volces.com",
+        region="cn-beijing"
+    )
+
+    try:
+        # 检查集合是否已存在
+        print(f"正在检查集合 '{collection_name}'...")
+        memory_service.get_collection(collection_name)
+        print(f"集合 '{collection_name}' 已存在，无需重复创建")
+        return False
+
+    except Exception as e:
+        if "collection not exist" not in str(e):
+            print(f"检查集合时发生错误: {str(e)}")
+            raise
+
+    # 创建新集合
+    print(f"开始创建集合 '{collection_name}'...")
+    try:
+        response = memory_service.create_collection(
+            collection_name=collection_name,
+            description=description,
+            builtin_event_types=["sys_event_v1", "sys_profile_collect_v1"],
+            builtin_entity_types=["sys_profile_v1"]
+        )
+        print(f"创建响应: {json.dumps(response, ensure_ascii=False, indent=2)}")
+        print(f"集合 '{collection_name}' 创建成功")
+
+        # 等待集合就绪
+        print("等待集合初始化完成...")
+        if wait_for_collection_ready(memory_service, collection_name):
+            print(f"集合 '{collection_name}' 已就绪，可以开始使用")
+            return True
+        else:
+            print(f"集合 '{collection_name}' 初始化超时")
+            return False
+
+    except Exception as e:
+        print(f"创建集合失败: {str(e)}")
+        raise
+
+if __name__ == "__main__":
+    # 可修改为目标集合名称
+    TARGET_COLLECTION = "emotional_support"
+    create_memory_collection(TARGET_COLLECTION)

dsLightRag/Volcengine/VikingDBMemoryService.py

@@ -73,7 +73,7 @@ def search_relevant_memories(memory_service, collection_name, user_id, query, li
     while True:
         try:
             filter_params = {
-                "user_id": [user_id],
+                "user_id": user_id,  # 修正为字符串类型
                 "memory_type": ["sys_event_v1", "sys_profile_v1"]
             }
             response = memory_service.search_memory(

dsLightRag/Volcengine/chat.py

@@ -84,11 +84,44 @@ def setup_memory_collection(collection_name="emotional_support"):
         memory_service, _ = initialize_services()
         ensure_collection_exists(memory_service, collection_name)
         print(f"记忆体 '{collection_name}' 创建/验证成功")
-        return memory_service
+        
+        # 添加集合就绪等待
+        print("等待集合准备就绪...")
+        if wait_for_collection_ready(memory_service, collection_name):
+            print(f"集合 '{collection_name}' 已就绪")
+            return memory_service
+        else:
+            print(f"集合 '{collection_name}' 未能就绪")
+            return None
     except Exception as e:
         print(f"记忆体创建失败: {e}")
         return None
 
+def wait_for_collection_ready(memory_service, collection_name, timeout=300, interval=10):
+    """
+    等待集合准备就绪
+    :param memory_service: 记忆库服务实例
+    :param collection_name: 集合名称
+    :param timeout: 超时时间(秒)
+    :param interval: 检查间隔(秒)
+    :return: True if ready, False if timeout
+    """
+    start_time = time.time()
+    while time.time() - start_time < timeout:
+        try:
+            collection_info = memory_service.get_collection(collection_name)
+            # 根据Volcengine API文档，状态可能在Status字段中，值可能为"READY"、"CREATING"等
+            status = collection_info.get("Status", "UNKNOWN")
+            print(f"集合 '{collection_name}' 当前状态: {status}")
+            if status == "READY":
+                return True
+            time.sleep(interval)
+        except Exception as e:
+            print(f"检查集合状态失败: {e}")
+            time.sleep(interval)
+    print(f"集合 '{collection_name}' 在{timeout}秒内未就绪")
+    return False
+
 def main():
     print("开始端到端记忆测试...")
     collection_name="emotional_support"
@@ -139,4 +172,9 @@ def main():
 
 if __name__ == "__main__":
     setup_memory_collection()
-#    main()
+#    main()
+"""
+memory_service = setup_memory_collection()
+if memory_service:
+    is_ready = wait_for_collection_ready(memory_service, "emotional_support")
+"""