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
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 68 additions and 17 deletions

View File

@ -24,6 +24,7 @@ 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):
@ -50,14 +51,11 @@ 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):
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)
@ -150,6 +148,7 @@ 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):
@ -202,6 +201,7 @@ 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):

View File

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

View File

@ -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