from applications.config import LOCAL_MODEL_CONFIG

from pymilvus import Collection, CollectionSchema, FieldSchema, DataType


collections = {}
for model_name, cfg in LOCAL_MODEL_CONFIG.items():
    col_name = model_name.replace("/", "_").replace("-", "_").lower()
    fields = [
        FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
        FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=1024),
        FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=cfg["dim"])
    ]
    schema = CollectionSchema(fields, description=f"{model_name} embeddings")
    collection = Collection(col_name, schema=schema)
    collection.load()
    collections[model_name] = collection