diff --git a/server/knowledge_base/kb_service/base.py b/server/knowledge_base/kb_service/base.py index a81474f..d823740 100644 --- a/server/knowledge_base/kb_service/base.py +++ b/server/knowledge_base/kb_service/base.py @@ -278,15 +278,10 @@ class KBServiceFactory: return DefaultKBService(kb_name) @staticmethod - def get_service_by_name(kb_name: str, - default_vs_type: SupportedVSType = SupportedVSType.FAISS, - default_embed_model: str = EMBEDDING_MODEL, - ) -> KBService: + def get_service_by_name(kb_name: str) -> KBService: _, vs_type, embed_model = load_kb_from_db(kb_name) - if vs_type is None: # faiss knowledge base not in db - vs_type = default_vs_type - if embed_model is None: - embed_model = default_embed_model + if _ is None: # kb not in db, just return None + return None return KBServiceFactory.get_service(kb_name, vs_type, embed_model) @staticmethod @@ -331,6 +326,9 @@ def get_kb_details() -> List[Dict]: def get_kb_file_details(kb_name: str) -> List[Dict]: kb = KBServiceFactory.get_service_by_name(kb_name) + if kb is None: + return [] + files_in_folder = list_files_from_folder(kb_name) files_in_db = kb.list_files() result = {} diff --git a/server/knowledge_base/migrate.py b/server/knowledge_base/migrate.py index 146dc77..0c21946 100644 --- a/server/knowledge_base/migrate.py +++ b/server/knowledge_base/migrate.py @@ -157,7 +157,7 @@ def prune_db_docs(kb_names: List[str]): """ for kb_name in kb_names: kb = KBServiceFactory.get_service_by_name(kb_name) - if kb and kb.exists(): + if kb is not None: files_in_db = kb.list_files() files_in_folder = list_files_from_folder(kb_name) files = list(set(files_in_db) - set(files_in_folder)) @@ -175,7 +175,7 @@ def prune_folder_files(kb_names: List[str]): """ for kb_name in kb_names: kb = KBServiceFactory.get_service_by_name(kb_name) - if kb and kb.exists(): + if kb is not None: files_in_db = kb.list_files() files_in_folder = list_files_from_folder(kb_name) files = list(set(files_in_folder) - set(files_in_db)) diff --git a/tests/test_migrate.py b/tests/test_migrate.py index d694b02..7195dd3 100644 --- a/tests/test_migrate.py +++ b/tests/test_migrate.py @@ -33,7 +33,7 @@ def test_recreate_vs(): folder2db([kb_name], "recreate_vs") kb = KBServiceFactory.get_service_by_name(kb_name) - assert kb.exists() + assert kb and kb.exists() files = kb.list_files() print(files)