update class method of KnowledgeBase and KnowledgeFile

This commit is contained in:
imClumsyPanda 2023-08-05 13:51:05 +08:00
parent 7beeaacea6
commit 14a5de6a29
2 changed files with 4 additions and 3 deletions

View File

@ -2,6 +2,7 @@ import urllib
from server.utils import BaseResponse, ListResponse from server.utils import BaseResponse, ListResponse
from server.knowledge_base.utils import validate_kb_name from server.knowledge_base.utils import validate_kb_name
from server.knowledge_base.knowledge_base import KnowledgeBase from server.knowledge_base.knowledge_base import KnowledgeBase
from configs.model_config import EMBEDDING_MODEL
async def list_kbs(): async def list_kbs():
@ -11,7 +12,7 @@ async def list_kbs():
async def create_kb(knowledge_base_name: str, async def create_kb(knowledge_base_name: str,
vector_store_type: str = "faiss", vector_store_type: str = "faiss",
embed_model: str = "m3e-base", embed_model: str = EMBEDDING_MODEL,
): ):
# Create selected knowledge base # Create selected knowledge base
if not validate_kb_name(knowledge_base_name): if not validate_kb_name(knowledge_base_name):

View File

@ -93,8 +93,8 @@ def delete_kb_from_db(kb_name):
class KnowledgeBase: class KnowledgeBase:
def __init__(self, def __init__(self,
knowledge_base_name: str, knowledge_base_name: str,
vector_store_type: str, vector_store_type: str = "faiss",
embed_model: str, embed_model: str = EMBEDDING_MODEL,
): ):
self.kb_name = knowledge_base_name self.kb_name = knowledge_base_name
if vector_store_type not in SUPPORTED_VS_TYPES: if vector_store_type not in SUPPORTED_VS_TYPES: