47 lines
1.7 KiB
Python
47 lines
1.7 KiB
Python
"""Milvus collection schemas for the mutual_fund knowledge base."""
|
|
from __future__ import annotations
|
|
|
|
from pymilvus import AsyncMilvusClient, DataType
|
|
|
|
from config.database.milvus import client as configured_milvus_client
|
|
|
|
|
|
KNOWLEDGE_COLLECTIONS = ("fin_faq", "fin_fund_doc", "fin_policy")
|
|
EMBEDDING_DIMENSION = 768
|
|
|
|
|
|
def build_knowledge_schema():
|
|
schema = AsyncMilvusClient.create_schema(auto_id=False, enable_dynamic_field=False)
|
|
schema.add_field("chunk_id", DataType.VARCHAR, is_primary=True, max_length=128)
|
|
schema.add_field("doc_id", DataType.VARCHAR, max_length=128)
|
|
schema.add_field("title", DataType.VARCHAR, max_length=512)
|
|
schema.add_field("section_title", DataType.VARCHAR, max_length=1024)
|
|
schema.add_field("text", DataType.VARCHAR, max_length=65535)
|
|
schema.add_field("strategy", DataType.VARCHAR, max_length=32)
|
|
schema.add_field("vector", DataType.FLOAT_VECTOR, dim=EMBEDDING_DIMENSION)
|
|
return schema
|
|
|
|
|
|
def build_knowledge_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_collections(milvus_client: AsyncMilvusClient | None = None) -> None:
|
|
client = milvus_client or configured_milvus_client()
|
|
schema = build_knowledge_schema()
|
|
index_params = build_knowledge_index_params()
|
|
for collection_name in KNOWLEDGE_COLLECTIONS:
|
|
if not await client.has_collection(collection_name):
|
|
await client.create_collection(
|
|
collection_name=collection_name,
|
|
schema=schema,
|
|
index_params=index_params,
|
|
)
|