Commit 2f2388af authored by Garfield Dai's avatar Garfield Dai

feat: update billing for datasets creation.

parent 5992c9cc
......@@ -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 FeatureService, FeatureModel
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
......
......@@ -25,7 +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=20)
documents_upload_quota: LimitationModel = LimitationModel(size=0, limit=50)
docs_processing: str = 'standard'
can_replace_logo: bool = False
......
......@@ -16,7 +16,7 @@ export const ALL_PLANS: Record<Plan, PlanInfo> = {
teamMembers: 1,
buildApps: 10,
vectorSpace: 5,
documentsUploadQuota: 20,
documentsUploadQuota: 50,
documentProcessingPriority: Priority.standard,
logHistory: 30,
customTools: unAvailable,
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment