Feat: Add documents limitation (#2662)

api/services/dataset_service.py

@@ -37,7 +37,7 @@ from services.errors.account import NoPermissionError
 from services.errors.dataset import DatasetNameDuplicateError
 from services.errors.document import DocumentIndexingError
 from services.errors.file import FileNotExistsError
-from services.feature_service import FeatureService
+from services.feature_service import FeatureModel, FeatureService
 from services.vector_service import VectorService
 from tasks.clean_notion_document_task import clean_notion_document_task
 from tasks.deal_dataset_vector_index_task import deal_dataset_vector_index_task
@@ -469,6 +469,9 @@ class DocumentService:
                 batch_upload_limit = int(current_app.config['BATCH_UPLOAD_LIMIT'])
                 if count > batch_upload_limit:
                     raise ValueError(f"You have reached the batch upload limit of {batch_upload_limit}.")
+
+                DocumentService.check_documents_upload_quota(count, features)
+
         # if dataset is empty, update dataset data_source_type
         if not dataset.data_source_type:
             dataset.data_source_type = document_data["data_source"]["type"]
@@ -619,6 +622,12 @@ class DocumentService:
 
         return documents, batch
 
+    @staticmethod
+    def check_documents_upload_quota(count: int, features: FeatureModel):
+        can_upload_size = features.documents_upload_quota.limit - features.documents_upload_quota.size
+        if count > can_upload_size:
+            raise ValueError(f'You have reached the limit of your subscription. Only {can_upload_size} documents can be uploaded.')
+
     @staticmethod
     def build_document(dataset: Dataset, process_rule_id: str, data_source_type: str, document_form: str,
                        document_language: str, data_source_info: dict, created_from: str, position: int,
@@ -763,6 +772,8 @@ class DocumentService:
             if count > batch_upload_limit:
                 raise ValueError(f"You have reached the batch upload limit of {batch_upload_limit}.")
 
+            DocumentService.check_documents_upload_quota(count, features)
+
         embedding_model = None
         dataset_collection_binding_id = None
         retrieval_model = None

api/services/feature_service.py

@@ -25,6 +25,7 @@ class FeatureModel(BaseModel):
     apps: LimitationModel = LimitationModel(size=0, limit=10)
     vector_space: LimitationModel = LimitationModel(size=0, limit=5)
     annotation_quota_limit: LimitationModel = LimitationModel(size=0, limit=10)
+    documents_upload_quota: LimitationModel = LimitationModel(size=0, limit=50)
     docs_processing: str = 'standard'
     can_replace_logo: bool = False
 
@@ -63,6 +64,9 @@ class FeatureService:
         features.vector_space.size = billing_info['vector_space']['size']
         features.vector_space.limit = billing_info['vector_space']['limit']
 
+        features.documents_upload_quota.size = billing_info['documents_upload_quota']['size']
+        features.documents_upload_quota.limit = billing_info['documents_upload_quota']['limit']
+
         features.annotation_quota_limit.size = billing_info['annotation_quota_limit']['size']
         features.annotation_quota_limit.limit = billing_info['annotation_quota_limit']['limit']