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

Co-authored-by: jyong <jyong@dify.ai>

api/controllers/console/datasets/datasets_segments.py

@@ -12,7 +12,11 @@ from controllers.console import api
 from controllers.console.app.error import ProviderNotInitializeError
 from controllers.console.datasets.error import InvalidActionError, NoFileUploadedError, TooManyFilesError
 from controllers.console.setup import setup_required
-from controllers.console.wraps import account_initialization_required, cloud_edition_billing_resource_check
+from controllers.console.wraps import (
+    account_initialization_required,
+    cloud_edition_billing_knowledge_limit_check,
+    cloud_edition_billing_resource_check,
+)
 from core.errors.error import LLMBadRequestError, ProviderTokenNotInitError
 from core.model_manager import ModelManager
 from core.model_runtime.entities.model_entities import ModelType
@@ -207,6 +211,7 @@ class DatasetDocumentSegmentAddApi(Resource):
     @login_required
     @account_initialization_required
     @cloud_edition_billing_resource_check('vector_space')
+    @cloud_edition_billing_knowledge_limit_check('add_segment')
     def post(self, dataset_id, document_id):
         # check dataset
         dataset_id = str(dataset_id)
@@ -357,6 +362,7 @@ class DatasetDocumentSegmentBatchImportApi(Resource):
     @login_required
     @account_initialization_required
     @cloud_edition_billing_resource_check('vector_space')
+    @cloud_edition_billing_knowledge_limit_check('add_segment')
     def post(self, dataset_id, document_id):
         # check dataset
         dataset_id = str(dataset_id)

api/controllers/console/wraps.py

@@ -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