修复因Milvus数据库初始化失败而导致整个服务无法启动的问题

bbit_ai/app/config/milvus.py

@@ -5,9 +5,21 @@ from utils.GlobalVariable import LOCAL_IP
 
 URI = "http://" + LOCAL_IP + ":19530"
 
-knVectorstore = Milvus(
+_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"},
+                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",
@@ -17,19 +29,20 @@ knVectorstore = Milvus(
                 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
-)
+                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

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)