diff --git a/api/apps/sdk/session.py b/api/apps/sdk/session.py index 5101c3831..81bc84aca 100644 --- a/api/apps/sdk/session.py +++ b/api/apps/sdk/session.py @@ -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//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//session/', 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//session//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//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//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) diff --git a/api/apps/system_app.py b/api/apps/system_app.py index e55a95037..35407ec06 100644 --- a/api/apps/system_app.py +++ b/api/apps/system_app.py @@ -14,12 +14,17 @@ # limitations under the License # import json +from datetime import datetime -from flask_login import login_required +from flask_login import login_required, current_user +from api.apps.api_app import generate_confirmation_token +from api.db.services.api_service import APITokenService from api.db.services.knowledgebase_service import KnowledgebaseService +from api.db.services.user_service import UserTenantService from api.settings import DATABASE_TYPE -from api.utils.api_utils import get_json_result +from api.utils import current_timestamp, datetime_format +from api.utils.api_utils import get_json_result, request, get_data_error_result, server_error_response from api.versions import get_rag_version from rag.settings import SVR_QUEUE_NAME from rag.utils.es_conn import ELASTICSEARCH @@ -88,3 +93,42 @@ def status(): res["task_executor"] = {"status": "red", "error": str(e)} return get_json_result(data=res) + + +@manager.route('/new_token', methods=['POST']) +@login_required +def new_token(): + try: + tenants = UserTenantService.query(user_id=current_user.id) + if not tenants: + return get_data_error_result(retmsg="Tenant not found!") + + tenant_id = tenants[0].tenant_id + obj = {"tenant_id": tenant_id, "token": generate_confirmation_token(tenant_id), + "create_time": current_timestamp(), + "create_date": datetime_format(datetime.now()), + "update_time": None, + "update_date": None + } + + if not APITokenService.save(**obj): + return get_data_error_result(retmsg="Fail to new a dialog!") + + return get_json_result(data=obj) + except Exception as e: + return server_error_response(e) + + +@manager.route('/token_list', methods=['GET']) +@login_required +def token_list(): + try: + tenants = UserTenantService.query(user_id=current_user.id) + if not tenants: + return get_data_error_result(retmsg="Tenant not found!") + + objs = APITokenService.query(tenant_id=tenants[0].tenant_id) + return get_json_result(data=[o.to_dict() for o in objs]) + except Exception as e: + return server_error_response(e) + diff --git a/api/db/db_models.py b/api/db/db_models.py index 8b547f8ec..9a6fa2527 100644 --- a/api/db/db_models.py +++ b/api/db/db_models.py @@ -1052,4 +1052,11 @@ def migrate_db(): ) except Exception as e: pass + try: + migrate( + migrator.alter_column_type('api_token', 'dialog_id', + CharField(max_length=32, null=True, index=True)) + ) + except Exception as e: + pass