mirror of
https://git.mirrors.martin98.com/https://github.com/langgenius/dify.git
synced 2025-07-08 12:11:49 +08:00

Eliminates unnecessary pre-calculation of token limits and recalculation of max tokens across multiple app runners, simplifying the logic for prompt handling. Updates tiktoken library from version 0.8.0 to 0.9.0 for improved tokenization performance. Increases default token limit in TokenBufferMemory to accommodate larger prompt messages. These changes streamline the token management process and leverage the latest improvements in the tiktoken library. Fixes potential token overflow issues and prepares the system for handling larger inputs more efficiently. Relates to internal optimization tasks. Signed-off-by: -LAN- <laipz8200@outlook.com>
170 lines
5.9 KiB
Python
170 lines
5.9 KiB
Python
import json
|
|
import os
|
|
from functools import wraps
|
|
|
|
from flask import abort, request
|
|
from flask_login import current_user # type: ignore
|
|
|
|
from configs import dify_config
|
|
from controllers.console.workspace.error import AccountNotInitializedError
|
|
from models.model import DifySetup
|
|
from services.feature_service import FeatureService, LicenseStatus
|
|
from services.operation_service import OperationService
|
|
|
|
from .error import NotInitValidateError, NotSetupError, UnauthorizedAndForceLogout
|
|
|
|
|
|
def account_initialization_required(view):
|
|
@wraps(view)
|
|
def decorated(*args, **kwargs):
|
|
# check account initialization
|
|
account = current_user
|
|
|
|
if account.status == "uninitialized":
|
|
raise AccountNotInitializedError()
|
|
|
|
return view(*args, **kwargs)
|
|
|
|
return decorated
|
|
|
|
|
|
def only_edition_cloud(view):
|
|
@wraps(view)
|
|
def decorated(*args, **kwargs):
|
|
if dify_config.EDITION != "CLOUD":
|
|
abort(404)
|
|
|
|
return view(*args, **kwargs)
|
|
|
|
return decorated
|
|
|
|
|
|
def only_edition_self_hosted(view):
|
|
@wraps(view)
|
|
def decorated(*args, **kwargs):
|
|
if dify_config.EDITION != "SELF_HOSTED":
|
|
abort(404)
|
|
|
|
return view(*args, **kwargs)
|
|
|
|
return decorated
|
|
|
|
|
|
def cloud_edition_billing_resource_check(resource: str):
|
|
def interceptor(view):
|
|
@wraps(view)
|
|
def decorated(*args, **kwargs):
|
|
features = FeatureService.get_features(current_user.current_tenant_id)
|
|
if features.billing.enabled:
|
|
members = features.members
|
|
apps = features.apps
|
|
vector_space = features.vector_space
|
|
documents_upload_quota = features.documents_upload_quota
|
|
annotation_quota_limit = features.annotation_quota_limit
|
|
if resource == "members" and 0 < members.limit <= members.size:
|
|
abort(403, "The number of members has reached the limit of your subscription.")
|
|
elif resource == "apps" and 0 < apps.limit <= apps.size:
|
|
abort(403, "The number of apps has reached the limit of your subscription.")
|
|
elif resource == "vector_space" and 0 < vector_space.limit <= vector_space.size:
|
|
abort(403, "The capacity of the vector space has reached the limit of your subscription.")
|
|
elif resource == "documents" and 0 < documents_upload_quota.limit <= documents_upload_quota.size:
|
|
# The api of file upload is used in the multiple places,
|
|
# so we need to check the source of the request from datasets
|
|
source = request.args.get("source")
|
|
if source == "datasets":
|
|
abort(403, "The number of documents has reached the limit of your subscription.")
|
|
else:
|
|
return view(*args, **kwargs)
|
|
elif resource == "workspace_custom" and not features.can_replace_logo:
|
|
abort(403, "The workspace custom feature has reached the limit of your subscription.")
|
|
elif resource == "annotation" and 0 < annotation_quota_limit.limit < annotation_quota_limit.size:
|
|
abort(403, "The annotation quota has reached the limit of your subscription.")
|
|
else:
|
|
return view(*args, **kwargs)
|
|
|
|
return view(*args, **kwargs)
|
|
|
|
return decorated
|
|
|
|
return interceptor
|
|
|
|
|
|
def cloud_edition_billing_knowledge_limit_check(resource: str):
|
|
def interceptor(view):
|
|
@wraps(view)
|
|
def decorated(*args, **kwargs):
|
|
features = FeatureService.get_features(current_user.current_tenant_id)
|
|
if features.billing.enabled:
|
|
if resource == "add_segment":
|
|
if features.billing.subscription.plan == "sandbox":
|
|
abort(
|
|
403,
|
|
"To unlock this feature and elevate your Dify experience, please upgrade to a paid plan.",
|
|
)
|
|
else:
|
|
return view(*args, **kwargs)
|
|
|
|
return view(*args, **kwargs)
|
|
|
|
return decorated
|
|
|
|
return interceptor
|
|
|
|
|
|
def cloud_utm_record(view):
|
|
@wraps(view)
|
|
def decorated(*args, **kwargs):
|
|
try:
|
|
features = FeatureService.get_features(current_user.current_tenant_id)
|
|
|
|
if features.billing.enabled:
|
|
utm_info = request.cookies.get("utm_info")
|
|
|
|
if utm_info:
|
|
utm_info_dict: dict = json.loads(utm_info)
|
|
OperationService.record_utm(current_user.current_tenant_id, utm_info_dict)
|
|
except Exception as e:
|
|
pass
|
|
return view(*args, **kwargs)
|
|
|
|
return decorated
|
|
|
|
|
|
def setup_required(view):
|
|
@wraps(view)
|
|
def decorated(*args, **kwargs):
|
|
# check setup
|
|
if dify_config.EDITION == "SELF_HOSTED" and os.environ.get("INIT_PASSWORD") and not DifySetup.query.first():
|
|
raise NotInitValidateError()
|
|
elif dify_config.EDITION == "SELF_HOSTED" and not DifySetup.query.first():
|
|
raise NotSetupError()
|
|
|
|
return view(*args, **kwargs)
|
|
|
|
return decorated
|
|
|
|
|
|
def enterprise_license_required(view):
|
|
@wraps(view)
|
|
def decorated(*args, **kwargs):
|
|
settings = FeatureService.get_system_features()
|
|
if settings.license.status in [LicenseStatus.INACTIVE, LicenseStatus.EXPIRED, LicenseStatus.LOST]:
|
|
raise UnauthorizedAndForceLogout("Your license is invalid. Please contact your administrator.")
|
|
|
|
return view(*args, **kwargs)
|
|
|
|
return decorated
|
|
|
|
|
|
def email_password_login_enabled(view):
|
|
@wraps(view)
|
|
def decorated(*args, **kwargs):
|
|
features = FeatureService.get_system_features()
|
|
if features.enable_email_password_login:
|
|
return view(*args, **kwargs)
|
|
|
|
# otherwise, return 403
|
|
abort(403)
|
|
|
|
return decorated
|