Underground-Digital/Workflow-Engine

1import json2import os3from functools import wraps4 5from flask import abort, request6from flask_login import current_user7 8from configs import dify_config9from controllers.console.workspace.error import AccountNotInitializedError10from models.model import DifySetup11from services.feature_service import FeatureService12from services.operation_service import OperationService13 14from .error import NotInitValidateError, NotSetupError15 16 17def account_initialization_required(view):18    @wraps(view)19    def decorated(*args, **kwargs):20        # check account initialization21        account = current_user22 23        if account.status == "uninitialized":24            raise AccountNotInitializedError()25 26        return view(*args, **kwargs)27 28    return decorated29 30 31def only_edition_cloud(view):32    @wraps(view)33    def decorated(*args, **kwargs):34        if dify_config.EDITION != "CLOUD":35            abort(404)36 37        return view(*args, **kwargs)38 39    return decorated40 41 42def only_edition_self_hosted(view):43    @wraps(view)44    def decorated(*args, **kwargs):45        if dify_config.EDITION != "SELF_HOSTED":46            abort(404)47 48        return view(*args, **kwargs)49 50    return decorated51 52 53def cloud_edition_billing_resource_check(resource: str):54    def interceptor(view):55        @wraps(view)56        def decorated(*args, **kwargs):57            features = FeatureService.get_features(current_user.current_tenant_id)58            if features.billing.enabled:59                members = features.members60                apps = features.apps61                vector_space = features.vector_space62                documents_upload_quota = features.documents_upload_quota63                annotation_quota_limit = features.annotation_quota_limit64                if resource == "members" and 0 < members.limit <= members.size:65                    abort(403, "The number of members has reached the limit of your subscription.")66                elif resource == "apps" and 0 < apps.limit <= apps.size:67                    abort(403, "The number of apps has reached the limit of your subscription.")68                elif resource == "vector_space" and 0 < vector_space.limit <= vector_space.size:69                    abort(403, "The capacity of the vector space has reached the limit of your subscription.")70                elif resource == "documents" and 0 < documents_upload_quota.limit <= documents_upload_quota.size:71                    # The api of file upload is used in the multiple places,72                    # so we need to check the source of the request from datasets73                    source = request.args.get("source")74                    if source == "datasets":75                        abort(403, "The number of documents has reached the limit of your subscription.")76                    else:77                        return view(*args, **kwargs)78                elif resource == "workspace_custom" and not features.can_replace_logo:79                    abort(403, "The workspace custom feature has reached the limit of your subscription.")80                elif resource == "annotation" and 0 < annotation_quota_limit.limit < annotation_quota_limit.size:81                    abort(403, "The annotation quota has reached the limit of your subscription.")82                else:83                    return view(*args, **kwargs)84 85            return view(*args, **kwargs)86 87        return decorated88 89    return interceptor90 91 92def cloud_edition_billing_knowledge_limit_check(resource: str):93    def interceptor(view):94        @wraps(view)95        def decorated(*args, **kwargs):96            features = FeatureService.get_features(current_user.current_tenant_id)97            if features.billing.enabled:98                if resource == "add_segment":99                    if features.billing.subscription.plan == "sandbox":100                        abort(101                            403,102                            "To unlock this feature and elevate your Dify experience, please upgrade to a paid plan.",103                        )104                else:105                    return view(*args, **kwargs)106 107            return view(*args, **kwargs)108 109        return decorated110 111    return interceptor112 113 114def cloud_utm_record(view):115    @wraps(view)116    def decorated(*args, **kwargs):117        try:118            features = FeatureService.get_features(current_user.current_tenant_id)119 120            if features.billing.enabled:121                utm_info = request.cookies.get("utm_info")122 123                if utm_info:124                    utm_info = json.loads(utm_info)125                    OperationService.record_utm(current_user.current_tenant_id, utm_info)126        except Exception as e:127            pass128        return view(*args, **kwargs)129 130    return decorated131 132 133def setup_required(view):134    @wraps(view)135    def decorated(*args, **kwargs):136        # check setup137        if dify_config.EDITION == "SELF_HOSTED" and os.environ.get("INIT_PASSWORD") and not DifySetup.query.first():138            raise NotInitValidateError()139        elif dify_config.EDITION == "SELF_HOSTED" and not DifySetup.query.first():140            raise NotSetupError()141 142        return view(*args, **kwargs)143 144    return decorated145