Mutual_Fund/rag/milvus_collections.py

"""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
from rag.embedding import EMBEDDING_DIMENSION


KNOWLEDGE_COLLECTIONS = ("fin_faq", "fin_fund_doc", "fin_policy")


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,
            )
            continue
        # 已存在的集合维度必须与当前 embedding 配置一致，否则入库/检索会在运行时失败
        desc = await client.describe_collection(collection_name)
        for field in desc.get("fields", []):
            if field.get("name") != "vector":
                continue
            dim = field.get("params", {}).get("dim")
            if dim is not None and int(dim) != EMBEDDING_DIMENSION:
                raise RuntimeError(
                    f"Milvus collection {collection_name!r} vector dim={dim}, "
                    f"expected {EMBEDDING_DIMENSION}; drop and recreate it"
                )