dsProject/dsLightRag/Volcengine/T2_CreateIndex.py

import json
import time
import logging
import sys
from Config.Config import VOLC_ACCESSKEY, VOLC_SECRETKEY
from VikingDBMemoryService import VikingDBMemoryService, VikingDBMemoryException

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger('CollectionMemory')

def main():
    logger.info("开始创建并等待集合就绪...")
    collection_name = "dsideal_collection"
    user_id = "system"
    assistant_id = "system"

    try:
        # 初始化服务
        logger.info("初始化记忆库服务...")
        memory_service = VikingDBMemoryService(ak=VOLC_ACCESSKEY, sk=VOLC_SECRETKEY)

        # 检查集合是否存在，不存在则创建
        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("添加测试数据以初始化索引...")
        session_id = f"init_session_{int(time.time())}"
        test_messages = [{
            "role": "user",
            "content": "初始化测试消息",
            "timestamp": int(time.time() * 1000)
        }]
        test_metadata = {
            "default_user_id": user_id,
            "default_assistant_id": assistant_id,
            "time": int(time.time() * 1000)
        }
        memory_service.add_session(
            collection_name=collection_name,
            session_id=session_id,
            messages=test_messages,
            metadata=test_metadata
        )

        # 等待索引就绪
        logger.info("开始等待索引构建完成...")
        max_retries = 30
        retry_interval = 10
        retry_count = 0

        while retry_count < max_retries:
            try:
                filter_params = {"user_id": [user_id], "memory_type": ["sys_event_v1"]}
                memory_service.search_memory(
                    collection_name=collection_name,
                    query="测试",
                    filter=filter_params,
                    limit=1
                )
                logger.info(f"集合 '{collection_name}' 索引构建完成，已就绪")
                sys.exit(0)
            except VikingDBMemoryException as e:
                error_msg = str(e)
                # 修复点：正确处理"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
                    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)
            except Exception as e:
                # 修复点：捕获所有异常，包括非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)}")
        sys.exit(1)

if __name__ == "__main__":
    main()