feat: billing enhancement 20231204 (#1691)

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

api/controllers/service_api/dataset/document.py

@@ -11,7 +11,7 @@ from controllers.service_api import api
 from controllers.service_api.app.error import ProviderNotInitializeError
 from controllers.service_api.dataset.error import ArchivedDocumentImmutableError, DocumentIndexingError, \
     NoFileUploadedError, TooManyFilesError
-from controllers.service_api.wraps import DatasetApiResource
+from controllers.service_api.wraps import DatasetApiResource, cloud_edition_billing_resource_check
 from libs.login import current_user
 from core.model_providers.error import ProviderTokenNotInitError
 from extensions.ext_database import db
@@ -24,6 +24,7 @@ from services.file_service import FileService
 class DocumentAddByTextApi(DatasetApiResource):
     """Resource for documents."""
 
+    @cloud_edition_billing_resource_check('vector_space', 'dataset')
     def post(self, tenant_id, dataset_id):
         """Create document by text."""
         parser = reqparse.RequestParser()
@@ -88,6 +89,7 @@ class DocumentAddByTextApi(DatasetApiResource):
 class DocumentUpdateByTextApi(DatasetApiResource):
     """Resource for update documents."""
 
+    @cloud_edition_billing_resource_check('vector_space', 'dataset')
     def post(self, tenant_id, dataset_id, document_id):
         """Update document by text."""
         parser = reqparse.RequestParser()
@@ -147,6 +149,7 @@ class DocumentUpdateByTextApi(DatasetApiResource):
 
 class DocumentAddByFileApi(DatasetApiResource):
     """Resource for documents."""
+    @cloud_edition_billing_resource_check('vector_space', 'dataset')
     def post(self, tenant_id, dataset_id):
         """Create document by upload file."""
         args = {}
@@ -212,6 +215,7 @@ class DocumentAddByFileApi(DatasetApiResource):
 class DocumentUpdateByFileApi(DatasetApiResource):
     """Resource for update documents."""
 
+    @cloud_edition_billing_resource_check('vector_space', 'dataset')
     def post(self, tenant_id, dataset_id, document_id):
         """Update document by upload file."""
         args = {}

api/controllers/service_api/dataset/segment.py

@@ -3,7 +3,7 @@ from flask_restful import reqparse, marshal
 from werkzeug.exceptions import NotFound
 from controllers.service_api import api
 from controllers.service_api.app.error import ProviderNotInitializeError
-from controllers.service_api.wraps import DatasetApiResource
+from controllers.service_api.wraps import DatasetApiResource, cloud_edition_billing_resource_check
 from core.model_providers.error import ProviderTokenNotInitError, LLMBadRequestError
 from core.model_providers.model_factory import ModelFactory
 from extensions.ext_database import db
@@ -14,6 +14,8 @@ from services.dataset_service import DatasetService, DocumentService, SegmentSer
 
 class SegmentApi(DatasetApiResource):
     """Resource for segments."""
+
+    @cloud_edition_billing_resource_check('vector_space', 'dataset')
     def post(self, tenant_id, dataset_id, document_id):
         """Create single segment."""
         # check dataset
@@ -144,6 +146,7 @@ class DatasetSegmentApi(DatasetApiResource):
         SegmentService.delete_segment(segment, document, dataset)
         return {'result': 'success'}, 200
 
+    @cloud_edition_billing_resource_check('vector_space', 'dataset')
     def post(self, tenant_id, dataset_id, document_id, segment_id):
         # check dataset
         dataset_id = str(dataset_id)

api/controllers/service_api/wraps.py

@@ -11,6 +11,7 @@ from libs.login import _get_user
 from extensions.ext_database import db
 from models.account import Tenant, TenantAccountJoin, Account
 from models.model import ApiToken, App
+from services.billing_service import BillingService
 
 
 def validate_app_token(view=None):
@@ -40,6 +41,33 @@ def validate_app_token(view=None):
     return decorator
 
 
+def cloud_edition_billing_resource_check(resource: str,
+                                         api_token_type: str,
+                                         error_msg: str = "You have reached the limit of your subscription."):
+    def interceptor(view):
+        def decorated(*args, **kwargs):
+            if current_app.config['EDITION'] == 'CLOUD':
+                api_token = validate_and_get_api_token(api_token_type)
+                billing_info = BillingService.get_info(api_token.tenant_id)
+
+                members = billing_info['members']
+                apps = billing_info['apps']
+                vector_space = billing_info['vector_space']
+
+                if resource == 'members' and 0 < members['limit'] <= members['size']:
+                    raise Unauthorized(error_msg)
+                elif resource == 'apps' and 0 < apps['limit'] <= apps['size']:
+                    raise Unauthorized(error_msg)
+                elif resource == 'vector_space' and 0 < vector_space['limit'] <= vector_space['size']:
+                    raise Unauthorized(error_msg)
+                else:
+                    return view(*args, **kwargs)
+
+            return view(*args, **kwargs)
+        return decorated
+    return interceptor
+
+
 def validate_dataset_token(view=None):
     def decorator(view):
         @wraps(view)