'commit'

dsLightRag/Volcengine/T2_CreateIndex.py

@@ -1,14 +1,20 @@
 import json
 import logging
+import time
 
 from Config.Config import VOLC_ACCESSKEY, VOLC_SECRETKEY
-from VikingDBMemoryService import VikingDBMemoryService, MEMORY_COLLECTION_NAME
+from VikingDBMemoryService import VikingDBMemoryService, MEMORY_COLLECTION_NAME, VikingDBMemoryException
+
+
 # 控制日志输出
 logger = logging.getLogger('CollectionMemory')
 logger.setLevel(logging.INFO)
-handler = logging.StreamHandler()
-handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s'))
-logger.addHandler(handler)
+
+# 只添加一次处理器，避免重复日志
+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="情感陪伴记忆库"):
@@ -44,15 +50,44 @@ def create_memory_collection(collection_name, description="情感陪伴记忆库
         logger.info(f"创建响应: {json.dumps(response, ensure_ascii=False, indent=2)}")
         logger.info(f"集合 '{collection_name}' 创建成功")
 
-        # 等待集合就绪
+        # 等待集合就绪 - 修改为模拟chat.py的索引就绪检查机制
         logger.info("等待集合初始化完成...")
-        # 将独立函数调用改为实例方法调用
-        if memory_service.wait_for_collection_ready():
-            logger.info(f"集合 '{collection_name}' 已就绪，可以开始使用")
-            return True
-        else:
-            logger.info(f"集合 '{collection_name}' 初始化超时")
-            return False
+        max_retries = 30  # 最多重试30次
+        retry_interval = 10  # 每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"]}
+                memory_service.search_memory(
+                    collection_name=collection_name,
+                    query="test",
+                    filter=filter_params,
+                    limit=1
+                )
+                # 如果没有抛出索引错误，则认为就绪
+                logger.info(f"集合 '{collection_name}' 索引构建完成，已就绪")
+                return True
+            except VikingDBMemoryException as e:
+                error_msg = str(e)
+                # 检查是否是索引未就绪相关错误
+                if "index not exist" in error_msg or "need to add messages" in error_msg:
+                    retry_count += 1
+                    logger.info(f"索引尚未就绪，等待中... (重试 {retry_count}/{max_retries})")
+                    time.sleep(retry_interval)
+                else:
+                    logger.error(f"检查索引就绪状态时发生错误: {str(e)}")
+                    return False
+            except Exception as e:
+                logger.error(f"检查过程发生意外错误: {str(e)}")
+                return False
+        logger.error(f"集合 '{collection_name}' 索引构建超时")
+        return False
 
     except Exception as e:
         logger.error(f"创建集合失败: {str(e)}")