87 lines
3.4 KiB
Python
87 lines
3.4 KiB
Python
|
"""NL2SQL 元数据向量的 Milvus 数据库和集合初始化。"""
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from pymilvus import AsyncMilvusClient, DataType
|
||
|
|
|
||
|
|
from config.settings import settings
|
||
|
|
|
||
|
|
NL2SQL_DATABASE_NAME = settings.milvus.db_name or "mutual_fund"
|
||
|
|
NL2SQL_COLLECTION = "data_agent_chunks"
|
||
|
|
|
||
|
|
|
||
|
|
def build_nl2sql_schema():
|
||
|
|
schema = AsyncMilvusClient.create_schema(auto_id=False, enable_dynamic_field=False)
|
||
|
|
schema.add_field("id", DataType.VARCHAR, is_primary=True, max_length=128)
|
||
|
|
schema.add_field("vector", DataType.FLOAT_VECTOR, dim=settings.llm.embed_dimensions)
|
||
|
|
schema.add_field("chunk_type", DataType.VARCHAR, max_length=32)
|
||
|
|
schema.add_field("table_name", DataType.VARCHAR, max_length=128)
|
||
|
|
schema.add_field("field_name", DataType.VARCHAR, max_length=128)
|
||
|
|
schema.add_field("text", DataType.VARCHAR, max_length=65535)
|
||
|
|
schema.add_field("is_valid", DataType.BOOL)
|
||
|
|
schema.add_field("is_deprecated", DataType.BOOL)
|
||
|
|
schema.add_field("content_hash", DataType.VARCHAR, max_length=128)
|
||
|
|
schema.add_field("created_at", DataType.INT64)
|
||
|
|
# Few-shot 字段允许元数据 chunk 缺省，避免影响现有表/字段向量写入。
|
||
|
|
schema.add_field("query", DataType.VARCHAR, max_length=4096, nullable=True)
|
||
|
|
schema.add_field("correct_sql", DataType.VARCHAR, max_length=16384, nullable=True)
|
||
|
|
schema.add_field("explanation", DataType.VARCHAR, max_length=65535, nullable=True)
|
||
|
|
schema.add_field("case_id", DataType.VARCHAR, max_length=128, nullable=True)
|
||
|
|
return schema
|
||
|
|
|
||
|
|
|
||
|
|
def build_nl2sql_index_params():
|
||
|
|
params = AsyncMilvusClient.prepare_index_params()
|
||
|
|
params.add_index(
|
||
|
|
field_name="vector",
|
||
|
|
index_type="HNSW",
|
||
|
|
metric_type="COSINE",
|
||
|
|
params={"M": 16, "efConstruction": 200},
|
||
|
|
)
|
||
|
|
return params
|
||
|
|
|
||
|
|
|
||
|
|
async def ensure_nl2sql_database(milvus_client: AsyncMilvusClient | None = None) -> None:
|
||
|
|
if milvus_client is None:
|
||
|
|
from config.database.milvus import client as configured_milvus_client
|
||
|
|
|
||
|
|
milvus_client = configured_milvus_client()
|
||
|
|
client = milvus_client
|
||
|
|
databases = await client.list_databases()
|
||
|
|
if NL2SQL_DATABASE_NAME not in databases:
|
||
|
|
await client.create_database(NL2SQL_DATABASE_NAME)
|
||
|
|
|
||
|
|
|
||
|
|
async def ensure_nl2sql_collection(milvus_client: AsyncMilvusClient | None = None) -> None:
|
||
|
|
if milvus_client is None:
|
||
|
|
from config.database.milvus import client as configured_milvus_client
|
||
|
|
|
||
|
|
milvus_client = configured_milvus_client()
|
||
|
|
client = milvus_client
|
||
|
|
await ensure_nl2sql_database(client)
|
||
|
|
if await client.has_collection(NL2SQL_COLLECTION):
|
||
|
|
return
|
||
|
|
await client.create_collection(
|
||
|
|
collection_name=NL2SQL_COLLECTION,
|
||
|
|
schema=build_nl2sql_schema(),
|
||
|
|
index_params=build_nl2sql_index_params(),
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
async def recreate_nl2sql_collection(
|
||
|
|
milvus_client: AsyncMilvusClient | None = None,
|
||
|
|
) -> None:
|
||
|
|
"""删除并按当前 Embedding 维度重新创建 NL2SQL 集合。"""
|
||
|
|
if milvus_client is None:
|
||
|
|
from config.database.milvus import client as configured_milvus_client
|
||
|
|
|
||
|
|
milvus_client = configured_milvus_client()
|
||
|
|
client = milvus_client
|
||
|
|
await ensure_nl2sql_database(client)
|
||
|
|
if await client.has_collection(NL2SQL_COLLECTION):
|
||
|
|
await client.drop_collection(collection_name=NL2SQL_COLLECTION)
|
||
|
|
await client.create_collection(
|
||
|
|
collection_name=NL2SQL_COLLECTION,
|
||
|
|
schema=build_nl2sql_schema(),
|
||
|
|
index_params=build_nl2sql_index_params(),
|
||
|
|
)
|