from langchain_milvus import Milvus

from config.llm import llmEmbeddings

URI = "http://ce_milvus:19530"

knVectorstore = Milvus(
    embedding_function=llmEmbeddings,
    connection_args={"uri": URI, "token": "root:Milvus", "db_name": "bbit_ai_lab"},
    collection_name="knowledge",
    index_params={"index_type": "FLAT", "metric_type": "L2"},
    consistency_level="Strong",
    auto_id=True,
    primary_field="id",
    text_field="text",
    vector_field="vector",
    partition_key_field="kn_id",
    enable_dynamic_field=True,
    drop_old=False,  # set to True if seeking to drop the collection with that name if it exists
)
memVectorstore = Milvus(
    embedding_function=llmEmbeddings,
    connection_args={"uri": URI, "token": "root:Milvus", "db_name": "bbit_ai_lab"},
    collection_name="memory",
    index_params={"index_type": "FLAT", "metric_type": "L2"},
    consistency_level="Strong",
    auto_id=True,
    primary_field="id",
    text_field="text",
    vector_field="vector",
    partition_key_field="ai_id",
    enable_dynamic_field=True,
    drop_old=False,  # set to True if seeking to drop the collection with that name if it exists
)