pengjunjie/ragflow
1
0
Fork 1
mirror of https://github.com/infiniflow/ragflow.git synced 2025-12-08 20:42:30 +08:00
Code Issues Packages Projects Releases Wiki Activity

Fix some issues in API (#2982)

Browse Source 
### 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
Download Patch File Download Diff File Expand all files Collapse all files

59
api/apps/sdk/chat.py
View File

@ -18,20 +18,21 @@ from flask import request
from api.db import StatusEnum
from api.db.services.dialog_service import DialogService
from api.db.services.knowledgebase_service import KnowledgebaseService
from api.db.services.llm_service import TenantLLMService
from api.db.services.llm_service import TenantLLMService
from api.db.services.user_service import TenantService
from api.utils import get_uuid
from api.utils.api_utils import get_error_data_result, token_required
from api.utils.api_utils import get_result
@manager.route('/chat', methods=['POST'])
@token_required
def create(tenant_id):
req=request.json
ids= req.get("knowledgebases")
ids= req.get("datasets")
if not ids:
return get_error_data_result(retmsg="`knowledgebases` is required")
return get_error_data_result(retmsg="`datasets` is required")
for kb_id in ids:
kbs = KnowledgebaseService.query(id=kb_id,tenant_id=tenant_id)
if not kbs:
@ -45,6 +46,8 @@ def create(tenant_id):
if llm:
if "model_name" in llm:
req["llm_id"] = llm.pop("model_name")
if not TenantLLMService.query(tenant_id=tenant_id,llm_name=req["llm_id"],model_type="chat"):
return get_error_data_result(f"`model_name` {req.get('llm_id')} doesn't exist")
req["llm_setting"] = req.pop("llm")
e, tenant = TenantService.get_by_id(tenant_id)
if not e:
@ -73,10 +76,10 @@ def create(tenant_id):
req["top_n"] = req.get("top_n", 6)
req["top_k"] = req.get("top_k", 1024)
req["rerank_id"] = req.get("rerank_id", "")
if req.get("llm_id"):
if not TenantLLMService.query(llm_name=req["llm_id"]):
return get_error_data_result(retmsg="the model_name does not exist.")
else:
if req.get("rerank_id"):
if not TenantLLMService.query(tenant_id=tenant_id,llm_name=req.get("rerank_id"),model_type="rerank"):
return get_error_data_result(f"`rerank_model` {req.get('rerank_id')} doesn't exist")
if not req.get("llm_id"):
req["llm_id"] = tenant.llm_id
if not req.get("name"):
return get_error_data_result(retmsg="`name` is required.")
@ -135,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["knowledgebases"] = req["knowledgebases"]
res["datasets"] = req["datasets"]
res["avatar"] = res.pop("icon")
return get_result(data=res)
@ -145,27 +148,32 @@ 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
if "knowledgebases" in req:
if not req.get("knowledgebases"):
return get_error_data_result(retmsg="`knowledgebases` can't be empty value")
kb_list = []
for kb in req.get("knowledgebases"):
if not kb["id"]:
return get_error_data_result(retmsg="knowledgebase needs id")
if not KnowledgebaseService.query(id=kb["id"], tenant_id=tenant_id):
return get_error_data_result(retmsg="you do not own the knowledgebase")
# if not DocumentService.query(kb_id=kb["id"]):
# return get_error_data_result(retmsg="There is a invalid knowledgebase")
kb_list.append(kb["id"])
req["kb_ids"] = kb_list
ids = req.get("datasets")
if "datasets" in req:
if not ids:
return get_error_data_result("`datasets` can't be empty")
if ids:
for kb_id in ids:
kbs = KnowledgebaseService.query(id=kb_id, tenant_id=tenant_id)
if not kbs:
return get_error_data_result(f"You don't own the dataset {kb_id}")
kb = kbs[0]
if kb.chunk_num == 0:
return get_error_data_result(f"The dataset {kb_id} doesn't own parsed file")
req["kb_ids"] = ids
llm = req.get("llm")
if llm:
if "model_name" in llm:
req["llm_id"] = llm.pop("model_name")
if not TenantLLMService.query(tenant_id=tenant_id,llm_name=req["llm_id"],model_type="chat"):
return get_error_data_result(f"`model_name` {req.get('llm_id')} doesn't exist")
req["llm_setting"] = req.pop("llm")
e, tenant = TenantService.get_by_id(tenant_id)
if not e:
return get_error_data_result(retmsg="Tenant not found!")
if req.get("rerank_model"):
if not TenantLLMService.query(tenant_id=tenant_id,llm_name=req.get("rerank_model"),model_type="rerank"):
return get_error_data_result(f"`rerank_model` {req.get('rerank_model')} doesn't exist")
# prompt
prompt = req.get("prompt")
key_mapping = {"parameters": "variables",
@ -185,9 +193,6 @@ def update(tenant_id,chat_id):
req["prompt_config"] = req.pop("prompt")
e, res = DialogService.get_by_id(chat_id)
res = res.to_json()
if "llm_id" in req:
if not TenantLLMService.query(llm_name=req["llm_id"]):
return get_error_data_result(retmsg="The `model_name` does not exist.")
if "name" in req:
if not req.get("name"):
return get_error_data_result(retmsg="`name` is not empty.")
@ -209,8 +214,8 @@ def update(tenant_id,chat_id):
# avatar
if "avatar" in req:
req["icon"] = req.pop("avatar")
if "knowledgebases" in req:
req.pop("knowledgebases")
if "datasets" in req:
req.pop("datasets")
if not DialogService.update_by_id(chat_id, req):
return get_error_data_result(retmsg="Chat not found!")
return get_result()
@ -279,7 +284,7 @@ def list_chat(tenant_id):
return get_error_data_result(retmsg=f"Don't exist the kb {kb_id}")
kb_list.append(kb[0].to_json())
del res["kb_ids"]
res["knowledgebases"] = kb_list
res["datasets"] = kb_list
res["avatar"] = res.pop("icon")
list_assts.append(res)
return get_result(data=list_assts)