add segment function billing check for SAAS env (#3082)

Co-authored-by: jyong <jyong@dify.ai>
This commit is contained in:
Jyong
2024-04-02 17:55:49 +08:00
committed by GitHub
parent 9c7e99e829
commit e12a0c154c
4 changed files with 62 additions and 9 deletions

View File

@@ -51,14 +51,12 @@ def cloud_edition_billing_resource_check(resource: str,
@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, error_msg)
elif resource == 'apps' and 0 < apps.limit <= apps.size:
@@ -80,7 +78,29 @@ def cloud_edition_billing_resource_check(resource: str,
return view(*args, **kwargs)
return view(*args, **kwargs)
return decorated
return interceptor
def cloud_edition_billing_knowledge_limit_check(resource: str,
error_msg: str = "To unlock this feature and elevate your Dify experience, please upgrade to a paid plan."):
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, error_msg)
else:
return view(*args, **kwargs)
return view(*args, **kwargs)
return decorated
return interceptor
@@ -99,4 +119,5 @@ def cloud_utm_record(view):
except Exception as e:
pass
return view(*args, **kwargs)
return decorated