You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
62 lines
2.3 KiB
62 lines
2.3 KiB
import asyncio
|
|
|
|
from pymilvus import FieldSchema, DataType, utility
|
|
|
|
from WxMini.Milvus.Utils.MilvusCollectionManager import MilvusCollectionManager
|
|
from WxMini.Milvus.Utils.MilvusConnectionPool import *
|
|
from WxMini.Utils.MySQLUtil import *
|
|
|
|
# 1. 使用连接池管理 Milvus 连接
|
|
milvus_pool = MilvusConnectionPool(host=MS_HOST, port=MS_PORT, max_connections=MS_MAX_CONNECTIONS)
|
|
|
|
# 2. 从连接池中获取一个连接
|
|
connection = milvus_pool.get_connection()
|
|
|
|
# 3. 初始化集合管理器
|
|
collection_name = MS_COLLECTION_NAME
|
|
collection_manager = MilvusCollectionManager(collection_name)
|
|
|
|
# 4. 判断集合是否存在，存在则删除
|
|
if utility.has_collection(collection_name):
|
|
print(f"集合 '{collection_name}' 已存在，正在删除...")
|
|
utility.drop_collection(collection_name)
|
|
print(f"集合 '{collection_name}' 已删除。")
|
|
|
|
# 5. 定义集合的字段和模式
|
|
fields = [
|
|
FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True), # 主键字段，自动生成 ID
|
|
FieldSchema(name="person_id", dtype=DataType.VARCHAR, max_length=64), # 会话 ID
|
|
FieldSchema(name="user_input", dtype=DataType.VARCHAR, max_length=2048), # 用户问题
|
|
FieldSchema(name="model_response", dtype=DataType.VARCHAR, max_length=2048), # 大模型反馈结果
|
|
FieldSchema(name="timestamp", dtype=DataType.VARCHAR, max_length=32), # 时间
|
|
FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=MS_DIMENSION) # 向量字段，维度为 200
|
|
]
|
|
schema_description = "Chat records collection with person_id , user input, model response, and timestamp"
|
|
|
|
# 6. 创建集合
|
|
print(f"正在创建集合 '{collection_name}'...")
|
|
collection_manager.create_collection(fields, schema_description)
|
|
print(f"集合 '{collection_name}' 创建成功。")
|
|
|
|
# 7. 释放连接
|
|
milvus_pool.release_connection(connection)
|
|
|
|
# 8. 关闭连接池
|
|
milvus_pool.close()
|
|
|
|
# 9、 清空MYSQL中的表
|
|
# 获取事件循环
|
|
loop = asyncio.get_event_loop()
|
|
|
|
# 初始化 MySQL 连接池
|
|
mysql_pool = loop.run_until_complete(init_mysql_pool())
|
|
print("MySQL 连接池已初始化。")
|
|
|
|
# 清空表
|
|
loop.run_until_complete(truncate_chat_log(mysql_pool))
|
|
|
|
# 关闭连接池
|
|
mysql_pool.close()
|
|
loop.run_until_complete(mysql_pool.wait_closed())
|
|
print("MySQL 连接池已关闭。")
|