diff --git a/bbit_ai/app/config/milvus.py b/bbit_ai/app/config/milvus.py index 6785ccd..5219099 100644 --- a/bbit_ai/app/config/milvus.py +++ b/bbit_ai/app/config/milvus.py @@ -5,31 +5,44 @@ from utils.GlobalVariable import LOCAL_IP URI = "http://" + LOCAL_IP + ":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 -) +_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 diff --git a/bbit_ai/app/db/milvus.py b/bbit_ai/app/db/milvus.py index 70a67c4..7b26e9a 100644 --- a/bbit_ai/app/db/milvus.py +++ b/bbit_ai/app/db/milvus.py @@ -3,7 +3,7 @@ from typing import List from langchain_core.documents import Document -from config.milvus import knVectorstore, memVectorstore +from config.milvus import get_kn_vectorstore, get_mem_vectorstore def get_knowledge_by_key_words(key_words: str, kn_ids: List[str]) -> str: @@ -17,7 +17,7 @@ def get_knowledge_by_key_words(key_words: str, kn_ids: List[str]) -> str: else: return "未找到相关的知识。" - result = knVectorstore.similarity_search( + result = get_kn_vectorstore().similarity_search( query=key_words, k=3, expr=expr # 可调节返回条数 ) @@ -46,7 +46,7 @@ def get_memory_by_key_words(key_words: str, ai_ids: List[str]) -> str: else: expr = "" # 不限制 kn_id todo 实际上应该不反悔任何内容 - result = memVectorstore.similarity_search( + result = get_mem_vectorstore().similarity_search( query=key_words, k=5, expr=expr # 可调节返回条数 ) @@ -65,7 +65,7 @@ def get_memory_by_key_words(key_words: str, ai_ids: List[str]) -> str: def get_knowledge_by_base_id(base_id: str): expr = f'kn_id == "{base_id}"' # base_id 会被替换 - result = knVectorstore.similarity_search( + result = get_kn_vectorstore().similarity_search( query="", k=100, expr=expr # 如果只想用过滤条件,可以传空字符串 ) return [ @@ -90,7 +90,7 @@ def add_knowledge(text: str, is_active: bool, base_id: str, user_id: str): }, ) ] - return knVectorstore.add_documents(docs) + return get_kn_vectorstore().add_documents(docs) def add_memory(ai_id: str, mem: str, user_id: str, is_active: bool): @@ -105,4 +105,4 @@ def add_memory(ai_id: str, mem: str, user_id: str, is_active: bool): }, ) ] - return memVectorstore.add_documents(docs) + return get_mem_vectorstore().add_documents(docs)