62 lines
2.3 KiB
Python
62 lines
2.3 KiB
Python
import json
|
|
import logging
|
|
|
|
from Config.Config import VOLC_ACCESSKEY, VOLC_SECRETKEY
|
|
from VikingDBMemoryService import VikingDBMemoryService, MEMORY_COLLECTION_NAME
|
|
from Volcengine.chat import wait_for_collection_ready
|
|
# 控制日志输出
|
|
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 wait_for_collection_ready(memory_service, collection_name):
|
|
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) |