rag_server/applications/utils/milvus/collection.py

from pymilvus import connections, CollectionSchema, Collection
from applications.utils.milvus.field import fields
from applications.config import MILVUS_CONFIG


connections.connect("default", **MILVUS_CONFIG)

schema = CollectionSchema(
    fields, description="Chunk multi-vector embeddings with metadata"
)
milvus_collection = Collection(name="chunk_multi_embeddings", schema=schema)

# create index
vector_index_params = {
    "index_type": "IVF_FLAT",
    "metric_type": "COSINE",
    "params": {"M": 16, "efConstruction": 200},
}

milvus_collection.create_index("vector_text", vector_index_params)
milvus_collection.create_index("vector_summary", vector_index_params)
milvus_collection.create_index("vector_questions", vector_index_params)

milvus_collection.load()

__all__ = ["milvus_collection"]