From 14a5de6a29c5925da9b2515318c64ef64dc871b5 Mon Sep 17 00:00:00 2001 From: imClumsyPanda Date: Sat, 5 Aug 2023 13:51:05 +0800 Subject: [PATCH] update class method of KnowledgeBase and KnowledgeFile --- server/knowledge_base/kb_api.py | 3 ++- server/knowledge_base/knowledge_base.py | 4 ++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/server/knowledge_base/kb_api.py b/server/knowledge_base/kb_api.py index f035838..d5ef57b 100644 --- a/server/knowledge_base/kb_api.py +++ b/server/knowledge_base/kb_api.py @@ -2,6 +2,7 @@ import urllib from server.utils import BaseResponse, ListResponse from server.knowledge_base.utils import validate_kb_name from server.knowledge_base.knowledge_base import KnowledgeBase +from configs.model_config import EMBEDDING_MODEL async def list_kbs(): @@ -11,7 +12,7 @@ async def list_kbs(): async def create_kb(knowledge_base_name: str, vector_store_type: str = "faiss", - embed_model: str = "m3e-base", + embed_model: str = EMBEDDING_MODEL, ): # Create selected knowledge base if not validate_kb_name(knowledge_base_name): diff --git a/server/knowledge_base/knowledge_base.py b/server/knowledge_base/knowledge_base.py index ad0aa6a..922dcc7 100644 --- a/server/knowledge_base/knowledge_base.py +++ b/server/knowledge_base/knowledge_base.py @@ -93,8 +93,8 @@ def delete_kb_from_db(kb_name): class KnowledgeBase: def __init__(self, knowledge_base_name: str, - vector_store_type: str, - embed_model: str, + vector_store_type: str = "faiss", + embed_model: str = EMBEDDING_MODEL, ): self.kb_name = knowledge_base_name if vector_store_type not in SUPPORTED_VS_TYPES: