refine API token application (#2847)

### What problem does this PR solve?

#2846

### Type of change

- [x] New Feature (non-breaking change which adds functionality)

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/<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)

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)
+

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