Spaces:
Build error
Build error
import json | |
import os | |
from functools import wraps | |
from flask import abort, request | |
from flask_login import current_user | |
from configs import dify_config | |
from controllers.console.workspace.error import AccountNotInitializedError | |
from models.model import DifySetup | |
from services.feature_service import FeatureService | |
from services.operation_service import OperationService | |
from .error import NotInitValidateError, NotSetupError | |
def account_initialization_required(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): | |
def decorated(*args, **kwargs): | |
if dify_config.EDITION != "CLOUD": | |
abort(404) | |
return view(*args, **kwargs) | |
return decorated | |
def only_edition_self_hosted(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): | |
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): | |
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): | |
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 = json.loads(utm_info) | |
OperationService.record_utm(current_user.current_tenant_id, utm_info) | |
except Exception as e: | |
pass | |
return view(*args, **kwargs) | |
return decorated | |
def setup_required(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 | |