mirror of
https://gitee.com/infiniflow/ragflow.git
synced 2025-12-06 15:29:03 +08:00
b6c4722687
### What problem does this PR solve? Try to make this more asynchronous. Verified in chat and agent scenarios, reducing blocking behavior. #11551, #11579. However, the impact of these changes still requires further investigation to ensure everything works as expected. ### Type of change - [x] Refactoring
118 lines
4.5 KiB
Python
118 lines
4.5 KiB
Python
#
|
|
# Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
#
|
|
from datetime import datetime, timedelta
|
|
from quart import request
|
|
from api.db.db_models import APIToken
|
|
from api.db.services.api_service import APITokenService, API4ConversationService
|
|
from api.db.services.user_service import UserTenantService
|
|
from api.utils.api_utils import generate_confirmation_token, get_data_error_result, get_json_result, get_request_json, server_error_response, validate_request
|
|
from common.time_utils import current_timestamp, datetime_format
|
|
from api.apps import login_required, current_user
|
|
|
|
|
|
@manager.route('/new_token', methods=['POST']) # noqa: F821
|
|
@login_required
|
|
async def new_token():
|
|
req = await get_request_json()
|
|
try:
|
|
tenants = UserTenantService.query(user_id=current_user.id)
|
|
if not tenants:
|
|
return get_data_error_result(message="Tenant not found!")
|
|
|
|
tenant_id = tenants[0].tenant_id
|
|
obj = {"tenant_id": tenant_id, "token": generate_confirmation_token(),
|
|
"create_time": current_timestamp(),
|
|
"create_date": datetime_format(datetime.now()),
|
|
"update_time": None,
|
|
"update_date": None
|
|
}
|
|
if req.get("canvas_id"):
|
|
obj["dialog_id"] = req["canvas_id"]
|
|
obj["source"] = "agent"
|
|
else:
|
|
obj["dialog_id"] = req["dialog_id"]
|
|
|
|
if not APITokenService.save(**obj):
|
|
return get_data_error_result(message="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']) # noqa: F821
|
|
@login_required
|
|
def token_list():
|
|
try:
|
|
tenants = UserTenantService.query(user_id=current_user.id)
|
|
if not tenants:
|
|
return get_data_error_result(message="Tenant not found!")
|
|
|
|
id = request.args["dialog_id"] if "dialog_id" in request.args else request.args["canvas_id"]
|
|
objs = APITokenService.query(tenant_id=tenants[0].tenant_id, dialog_id=id)
|
|
return get_json_result(data=[o.to_dict() for o in objs])
|
|
except Exception as e:
|
|
return server_error_response(e)
|
|
|
|
|
|
@manager.route('/rm', methods=['POST']) # noqa: F821
|
|
@validate_request("tokens", "tenant_id")
|
|
@login_required
|
|
async def rm():
|
|
req = await get_request_json()
|
|
try:
|
|
for token in req["tokens"]:
|
|
APITokenService.filter_delete(
|
|
[APIToken.tenant_id == req["tenant_id"], APIToken.token == token])
|
|
return get_json_result(data=True)
|
|
except Exception as e:
|
|
return server_error_response(e)
|
|
|
|
|
|
@manager.route('/stats', methods=['GET']) # noqa: F821
|
|
@login_required
|
|
def stats():
|
|
try:
|
|
tenants = UserTenantService.query(user_id=current_user.id)
|
|
if not tenants:
|
|
return get_data_error_result(message="Tenant not found!")
|
|
objs = API4ConversationService.stats(
|
|
tenants[0].tenant_id,
|
|
request.args.get(
|
|
"from_date",
|
|
(datetime.now() -
|
|
timedelta(
|
|
days=7)).strftime("%Y-%m-%d 00:00:00")),
|
|
request.args.get(
|
|
"to_date",
|
|
datetime.now().strftime("%Y-%m-%d %H:%M:%S")),
|
|
"agent" if "canvas_id" in request.args else None)
|
|
|
|
res = {"pv": [], "uv": [], "speed": [], "tokens": [], "round": [], "thumb_up": []}
|
|
|
|
for obj in objs:
|
|
dt = obj["dt"]
|
|
res["pv"].append((dt, obj["pv"]))
|
|
res["uv"].append((dt, obj["uv"]))
|
|
res["speed"].append((dt, float(obj["tokens"]) / (float(obj["duration"]) + 0.1))) # +0.1 to avoid division by zero
|
|
res["tokens"].append((dt, float(obj["tokens"]) / 1000.0)) # convert to thousands
|
|
res["round"].append((dt, obj["round"]))
|
|
res["thumb_up"].append((dt, obj["thumb_up"]))
|
|
|
|
return get_json_result(data=res)
|
|
except Exception as e:
|
|
return server_error_response(e)
|