升级新库

bbit_ai/app/db/milvus.py

@@ -1,9 +1,10 @@
-from config.milvus import knVectorstore,memVectorstore
-from langchain.schema import Document
 from datetime import datetime
 from typing import List
 
-from typing import List, Dict, Any
+from langchain_core.documents import Document
+
+from config.milvus import knVectorstore, memVectorstore
+
 
 def get_knowledge_by_key_words(key_words: str, kn_ids: List[str]) -> str:
     """
@@ -15,13 +16,11 @@ def get_knowledge_by_key_words(key_words: str, kn_ids: List[str]) -> str:
         expr = f"({ids_expr})"
     else:
         return "未找到相关的知识。"
-    
+
     result = knVectorstore.similarity_search(
-        query=key_words,
-        k=3,  # 可调节返回条数
-        expr=expr
+        query=key_words, k=3, expr=expr  # 可调节返回条数
     )
-    
+
     # 整理成字符串
     doc_texts = []
     for idx, doc in enumerate(result, start=1):
@@ -29,14 +28,14 @@ def get_knowledge_by_key_words(key_words: str, kn_ids: List[str]) -> str:
         if text:
             # 可以加个编号，便于LLM区分
             doc_texts.append(f"[文档{idx}]: {text}")
-    
+
     # 拼成一个大字符串，用换行隔开
     combined_text = "\n\n".join(doc_texts)
     return combined_text
 
 
 def get_memory_by_key_words(key_words: str, ai_ids: List[str]) -> str:
-    print("ai_id是：" , ai_ids)
+    print("ai_id是：", ai_ids)
     """
     根据关键词和 ai_ids 列表，在知识库中检索相关内容，并返回整理后的文本字符串
     """
@@ -46,13 +45,11 @@ def get_memory_by_key_words(key_words: str, ai_ids: List[str]) -> str:
         expr = f"({ids_expr})"
     else:
         expr = ""  # 不限制 kn_id todo 实际上应该不反悔任何内容
-    
+
     result = memVectorstore.similarity_search(
-        query=key_words,
-        k=5,  # 可调节返回条数
-        expr=expr
+        query=key_words, k=5, expr=expr  # 可调节返回条数
     )
-    
+
     # 整理成字符串
     doc_texts = []
     for idx, doc in enumerate(result, start=1):
@@ -60,16 +57,16 @@ def get_memory_by_key_words(key_words: str, ai_ids: List[str]) -> str:
         if text:
             # 可以加个编号，便于LLM区分
             doc_texts.append(f"[记忆{idx}]: {text}")
-    
+
     # 拼成一个大字符串，用换行隔开
     combined_text = "\n\n".join(doc_texts)
     return combined_text
+
+
 def get_knowledge_by_base_id(base_id: str):
     expr = f'kn_id == "{base_id}"'  # base_id 会被替换
-    result =  knVectorstore.similarity_search(
-        query="",  # 如果只想用过滤条件，可以传空字符串
-        k=100,
-        expr=expr
+    result = knVectorstore.similarity_search(
+        query="", k=100, expr=expr  # 如果只想用过滤条件，可以传空字符串
     )
     return [
         {
@@ -80,6 +77,7 @@ def get_knowledge_by_base_id(base_id: str):
         for doc in result
     ]
 
+
 def add_knowledge(text: str, is_active: bool, base_id: str, user_id: str):
     docs = [
         Document(
@@ -89,12 +87,13 @@ def add_knowledge(text: str, is_active: bool, base_id: str, user_id: str):
                 "created_by": str(user_id),
                 "created_at": datetime.now().isoformat(),
                 "is_active": is_active,
-            }
+            },
         )
     ]
     return knVectorstore.add_documents(docs)
 
-def add_memory(ai_id:str,mem: str, user_id: str,is_active: bool):
+
+def add_memory(ai_id: str, mem: str, user_id: str, is_active: bool):
     docs = [
         Document(
             page_content=mem,
@@ -103,7 +102,7 @@ def add_memory(ai_id:str,mem: str, user_id: str,is_active: bool):
                 "created_by": str(user_id),
                 "created_at": datetime.now().isoformat(),
                 "is_active": is_active,
-            }
+            },
         )
     ]
     return memVectorstore.add_documents(docs)