'commit'

dsLightRag/Volcengine/T2_CreateIndex.py

@@ -1,75 +1,58 @@
 import json
-import logging
 import time
-
+import logging
+import sys
 from Config.Config import VOLC_ACCESSKEY, VOLC_SECRETKEY
-from VikingDBMemoryService import VikingDBMemoryService, MEMORY_COLLECTION_NAME, VikingDBMemoryException
+from VikingDBMemoryService import VikingDBMemoryService, VikingDBMemoryException
 
-
-# 控制日志输出
+# 配置日志
+logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger('CollectionMemory')
-logger.setLevel(logging.INFO)
 
-# 只添加一次处理器，避免重复日志
-if not logger.handlers:
-    handler = logging.StreamHandler()
-    handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s'))
-    logger.addHandler(handler)
-
-
-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"
-    )
+def main():
+    logger.info("开始创建并等待集合就绪...")
+    collection_name = "dsideal_collection"
+    user_id = "system"
+    assistant_id = "system"
 
     try:
-        # 检查集合是否已存在
-        logger.info(f"正在检查集合 '{collection_name}'...")
-        memory_service.get_collection(collection_name)
-        logger.info(f"集合 '{collection_name}' 已存在，无需重复创建")
-        return False
+        # 初始化服务
+        logger.info("初始化记忆库服务...")
+        memory_service = VikingDBMemoryService(ak=VOLC_ACCESSKEY, sk=VOLC_SECRETKEY)
 
-    except Exception as e:
-        if "collection not exist" not in str(e):
-            logger.info(f"检查集合时发生错误: {str(e)}")
-            raise
+        # 检查集合是否存在，不存在则创建
+        logger.info(f"检查集合 '{collection_name}' 是否存在...")
+        try:
+            memory_service.get_collection(collection_name)
+            logger.info(f"集合 '{collection_name}' 已存在")
+        except Exception as e:
+            error_message = str(e)
+            if "collection not exist" in error_message:
+                logger.info(f"集合 '{collection_name}' 不存在，正在创建...")
+                memory_service.create_collection(
+                    collection_name=collection_name,
+                    description="中文情感陪伴场景测试",
+                    builtin_event_types=["sys_event_v1", "sys_profile_collect_v1"],
+                    builtin_entity_types=["sys_profile_v1"]
+                )
+                logger.info(f"集合 '{collection_name}' 创建成功")
+            else:
+                logger.error(f"检查集合时出错: {e}")
+                sys.exit(1)
 
-    # 创建新集合
-    logger.info(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"]
-        )
-        logger.info(f"创建响应: {json.dumps(response, ensure_ascii=False, indent=2)}")
-        logger.info(f"集合 '{collection_name}' 创建成功")
-
-        # 添加测试消息，触发索引构建
-        logger.info("添加测试消息以初始化索引...")
-        # 生成唯一会话ID
-        session_id = f"test_session_{int(time.time())}"
-        # 构造测试消息（符合sys_event_v1类型要求）
+        # 添加测试数据
+        logger.info("添加测试数据以初始化索引...")
+        session_id = f"init_session_{int(time.time())}"
         test_messages = [{
             "role": "user",
-            "content": "initial test message",
-            "memory_type": "sys_event_v1",
-            "user_id": "test_user"
+            "content": "初始化测试消息",
+            "timestamp": int(time.time() * 1000)
         }]
-        # 构造元数据（包含API要求的必填字段）
         test_metadata = {
-            "default_user_id": "test_user",  # API必填字段
-            "default_assistant_id": "test_assistant",  # API必填字段
-            "user_id": "test_user",
-            "session_id": session_id,
-            "time": int(time.time() * 1000)  # 修正字段名并使用毫秒级时间戳
+            "default_user_id": user_id,
+            "default_assistant_id": assistant_id,
+            "time": int(time.time() * 1000)
         }
-        # 使用正确的add_session方法添加会话
         memory_service.add_session(
             collection_name=collection_name,
             session_id=session_id,
@@ -77,51 +60,52 @@ def create_memory_collection(collection_name, description="情感陪伴记忆库
             metadata=test_metadata
         )
 
-        # 等待集合就绪 - 修改为模拟chat.py的索引就绪检查机制
-        logger.info("等待集合初始化完成...")
-        max_retries = 30  # 最多重试30次
-        retry_interval = 10  # 每10秒重试一次
+        # 等待索引就绪
+        logger.info("开始等待索引构建完成...")
+        max_retries = 30
+        retry_interval = 10
         retry_count = 0
-        
-        # 增加初始延迟，避免创建后立即检查
-        logger.info(f"初始延迟30秒，等待索引构建...")
-        time.sleep(30)
-        
+
         while retry_count < max_retries:
             try:
-                # 尝试执行需要索引的操作（模拟chat.py中的搜索逻辑）
-                filter_params = {
-                    "memory_type": ["sys_event_v1"],
-                    "user_id": "test_user"  # 添加测试用户ID满足API要求
-                }
+                filter_params = {"user_id": [user_id], "memory_type": ["sys_event_v1"]}
                 memory_service.search_memory(
                     collection_name=collection_name,
-                    query="test",
+                    query="测试",
                     filter=filter_params,
                     limit=1
                 )
-                # 如果没有抛出索引错误，则认为就绪
                 logger.info(f"集合 '{collection_name}' 索引构建完成，已就绪")
-                return True
+                sys.exit(0)
             except VikingDBMemoryException as e:
                 error_msg = str(e)
-                # 检查是否是索引未就绪相关错误
-                if "index not exist" in error_msg or "need to add messages" in error_msg:
+                # 修复点：正确处理"index not ready"错误，进行重试
+                if "index not ready" in error_msg or "index not exist" in error_msg or "need to add messages" in error_msg:
                     retry_count += 1
-                    logger.info(f"索引尚未就绪，等待中... (重试 {retry_count}/{max_retries})")
+                    remaining = max_retries - retry_count
+                    logger.info(f"索引尚未就绪，将重试({retry_count}/{max_retries})，剩余{remaining}次...")
                     time.sleep(retry_interval)
                 else:
-                    logger.error(f"检查索引就绪状态时发生错误: {str(e)}")
-                    return False
+                    logger.error(f"检查索引状态时发生错误: {str(e)}")
+                    sys.exit(1)
             except Exception as e:
-                logger.error(f"检查过程发生意外错误: {str(e)}")
-                return False
-        logger.error(f"集合 '{collection_name}' 索引构建超时")
-        return False
+                # 修复点：捕获所有异常，包括非VikingDBMemoryException异常
+                error_msg = str(e)
+                if "index not ready" in error_msg or "index not exist" in error_msg or "need to add messages" in error_msg:
+                    retry_count += 1
+                    remaining = max_retries - retry_count
+                    logger.info(f"索引尚未就绪，将重试({retry_count}/{max_retries})，剩余{remaining}次...")
+                    time.sleep(retry_interval)
+                else:
+                    logger.error(f"检查索引状态时发生未知错误: {str(e)}")
+                    sys.exit(1)
+
+        logger.error(f"达到最大重试次数({max_retries})，索引仍未就绪")
+        sys.exit(1)
 
     except Exception as e:
-        logger.error(f"创建集合失败: {str(e)}")
-        raise
+        logger.error(f"系统异常: {str(e)}")
+        sys.exit(1)
 
 if __name__ == "__main__":
-    create_memory_collection(MEMORY_COLLECTION_NAME)
+    main()