'commit'

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)