from langchain_milvus import Milvus

from config.llm import llmEmbeddings
from utils.GlobalVariable import LOCAL_IP

URI = "http://" + LOCAL_IP + ":19530"

_knVectorstore = None
_memVectorstore = None


def get_kn_vectorstore():
    global _knVectorstore
    if _knVectorstore is None:
        try:
            _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,
            )
        except Exception as e:
            print(f"[Milvus Init Failed] knowledge: {e}")
            _knVectorstore = None
    return _knVectorstore


def get_mem_vectorstore():
    global _memVectorstore
    if _memVectorstore is None:
        try:
            _memVectorstore = Milvus(...)
        except Exception as e:
            print(f"[Milvus Init Failed] memory: {e}")
            _memVectorstore = None
    return _memVectorstore