dify/api/controllers/console/wraps.py

# -*- coding:utf-8 -*-
from functools import wraps

from flask import current_app, abort
from flask_login import current_user

from controllers.console.workspace.error import AccountNotInitializedError
from services.billing_service import BillingService


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 current_app.config['EDITION'] != 'CLOUD':
            abort(404)

        return view(*args, **kwargs)

    return decorated


def only_edition_self_hosted(view):
    @wraps(view)
    def decorated(*args, **kwargs):
        if current_app.config['EDITION'] != 'SELF_HOSTED':
            abort(404)

        return view(*args, **kwargs)

    return decorated


def cloud_edition_billing_resource_check(resource: str,
                                         error_msg: str = "You have reached the limit of your subscription."):
    def interceptor(view):
        @wraps(view)
        def decorated(*args, **kwargs):
            if current_app.config['EDITION'] == 'CLOUD':
                tenant_id = current_user.current_tenant_id
                billing_info = BillingService.get_info(tenant_id)
                members = billing_info['members']
                apps = billing_info['apps']
                vector_space = billing_info['vector_space']
                annotation_quota_limit = billing_info['annotation_quota_limit']

                if resource == 'members' and 0 < members['limit'] <= members['size']:
                    abort(403, error_msg)
                elif resource == 'apps' and 0 < apps['limit'] <= apps['size']:
                    abort(403, error_msg)
                elif resource == 'vector_space' and 0 < vector_space['limit'] <= vector_space['size']:
                    abort(403, error_msg)
                elif resource == 'workspace_custom' and not billing_info['can_replace_logo']:
                    abort(403, error_msg)
                elif resource == 'annotation' and 0 < annotation_quota_limit['limit'] <= annotation_quota_limit['size']:
                    abort(403, error_msg)
                else:
                    return view(*args, **kwargs)

            return view(*args, **kwargs)
        return decorated
    return interceptor