Fix some issues in API and test (#3001)

### What problem does this PR solve?

Fix some issues in API and test

### 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-24 20:05:21 +08:00
committed by GitHub
parent e997b42504
commit 161c7a231b
6 changed files with 64 additions and 76 deletions

View File

@ -30,9 +30,9 @@ from api.utils.api_utils import get_result
@token_required
def create(tenant_id):
req=request.json
ids= req.get("datasets")
ids= req.get("dataset_ids")
if not ids:
return get_error_data_result(retmsg="`datasets` is required")
return get_error_data_result(retmsg="`dataset_ids` is required")
for kb_id in ids:
kbs = KnowledgebaseService.query(id=kb_id,tenant_id=tenant_id)
if not kbs:
@ -138,7 +138,7 @@ def create(tenant_id):
res["llm"] = res.pop("llm_setting")
res["llm"]["model_name"] = res.pop("llm_id")
del res["kb_ids"]
res["datasets"] = req["datasets"]
res["dataset_ids"] = req["dataset_ids"]
res["avatar"] = res.pop("icon")
return get_result(data=res)
@ -148,8 +148,8 @@ def update(tenant_id,chat_id):
if not DialogService.query(tenant_id=tenant_id, id=chat_id, status=StatusEnum.VALID.value):
return get_error_data_result(retmsg='You do not own the chat')
req =request.json
ids = req.get("datasets")
if "datasets" in req:
ids = req.get("dataset_ids")
if "dataset_ids" in req:
if not ids:
return get_error_data_result("`datasets` can't be empty")
if ids:
@ -214,8 +214,8 @@ def update(tenant_id,chat_id):
# avatar
if "avatar" in req:
req["icon"] = req.pop("avatar")
if "datasets" in req:
req.pop("datasets")
if "dataset_ids" in req:
req.pop("dataset_ids")
if not DialogService.update_by_id(chat_id, req):
return get_error_data_result(retmsg="Chat not found!")
return get_result()

View File

@ -550,33 +550,32 @@ def update_chunk(tenant_id,dataset_id,document_id,chunk_id):
@token_required
def retrieval_test(tenant_id):
req = request.json
if not req.get("datasets"):
if not req.get("dataset_ids"):
return get_error_data_result("`datasets` is required.")
kb_ids = req["datasets"]
kb_ids = req["dataset_ids"]
if not isinstance(kb_ids,list):
return get_error_data_result("`datasets` should be a list")
kbs = KnowledgebaseService.get_by_ids(kb_ids)
embd_nms = list(set([kb.embd_id for kb in kbs]))
if len(embd_nms) != 1:
return get_result(
retmsg='Knowledge bases use different embedding models or does not exist."',
retcode=RetCode.AUTHENTICATION_ERROR)
if isinstance(kb_ids, str): kb_ids = [kb_ids]
for id in kb_ids:
if not KnowledgebaseService.query(id=id,tenant_id=tenant_id):
return get_error_data_result(f"You don't own the dataset {id}.")
embd_nms = list(set([kb.embd_id for kb in kbs]))
if len(embd_nms) != 1:
return get_result(
retmsg='Datasets use different embedding models."',
retcode=RetCode.AUTHENTICATION_ERROR)
if "question" not in req:
return get_error_data_result("`question` is required.")
page = int(req.get("offset", 1))
size = int(req.get("limit", 1024))
question = req["question"]
doc_ids = req.get("documents", [])
if not isinstance(req.get("documents"),list):
doc_ids = req.get("document_ids", [])
if not isinstance(doc_ids,list):
return get_error_data_result("`documents` should be a list")
doc_ids_list=KnowledgebaseService.list_documents_by_ids(kb_ids)
for doc_id in doc_ids:
if doc_id not in doc_ids_list:
return get_error_data_result(f"You don't own the document {doc_id}")
return get_error_data_result(f"The datasets don't own the document {doc_id}")
similarity_threshold = float(req.get("similarity_threshold", 0.2))
vector_similarity_weight = float(req.get("vector_similarity_weight", 0.3))
top = int(req.get("top_k", 1024))