From c4a39953103cc70e9371abe6bfa9d9ec61bd8dac Mon Sep 17 00:00:00 2001
From: liunux4odoo <liunu@qq.com>
Date: Tue, 8 Aug 2023 16:40:18 +0800
Subject: [PATCH] move commonto tools from kbservice.base to
 server.knowledge_base.utils

---
 server/knowledge_base/kb_api.py          |  2 +-
 server/knowledge_base/kb_doc_api.py      |  4 ++--
 server/knowledge_base/kb_service/base.py | 15 +--------------
 server/knowledge_base/utils.py           |  5 ++++-
 4 files changed, 8 insertions(+), 18 deletions(-)

diff --git a/server/knowledge_base/kb_api.py b/server/knowledge_base/kb_api.py
index 3abee4b..0ae7ac1 100644
--- a/server/knowledge_base/kb_api.py
+++ b/server/knowledge_base/kb_api.py
@@ -1,7 +1,7 @@
 import urllib
 from server.utils import BaseResponse, ListResponse
 from server.knowledge_base.utils import validate_kb_name, KBServiceFactory
-from server.knowledge_base.kb_service.base import list_kbs_from_db
+from server.db.repository.knowledge_base_repository import list_kbs_from_db
 from configs.model_config import EMBEDDING_MODEL
 
 
diff --git a/server/knowledge_base/kb_doc_api.py b/server/knowledge_base/kb_doc_api.py
index fc8a782..7721ebe 100644
--- a/server/knowledge_base/kb_doc_api.py
+++ b/server/knowledge_base/kb_doc_api.py
@@ -5,8 +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
-from server.knowledge_base.kb_service.base import SupportedVSType, list_docs_from_folder
+from server.knowledge_base.utils import KnowledgeFile, KBServiceFactory, list_docs_from_folder
+from server.knowledge_base.kb_service.base import SupportedVSType
 from server.knowledge_base.kb_service.faiss_kb_service import refresh_vs_cache
 
 
diff --git a/server/knowledge_base/kb_service/base.py b/server/knowledge_base/kb_service/base.py
index 6f78649..0b542e5 100644
--- a/server/knowledge_base/kb_service/base.py
+++ b/server/knowledge_base/kb_service/base.py
@@ -12,7 +12,7 @@ from server.db.repository.knowledge_file_repository import add_doc_to_db, delete
     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, KnowledgeFile)
+from server.knowledge_base.utils import (get_kb_path, get_doc_path, load_embeddings, KnowledgeFile)
 from typing import List
 
 
@@ -22,19 +22,6 @@ class SupportedVSType:
     DEFAULT = 'default'
 
 
-def list_docs_from_folder(kb_name: str):
-    doc_path = get_doc_path(kb_name)
-    return [file for file in os.listdir(doc_path)
-            if os.path.isfile(os.path.join(doc_path, file))]
-
-
-@lru_cache(1)
-def load_embeddings(model: str, device: str):
-    embeddings = HuggingFaceEmbeddings(model_name=embedding_model_dict[model],
-                                       model_kwargs={'device': device})
-    return embeddings
-
-
 class KBService(ABC):
 
     def __init__(self,
diff --git a/server/knowledge_base/utils.py b/server/knowledge_base/utils.py
index 00344e1..91e9b40 100644
--- a/server/knowledge_base/utils.py
+++ b/server/knowledge_base/utils.py
@@ -6,7 +6,6 @@ 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
-from server.knowledge_base.kb_service.milvus_kb_service import MilvusKBService
 
 
 def validate_kb_name(knowledge_base_id: str) -> bool:
@@ -27,6 +26,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_docs_from_folder(kb_name: str):
+    doc_path = get_doc_path(kb_name)
+    return [file for file in os.listdir(doc_path)
+            if os.path.isfile(os.path.join(doc_path, file))]
 
 @lru_cache(1)
 def load_embeddings(model: str, device: str):