dsProject/dsLightRag/Volcengine/T2_CreateIndex.py

import json
import logging

from Config.Config import VOLC_ACCESSKEY, VOLC_SECRETKEY
from VikingDBMemoryService import VikingDBMemoryService, MEMORY_COLLECTION_NAME
# 控制日志输出
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)


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:
        # 检查集合是否已存在
        logger.info(f"正在检查集合 '{collection_name}'...")
        memory_service.get_collection(collection_name)
        logger.info(f"集合 '{collection_name}' 已存在，无需重复创建")
        return False

    except Exception as e:
        if "collection not exist" not in str(e):
            logger.info(f"检查集合时发生错误: {str(e)}")
            raise

    # 创建新集合
    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("等待集合初始化完成...")
        # 将独立函数调用改为实例方法调用
        if memory_service.wait_for_collection_ready():
            logger.info(f"集合 '{collection_name}' 已就绪，可以开始使用")
            return True
        else:
            logger.info(f"集合 '{collection_name}' 初始化超时")
            return False

    except Exception as e:
        logger.error(f"创建集合失败: {str(e)}")
        raise

if __name__ == "__main__":
    create_memory_collection(MEMORY_COLLECTION_NAME)