Fix some issues in API (#2982)

### What problem does this PR solve?

Fix some issues in API

### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)

---------

Co-authored-by: liuhua <10215101452@stu.ecun.edu.cn>
This commit is contained in:
liuhua
2024-10-23 12:02:18 +08:00
committed by GitHub
parent 43b959fe58
commit 8714754afc
14 changed files with 206 additions and 106 deletions

View File

@ -15,17 +15,17 @@
#
from flask import request
from api.db import StatusEnum, FileSource
from api.db.db_models import File
from api.db.services.document_service import DocumentService
from api.db.services.file2document_service import File2DocumentService
from api.db.services.file_service import FileService
from api.db.services.knowledgebase_service import KnowledgebaseService
from api.db.services.llm_service import TenantLLMService
from api.db.services.user_service import TenantService
from api.settings import RetCode
from api.utils import get_uuid
from api.utils.api_utils import get_result, token_required, get_error_data_result, valid
from api.utils.api_utils import get_result, token_required, get_error_data_result, valid,get_parser_config
@manager.route('/dataset', methods=['POST'])
@ -36,15 +36,17 @@ def create(tenant_id):
permission = req.get("permission")
language = req.get("language")
chunk_method = req.get("chunk_method")
valid_permission = ("me", "team")
valid_language =("Chinese", "English")
valid_chunk_method = ("naive","manual","qa","table","paper","book","laws","presentation","picture","one","knowledge_graph","email")
parser_config = req.get("parser_config")
valid_permission = {"me", "team"}
valid_language ={"Chinese", "English"}
valid_chunk_method = {"naive","manual","qa","table","paper","book","laws","presentation","picture","one","knowledge_graph","email"}
check_validation=valid(permission,valid_permission,language,valid_language,chunk_method,valid_chunk_method)
if check_validation:
return check_validation
if "tenant_id" in req or "embedding_model" in req:
req["parser_config"]=get_parser_config(chunk_method,parser_config)
if "tenant_id" in req:
return get_error_data_result(
retmsg="`tenant_id` or `embedding_model` must not be provided")
retmsg="`tenant_id` must not be provided")
chunk_count=req.get("chunk_count")
document_count=req.get("document_count")
if chunk_count or document_count:
@ -59,9 +61,13 @@ def create(tenant_id):
retmsg="`name` is not empty string!")
if KnowledgebaseService.query(name=req["name"], tenant_id=tenant_id, status=StatusEnum.VALID.value):
return get_error_data_result(
retmsg="Duplicated knowledgebase name in creating dataset.")
retmsg="Duplicated dataset name in creating dataset.")
req["tenant_id"] = req['created_by'] = tenant_id
req['embedding_model'] = t.embd_id
if not req.get("embedding_model"):
req['embedding_model'] = t.embd_id
else:
if not TenantLLMService.query(tenant_id=tenant_id,model_type="embedding", llm_name=req.get("embedding_model")):
return get_error_data_result(f"`embedding_model` {req.get('embedding_model')} doesn't exist")
key_mapping = {
"chunk_num": "chunk_count",
"doc_num": "document_count",
@ -116,10 +122,12 @@ def update(tenant_id,dataset_id):
permission = req.get("permission")
language = req.get("language")
chunk_method = req.get("chunk_method")
valid_permission = ("me", "team")
valid_language =("Chinese", "English")
valid_chunk_method = ("naive","manual","qa","table","paper","book","laws","presentation","picture","one","knowledge_graph","email")
check_validation=valid(permission,valid_permission,language,valid_language,chunk_method,valid_chunk_method)
parser_config = req.get("parser_config")
valid_permission = {"me", "team"}
valid_language = {"Chinese", "English"}
valid_chunk_method = {"naive", "manual", "qa", "table", "paper", "book", "laws", "presentation", "picture", "one",
"knowledge_graph", "email"}
check_validation = valid(permission, valid_permission, language, valid_language, chunk_method, valid_chunk_method)
if check_validation:
return check_validation
if "tenant_id" in req:
@ -142,10 +150,16 @@ def update(tenant_id,dataset_id):
return get_error_data_result(
retmsg="If `chunk_count` is not 0, `chunk_method` is not changeable.")
req['parser_id'] = req.pop('chunk_method')
if req['parser_id'] != kb.parser_id:
req["parser_config"] = get_parser_config(chunk_method, parser_config)
if "embedding_model" in req:
if kb.chunk_num != 0 and req['embedding_model'] != kb.embd_id:
return get_error_data_result(
retmsg="If `chunk_count` is not 0, `embedding_method` is not changeable.")
if not req.get("embedding_model"):
return get_error_data_result("`embedding_model` can't be empty")
if not TenantLLMService.query(tenant_id=tenant_id,model_type="embedding", llm_name=req.get("embedding_model")):
return get_error_data_result(f"`embedding_model` {req.get('embedding_model')} doesn't exist")
req['embd_id'] = req.pop('embedding_model')
if "name" in req:
req["name"] = req["name"].strip()
@ -153,7 +167,7 @@ def update(tenant_id,dataset_id):
and len(KnowledgebaseService.query(name=req["name"], tenant_id=tenant_id,
status=StatusEnum.VALID.value)) > 0:
return get_error_data_result(
retmsg="Duplicated knowledgebase name in updating dataset.")
retmsg="Duplicated dataset name in updating dataset.")
if not KnowledgebaseService.update_by_id(kb.id, req):
return get_error_data_result(retmsg="Update dataset error.(Database error)")
return get_result(retcode=RetCode.SUCCESS)