Merge pull request #1233 from Pickupppp/master

修复 milvus 作为默认数据库初始化异常和调用search方法异常
chatchat-space · Aug 24, 2023 · 67ac532 · 67ac532
2 parents bc3dbb4 + 050a585
commit 67ac532
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/server/knowledge_base/kb_service/milvus_kb_service.py b/server/knowledge_base/kb_service/milvus_kb_service.py
@@ -45,7 +45,7 @@ def do_init(self):
     def do_drop_kb(self):
         self.milvus.col.drop()
 
-    def do_search(self, query: str, top_k: int, embeddings: Embeddings):
+    def do_search(self, query: str, top_k: int,score_threshold: float, embeddings: Embeddings):
         # todo: support score threshold
         self._load_milvus(embeddings=embeddings)
         return self.milvus.similarity_search_with_score(query, top_k)
@@ -70,7 +70,8 @@ def do_delete_doc(self, kb_file: KnowledgeFile):
         self.milvus.col.delete(expr=f'pk in {delete_list}')
 
     def do_clear_vs(self):
-        self.milvus.col.drop()
+        if not self.milvus.col:
+            self.milvus.col.drop()
 
 
 if __name__ == '__main__':