move KBServiceFactory to server.knowledge_base.kb_service.base

configs/model_config.py.example

@@ -242,7 +242,7 @@ KB_ROOT_PATH = os.path.join(os.path.dirname(os.path.dirname(__file__)), "knowled
 # 数据库默认存储路径。
 # 如果使用sqlite，可以直接修改DB_ROOT_PATH；如果使用其它数据库，请直接修改SQLALCHEMY_DATABASE_URI。
 DB_ROOT_PATH = os.path.join(KB_ROOT_PATH, "info.db")
-SQLALCHEMY_DATABASE_URI = f"sqlite://{DB_ROOT_PATH}"
+SQLALCHEMY_DATABASE_URI = f"sqlite:///{DB_ROOT_PATH}"
 
 # 缓存向量库数量
 CACHED_VS_NUM = 1

server/chat/knowledge_base_chat.py

@@ -10,8 +10,7 @@ from langchain.callbacks import AsyncIteratorCallbackHandler
 from typing import AsyncIterable
 import asyncio
 from langchain.prompts import PromptTemplate
-from server.knowledge_base.utils import KBServiceFactory
+from server.knowledge_base.kb_service.base import KBService, KBServiceFactory
-from server.knowledge_base.kb_service.base import KBService
 import json
 
 

server/knowledge_base/__init__.py

@@ -1,3 +1,3 @@
-from .kb_api import list_kbs, create_kb, delete_kb
+# from .kb_api import list_kbs, create_kb, delete_kb
-from .kb_doc_api import list_docs, upload_doc, delete_doc, update_doc, download_doc, recreate_vector_store
+# from .kb_doc_api import list_docs, upload_doc, delete_doc, update_doc, download_doc, recreate_vector_store
-from .utils import KnowledgeFile, KBServiceFactory
+# from .utils import KnowledgeFile, KBServiceFactory

server/knowledge_base/kb_api.py

@@ -1,6 +1,7 @@
 import urllib
 from server.utils import BaseResponse, ListResponse
-from server.knowledge_base.utils import validate_kb_name, KBServiceFactory
+from server.knowledge_base.utils import validate_kb_name
+from server.knowledge_base.kb_service.base import KBServiceFactory
 from server.db.repository.knowledge_base_repository import list_kbs_from_db
 from configs.model_config import EMBEDDING_MODEL
 

server/knowledge_base/kb_doc_api.py

@@ -5,7 +5,8 @@ from server.utils import BaseResponse, ListResponse
 from server.knowledge_base.utils import (validate_kb_name)
 from fastapi.responses import StreamingResponse
 import json
-from server.knowledge_base.utils import KnowledgeFile, KBServiceFactory, list_docs_from_folder
+from server.knowledge_base.utils import KnowledgeFile, list_docs_from_folder
+from server.knowledge_base.kb_service.base import KBServiceFactory
 from server.knowledge_base.kb_service.base import SupportedVSType
 from server.knowledge_base.kb_service.faiss_kb_service import refresh_vs_cache
 

server/knowledge_base/kb_service/base.py

@@ -7,13 +7,13 @@ from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.embeddings.base import Embeddings
 from langchain.docstore.document import Document
 
-from server.db.repository.knowledge_base_repository import add_kb_to_db, delete_kb_from_db, list_kbs_from_db, kb_exists
+from server.db.repository.knowledge_base_repository import add_kb_to_db, delete_kb_from_db, list_kbs_from_db, kb_exists, load_kb_from_db
 from server.db.repository.knowledge_file_repository import add_doc_to_db, delete_file_from_db, doc_exists, \
     list_docs_from_db
 from configs.model_config import (DB_ROOT_PATH, kbs_config, VECTOR_SEARCH_TOP_K,
                                   embedding_model_dict, EMBEDDING_DEVICE, EMBEDDING_MODEL)
 from server.knowledge_base.utils import (get_kb_path, get_doc_path, load_embeddings, KnowledgeFile)
-from typing import List
+from typing import List, Union
 
 
 class SupportedVSType:
@@ -162,3 +162,33 @@ class KBService(ABC):
         从知识库删除全部向量子类实自己逻辑
         """
         pass
+
+
+class KBServiceFactory:
+
+    @staticmethod
+    def get_service(kb_name: str,
+                    vector_store_type: Union[str, SupportedVSType],
+                    embed_model: str = EMBEDDING_MODEL,
+                    ) -> KBService:
+        if isinstance(vector_store_type, str):
+            vector_store_type = getattr(SupportedVSType, vector_store_type.upper())
+        if SupportedVSType.FAISS == vector_store_type:
+            from server.knowledge_base.kb_service.faiss_kb_service import FaissKBService
+            return FaissKBService(kb_name, embed_model=embed_model)
+        elif SupportedVSType.MILVUS == vector_store_type:
+            from server.knowledge_base.kb_service.milvus_kb_service import MilvusKBService
+            return MilvusKBService(kb_name, embed_model=embed_model) # other milvus parameters are set in model_config.kbs_config
+        elif SupportedVSType.DEFAULT == vector_store_type: # kb_exists of default kbservice is False, to make validation easier.
+            return DefaultKBService(kb_name)
+
+    @staticmethod
+    def get_service_by_name(kb_name: str
+                            ) -> KBService:
+        kb_name, vs_type, embed_model = load_kb_from_db(kb_name)
+        return KBServiceFactory.get_service(kb_name, vs_type, embed_model)
+
+    @staticmethod
+    def get_default():
+        return KBServiceFactory.get_service("default", SupportedVSType.DEFAULT)
+

server/knowledge_base/utils.py

@@ -3,9 +3,6 @@ import os
 from langchain.embeddings.huggingface import HuggingFaceEmbeddings
 from configs.model_config import (embedding_model_dict, KB_ROOT_PATH, EMBEDDING_MODEL, kbs_config)
 from functools import lru_cache
-from server.knowledge_base.kb_service.base import KBService, SupportedVSType
-from server.db.repository.knowledge_base_repository import load_kb_from_db
-from server.knowledge_base.kb_service.default_kb_service import DefaultKBService
 
 
 def validate_kb_name(knowledge_base_id: str) -> bool:
@@ -26,6 +23,10 @@ def get_vs_path(knowledge_base_name: str):
 def get_file_path(knowledge_base_name: str, doc_name: str):
     return os.path.join(get_doc_path(knowledge_base_name), doc_name)
 
+def list_kbs_from_folder():
+    return [f for f in os.listdir(KB_ROOT_PATH)
+            if os.path.isdir(os.path.join(KB_ROOT_PATH, f))]
+
 def list_docs_from_folder(kb_name: str):
     doc_path = get_doc_path(kb_name)
     return [file for file in os.listdir(doc_path)
@@ -78,33 +79,3 @@ class KnowledgeFile:
         TextSplitter = getattr(sys.modules['langchain.text_splitter'], self.text_splitter_name)
         text_splitter = TextSplitter(chunk_size=500, chunk_overlap=200)
         return loader.load_and_split(text_splitter)
-
-
-class KBServiceFactory:
-
-    @staticmethod
-    def get_service(kb_name: str,
-                    vector_store_type: Union[str, SupportedVSType],
-                    embed_model: str = EMBEDDING_MODEL,
-                    ) -> KBService:
-        if isinstance(vector_store_type, str):
-            vector_store_type = getattr(SupportedVSType, vector_store_type.upper())
-        if SupportedVSType.FAISS == vector_store_type:
-            from server.knowledge_base.kb_service.faiss_kb_service import FaissKBService
-            return FaissKBService(kb_name, embed_model=embed_model)
-        elif SupportedVSType.MILVUS == vector_store_type:
-            from server.knowledge_base.kb_service.milvus_kb_service import MilvusKBService
-            return MilvusKBService(kb_name, embed_model=embed_model) # other milvus parameters are set in model_config.kbs_config
-        elif SupportedVSType.DEFAULT == vector_store_type: # kb_exists of default kbservice is False, to make validation easier.
-            return DefaultKBService(kb_name)
-
-    @staticmethod
-    def get_service_by_name(kb_name: str
-                            ) -> KBService:
-        kb_name, vs_type, embed_model = load_kb_from_db(kb_name)
-        return KBServiceFactory.get_service(kb_name, vs_type, embed_model)
-
-    @staticmethod
-    def get_default():
-        return KBServiceFactory.get_service("default", SupportedVSType.DEFAULT)
-