refine API token application (#2847)

### What problem does this PR solve?

#2846

### Type of change

- [x] New Feature (non-breaking change which adds functionality)
This commit is contained in:
Kevin Hu
2024-10-15 15:47:40 +08:00
committed by GitHub
parent fcabdf7745
commit ce495e4e3e
3 changed files with 68 additions and 17 deletions

View File

@ -24,9 +24,10 @@ from api.utils import get_uuid
from api.utils.api_utils import get_error_data_result
from api.utils.api_utils import get_result, token_required
@manager.route('/chat/<chat_id>/session', methods=['POST'])
@token_required
def create(tenant_id,chat_id):
def create(tenant_id, chat_id):
req = request.json
req["dialog_id"] = chat_id
dia = DialogService.query(tenant_id=tenant_id, id=req["dialog_id"], status=StatusEnum.VALID.value)
@ -50,17 +51,14 @@ def create(tenant_id,chat_id):
del conv["reference"]
return get_result(data=conv)
@manager.route('/chat/<chat_id>/session/<session_id>', methods=['PUT'])
@token_required
def update(tenant_id,chat_id,session_id):
def update(tenant_id, chat_id, session_id):
req = request.json
if "dialog_id" in req and req.get("dialog_id") != chat_id:
return get_error_data_result(retmsg="Can't change chat_id")
if "chat_id" in req and req.get("chat_id") != chat_id:
return get_error_data_result(retmsg="Can't change chat_id")
req["dialog_id"] = chat_id
conv_id = session_id
conv = ConversationService.query(id=conv_id,dialog_id=chat_id)
conv = ConversationService.query(id=conv_id, dialog_id=chat_id)
if not conv:
return get_error_data_result(retmsg="Session does not exist")
if not DialogService.query(id=chat_id, tenant_id=tenant_id, status=StatusEnum.VALID.value):
@ -78,14 +76,14 @@ def update(tenant_id,chat_id,session_id):
@manager.route('/chat/<chat_id>/session/<session_id>/completion', methods=['POST'])
@token_required
def completion(tenant_id,chat_id,session_id):
def completion(tenant_id, chat_id, session_id):
req = request.json
# req = {"conversation_id": "9aaaca4c11d311efa461fa163e197198", "messages": [
# {"role": "user", "content": "上海有吗?"}
# ]}
if not req.get("question"):
return get_error_data_result(retmsg="Please input your question.")
conv = ConversationService.query(id=session_id,dialog_id=chat_id)
conv = ConversationService.query(id=session_id, dialog_id=chat_id)
if not conv:
return get_error_data_result(retmsg="Session does not exist")
conv = conv[0]
@ -125,7 +123,7 @@ def completion(tenant_id,chat_id,session_id):
try:
for ans in chat(dia, msg, **req):
fillin_conv(ans)
yield "data:" + json.dumps({"code": 0, "data": ans}, ensure_ascii=False) + "\n\n"
yield "data:" + json.dumps({"code": 0, "data": ans}, ensure_ascii=False) + "\n\n"
ConversationService.update_by_id(conv.id, conv.to_dict())
except Exception as e:
yield "data:" + json.dumps({"code": 500, "message": str(e),
@ -150,14 +148,15 @@ def completion(tenant_id,chat_id,session_id):
break
return get_result(data=answer)
@manager.route('/chat/<chat_id>/session', methods=['GET'])
@token_required
def list(chat_id,tenant_id):
def list(chat_id, tenant_id):
if not DialogService.query(tenant_id=tenant_id, id=chat_id, status=StatusEnum.VALID.value):
return get_error_data_result(retmsg=f"You don't own the assistant {chat_id}.")
id = request.args.get("id")
name = request.args.get("name")
session = ConversationService.query(id=id,name=name,dialog_id=chat_id)
session = ConversationService.query(id=id, name=name, dialog_id=chat_id)
if not session:
return get_error_data_result(retmsg="The session doesn't exist")
page_number = int(request.args.get("page", 1))
@ -167,7 +166,7 @@ def list(chat_id,tenant_id):
desc = False
else:
desc = True
convs = ConversationService.get_list(chat_id,page_number,items_per_page,orderby,desc,id,name)
convs = ConversationService.get_list(chat_id, page_number, items_per_page, orderby, desc, id, name)
if not convs:
return get_result(data=[])
for conv in convs:
@ -202,16 +201,17 @@ def list(chat_id,tenant_id):
del conv["reference"]
return get_result(data=convs)
@manager.route('/chat/<chat_id>/session', methods=["DELETE"])
@token_required
def delete(tenant_id,chat_id):
def delete(tenant_id, chat_id):
if not DialogService.query(id=chat_id, tenant_id=tenant_id, status=StatusEnum.VALID.value):
return get_error_data_result(retmsg="You don't own the chat")
ids = request.json.get("ids")
if not ids:
return get_error_data_result(retmsg="`ids` is required in deleting operation")
for id in ids:
conv = ConversationService.query(id=id,dialog_id=chat_id)
conv = ConversationService.query(id=id, dialog_id=chat_id)
if not conv:
return get_error_data_result(retmsg="The chat doesn't own the session")
ConversationService.delete_by_id(id)