estel/dify
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

feat: billing enhancement 20231204 (#1691)

Browse Source 
Co-authored-by: jyong <jyong@dify.ai>
This commit is contained in:
Garfield Dai
2023-12-05 16:53:55 +08:00
committed by GitHub
parent cb3a55dae6
commit 7b8a10f3ea
12 changed files with 86 additions and 81 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
api/controllers/console/app/app.py
View File

@@ -12,7 +12,7 @@ from constants.model_template import model_templates, demo_model_templates
from controllers.console import api
from controllers.console.app.error import AppNotFoundError, ProviderNotInitializeError
from controllers.console.setup import setup_required
from controllers.console.wraps import account_initialization_required
from controllers.console.wraps import account_initialization_required, cloud_edition_billing_resource_check
from core.model_providers.error import ProviderTokenNotInitError, LLMBadRequestError
from core.model_providers.model_factory import ModelFactory
from core.model_providers.model_provider_factory import ModelProviderFactory
@@ -57,6 +57,7 @@ class AppListApi(Resource):
@login_required
@account_initialization_required
@marshal_with(app_detail_fields)
@cloud_edition_billing_resource_check('apps')
def post(self):
"""Create app"""
parser = reqparse.RequestParser()

29
api/controllers/console/datasets/datasets_document.py
View File

@@ -16,7 +16,7 @@ from controllers.console.app.error import ProviderNotInitializeError, ProviderQu
from controllers.console.datasets.error import DocumentAlreadyFinishedError, InvalidActionError, DocumentIndexingError, \
InvalidMetadataError, ArchivedDocumentImmutableError
from controllers.console.setup import setup_required
from controllers.console.wraps import account_initialization_required
from controllers.console.wraps import account_initialization_required, cloud_edition_billing_resource_check
from core.indexing_runner import IndexingRunner
from core.model_providers.error import ProviderTokenNotInitError, QuotaExceededError, ModelCurrentlyNotSupportError, \
LLMBadRequestError
@@ -194,6 +194,7 @@ class DatasetDocumentListApi(Resource):
@login_required
@account_initialization_required
@marshal_with(documents_and_batch_fields)
@cloud_edition_billing_resource_check('vector_space')
def post(self, dataset_id):
dataset_id = str(dataset_id)
@@ -252,6 +253,7 @@ class DatasetInitApi(Resource):
@login_required
@account_initialization_required
@marshal_with(dataset_and_document_fields)
@cloud_edition_billing_resource_check('vector_space')
def post(self):
# The role of the current user in the ta table must be admin or owner
if current_user.current_tenant.current_role not in ['admin', 'owner']:
@@ -693,6 +695,7 @@ class DocumentStatusApi(DocumentResource):
@setup_required
@login_required
@account_initialization_required
@cloud_edition_billing_resource_check('vector_space')
def patch(self, dataset_id, document_id, action):
dataset_id = str(dataset_id)
document_id = str(document_id)
@@ -770,14 +773,6 @@ class DocumentStatusApi(DocumentResource):
if not document.archived:
raise InvalidActionError('Document is not archived.')
# check document limit
if current_app.config['EDITION'] == 'CLOUD':
documents_count = DocumentService.get_tenant_documents_count()
total_count = documents_count + 1
tenant_document_count = int(current_app.config['TENANT_DOCUMENT_COUNT'])
if total_count > tenant_document_count:
raise ValueError(f"All your documents have overed limit {tenant_document_count}.")
document.archived = False
document.archived_at = None
document.archived_by = None
@@ -856,21 +851,6 @@ class DocumentRecoverApi(DocumentResource):
return {'result': 'success'}, 204
class DocumentLimitApi(DocumentResource):
@setup_required
@login_required
@account_initialization_required
def get(self):
"""get document limit"""
documents_count = DocumentService.get_tenant_documents_count()
tenant_document_count = int(current_app.config['TENANT_DOCUMENT_COUNT'])
return {
'documents_count': documents_count,
'documents_limit': tenant_document_count
}, 200
api.add_resource(GetProcessRuleApi, '/datasets/process-rule')
api.add_resource(DatasetDocumentListApi,
'/datasets/<uuid:dataset_id>/documents')
@@ -896,4 +876,3 @@ api.add_resource(DocumentStatusApi,
'/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/status/<string:action>')
api.add_resource(DocumentPauseApi, '/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/processing/pause')
api.add_resource(DocumentRecoverApi, '/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/processing/resume')
api.add_resource(DocumentLimitApi, '/datasets/limit')

6
api/controllers/console/datasets/datasets_segments.py
View File

@@ -11,7 +11,7 @@ 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
from controllers.console.wraps import account_initialization_required, cloud_edition_billing_resource_check
from core.model_providers.error import LLMBadRequestError, ProviderTokenNotInitError
from core.model_providers.model_factory import ModelFactory
from libs.login import login_required
@@ -114,6 +114,7 @@ class DatasetDocumentSegmentApi(Resource):
@setup_required
@login_required
@account_initialization_required
@cloud_edition_billing_resource_check('vector_space')
def patch(self, dataset_id, segment_id, action):
dataset_id = str(dataset_id)
dataset = DatasetService.get_dataset(dataset_id)
@@ -200,6 +201,7 @@ class DatasetDocumentSegmentAddApi(Resource):
@setup_required
@login_required
@account_initialization_required
@cloud_edition_billing_resource_check('vector_space')
def post(self, dataset_id, document_id):
# check dataset
dataset_id = str(dataset_id)
@@ -250,6 +252,7 @@ class DatasetDocumentSegmentUpdateApi(Resource):
@setup_required
@login_required
@account_initialization_required
@cloud_edition_billing_resource_check('vector_space')
def patch(self, dataset_id, document_id, segment_id):
# check dataset
dataset_id = str(dataset_id)
@@ -344,6 +347,7 @@ class DatasetDocumentSegmentBatchImportApi(Resource):
@setup_required
@login_required
@account_initialization_required
@cloud_edition_billing_resource_check('vector_space')
def post(self, dataset_id, document_id):
# check dataset
dataset_id = str(dataset_id)

2
api/controllers/console/explore/installed_app.py
View File

@@ -14,6 +14,7 @@ from extensions.ext_database import db
from fields.installed_app_fields import installed_app_list_fields
from models.model import App, InstalledApp, RecommendedApp
from services.account_service import TenantService
from controllers.console.wraps import cloud_edition_billing_resource_check
class InstalledAppsListApi(Resource):
@@ -47,6 +48,7 @@ class InstalledAppsListApi(Resource):
@login_required
@account_initialization_required
@cloud_edition_billing_resource_check('apps')
def post(self):
parser = reqparse.RequestParser()
parser.add_argument('app_id', type=str, required=True, help='Invalid app_id')

3
api/controllers/console/workspace/members.py
View File

@@ -7,7 +7,7 @@ from flask_restful import Resource, reqparse, marshal_with, abort, fields, marsh
import services
from controllers.console import api
from controllers.console.setup import setup_required
from controllers.console.wraps import account_initialization_required
from controllers.console.wraps import account_initialization_required, cloud_edition_billing_resource_check
from libs.helper import TimestampField
from extensions.ext_database import db
from models.account import Account, TenantAccountJoin
@@ -47,6 +47,7 @@ class MemberInviteEmailApi(Resource):
@setup_required
@login_required
@account_initialization_required
@cloud_edition_billing_resource_check('members')
def post(self):
parser = reqparse.RequestParser()
parser.add_argument('emails', type=str, required=True, location='json', action='append')

28
api/controllers/console/wraps.py
View File

@@ -5,6 +5,7 @@ from flask import current_app, abort
from flask_login import current_user
from controllers.console.workspace.error import AccountNotInitializedError
from services.billing_service import BillingService
def account_initialization_required(view):
@@ -41,3 +42,30 @@ def only_edition_self_hosted(view):
return view(*args, **kwargs)
return decorated
def cloud_edition_billing_resource_check(resource: str,
error_msg: str = "You have reached the limit of your subscription."):
def interceptor(view):
@wraps(view)
def decorated(*args, **kwargs):
if current_app.config['EDITION'] == 'CLOUD':
tenant_id = current_user.current_tenant_id
billing_info = BillingService.get_info(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']:
abort(403, error_msg)
elif resource == 'apps' and 0 < apps['limit'] <= apps['size']:
abort(403, error_msg)
elif resource == 'vector_space' and 0 < vector_space['limit'] <= vector_space['size']:
abort(403, error_msg)
else:
return view(*args, **kwargs)
return view(*args, **kwargs)
return decorated
return interceptor