move commonto tools from kbservice.base to server.knowledge_base.utils

This commit is contained in:
liunux4odoo 2023-08-08 16:40:18 +08:00
parent 584e3a9234
commit c4a3995310
4 changed files with 8 additions and 18 deletions

View File

@ -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

View File

@ -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

View File

@ -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,

View File

@ -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):