diff --git a/api/controllers/console/billing/billing.py b/api/controllers/console/billing/billing.py index 705f5970dd..436d29df83 100644 --- a/api/controllers/console/billing/billing.py +++ b/api/controllers/console/billing/billing.py @@ -2,6 +2,7 @@ from flask_restx import Resource, reqparse from controllers.console import console_ns from controllers.console.wraps import account_initialization_required, only_edition_cloud, setup_required +from enums.cloud_plan import CloudPlan from libs.login import current_account_with_tenant, login_required from services.billing_service import BillingService @@ -16,7 +17,13 @@ class Subscription(Resource): current_user, current_tenant_id = current_account_with_tenant() parser = ( reqparse.RequestParser() - .add_argument("plan", type=str, required=True, location="args", choices=["professional", "team"]) + .add_argument( + "plan", + type=str, + required=True, + location="args", + choices=[CloudPlan.PROFESSIONAL, CloudPlan.TEAM], + ) .add_argument("interval", type=str, required=True, location="args", choices=["month", "year"]) ) args = parser.parse_args() diff --git a/api/controllers/console/workspace/workspace.py b/api/controllers/console/workspace/workspace.py index f9856df9ea..70d5ef242e 100644 --- a/api/controllers/console/workspace/workspace.py +++ b/api/controllers/console/workspace/workspace.py @@ -21,6 +21,7 @@ from controllers.console.wraps import ( cloud_edition_billing_resource_check, setup_required, ) +from enums.cloud_plan import CloudPlan from extensions.ext_database import db from libs.helper import TimestampField from libs.login import current_account_with_tenant, login_required @@ -83,7 +84,7 @@ class TenantListApi(Resource): "name": tenant.name, "status": tenant.status, "created_at": tenant.created_at, - "plan": features.billing.subscription.plan if features.billing.enabled else "sandbox", + "plan": features.billing.subscription.plan if features.billing.enabled else CloudPlan.SANDBOX, "current": tenant.id == current_tenant_id if current_tenant_id else False, } diff --git a/api/controllers/console/wraps.py b/api/controllers/console/wraps.py index 8572a6dc9b..9b485544db 100644 --- a/api/controllers/console/wraps.py +++ b/api/controllers/console/wraps.py @@ -10,6 +10,7 @@ from flask import abort, request from configs import dify_config from controllers.console.workspace.error import AccountNotInitializedError +from enums.cloud_plan import CloudPlan from extensions.ext_database import db from extensions.ext_redis import redis_client from libs.login import current_account_with_tenant @@ -133,7 +134,7 @@ def cloud_edition_billing_knowledge_limit_check(resource: str): features = FeatureService.get_features(current_tenant_id) if features.billing.enabled: if resource == "add_segment": - if features.billing.subscription.plan == "sandbox": + if features.billing.subscription.plan == CloudPlan.SANDBOX: abort( 403, "To unlock this feature and elevate your Dify experience, please upgrade to a paid plan.", diff --git a/api/controllers/service_api/wraps.py b/api/controllers/service_api/wraps.py index 638ab528f3..fe1e2c419b 100644 --- a/api/controllers/service_api/wraps.py +++ b/api/controllers/service_api/wraps.py @@ -13,6 +13,7 @@ from sqlalchemy import select, update from sqlalchemy.orm import Session from werkzeug.exceptions import Forbidden, NotFound, Unauthorized +from enums.cloud_plan import CloudPlan from extensions.ext_database import db from extensions.ext_redis import redis_client from libs.datetime_utils import naive_utc_now @@ -138,7 +139,7 @@ def cloud_edition_billing_knowledge_limit_check(resource: str, api_token_type: s features = FeatureService.get_features(api_token.tenant_id) if features.billing.enabled: if resource == "add_segment": - if features.billing.subscription.plan == "sandbox": + if features.billing.subscription.plan == CloudPlan.SANDBOX: raise Forbidden( "To unlock this feature and elevate your Dify experience, please upgrade to a paid plan." ) diff --git a/api/core/app/apps/pipeline/pipeline_generator.py b/api/core/app/apps/pipeline/pipeline_generator.py index f8bfbce37a..c36d34f571 100644 --- a/api/core/app/apps/pipeline/pipeline_generator.py +++ b/api/core/app/apps/pipeline/pipeline_generator.py @@ -40,6 +40,7 @@ from core.workflow.repositories.draft_variable_repository import DraftVariableSa from core.workflow.repositories.workflow_execution_repository import WorkflowExecutionRepository from core.workflow.repositories.workflow_node_execution_repository import WorkflowNodeExecutionRepository from core.workflow.variable_loader import DUMMY_VARIABLE_LOADER, VariableLoader +from enums.cloud_plan import CloudPlan from extensions.ext_database import db from extensions.ext_redis import redis_client from libs.flask_utils import preserve_flask_contexts @@ -255,7 +256,7 @@ class PipelineGenerator(BaseAppGenerator): json_text = json.dumps(text) upload_file = FileService(db.engine).upload_text(json_text, name, user.id, dataset.tenant_id) features = FeatureService.get_features(dataset.tenant_id) - if features.billing.enabled and features.billing.subscription.plan == "sandbox": + if features.billing.enabled and features.billing.subscription.plan == CloudPlan.SANDBOX: tenant_pipeline_task_key = f"tenant_pipeline_task:{dataset.tenant_id}" tenant_self_pipeline_task_queue = f"tenant_self_pipeline_task_queue:{dataset.tenant_id}" diff --git a/api/enums/__init__.py b/api/enums/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/api/enums/cloud_plan.py b/api/enums/cloud_plan.py new file mode 100644 index 0000000000..927cff5471 --- /dev/null +++ b/api/enums/cloud_plan.py @@ -0,0 +1,15 @@ +from enum import StrEnum, auto + + +class CloudPlan(StrEnum): + """ + Enum representing user plan types in the cloud platform. + + SANDBOX: Free/default plan with limited features + PROFESSIONAL: Professional paid plan + TEAM: Team collaboration paid plan + """ + + SANDBOX = auto() + PROFESSIONAL = auto() + TEAM = auto() diff --git a/api/schedule/clean_messages.py b/api/schedule/clean_messages.py index 65038dce4d..352a84b592 100644 --- a/api/schedule/clean_messages.py +++ b/api/schedule/clean_messages.py @@ -7,6 +7,7 @@ from sqlalchemy.exc import SQLAlchemyError import app from configs import dify_config +from enums.cloud_plan import CloudPlan from extensions.ext_database import db from extensions.ext_redis import redis_client from models.model import ( @@ -63,7 +64,7 @@ def clean_messages(): plan = features.billing.subscription.plan else: plan = plan_cache.decode() - if plan == "sandbox": + if plan == CloudPlan.SANDBOX: # clean related message db.session.query(MessageFeedback).where(MessageFeedback.message_id == message.id).delete( synchronize_session=False diff --git a/api/schedule/clean_unused_datasets_task.py b/api/schedule/clean_unused_datasets_task.py index 9efd46ba5d..d9fb6a24f1 100644 --- a/api/schedule/clean_unused_datasets_task.py +++ b/api/schedule/clean_unused_datasets_task.py @@ -9,6 +9,7 @@ from sqlalchemy.exc import SQLAlchemyError import app from configs import dify_config from core.rag.index_processor.index_processor_factory import IndexProcessorFactory +from enums.cloud_plan import CloudPlan from extensions.ext_database import db from extensions.ext_redis import redis_client from models.dataset import Dataset, DatasetAutoDisableLog, DatasetQuery, Document @@ -35,7 +36,7 @@ def clean_unused_datasets_task(): }, { "clean_day": datetime.datetime.now() - datetime.timedelta(days=dify_config.PLAN_PRO_CLEAN_DAY_SETTING), - "plan_filter": "sandbox", + "plan_filter": CloudPlan.SANDBOX, "add_logs": False, }, ] diff --git a/api/schedule/mail_clean_document_notify_task.py b/api/schedule/mail_clean_document_notify_task.py index b70707b17e..d738bf46fa 100644 --- a/api/schedule/mail_clean_document_notify_task.py +++ b/api/schedule/mail_clean_document_notify_task.py @@ -7,6 +7,7 @@ from sqlalchemy import select import app from configs import dify_config +from enums.cloud_plan import CloudPlan from extensions.ext_database import db from extensions.ext_mail import mail from libs.email_i18n import EmailType, get_email_i18n_service @@ -45,7 +46,7 @@ def mail_clean_document_notify_task(): for tenant_id, tenant_dataset_auto_disable_logs in dataset_auto_disable_logs_map.items(): features = FeatureService.get_features(tenant_id) plan = features.billing.subscription.plan - if plan != "sandbox": + if plan != CloudPlan.SANDBOX: knowledge_details = [] # check tenant tenant = db.session.query(Tenant).where(Tenant.id == tenant_id).first() diff --git a/api/services/app_generate_service.py b/api/services/app_generate_service.py index b462ddf236..25ee8223c2 100644 --- a/api/services/app_generate_service.py +++ b/api/services/app_generate_service.py @@ -10,6 +10,7 @@ from core.app.apps.completion.app_generator import CompletionAppGenerator from core.app.apps.workflow.app_generator import WorkflowAppGenerator from core.app.entities.app_invoke_entities import InvokeFrom from core.app.features.rate_limiting import RateLimit +from enums.cloud_plan import CloudPlan from libs.helper import RateLimiter from models.model import Account, App, AppMode, EndUser from models.workflow import Workflow @@ -44,7 +45,7 @@ class AppGenerateService: if dify_config.BILLING_ENABLED: # check if it's free plan limit_info = BillingService.get_info(app_model.tenant_id) - if limit_info["subscription"]["plan"] == "sandbox": + if limit_info["subscription"]["plan"] == CloudPlan.SANDBOX: if cls.system_rate_limiter.is_rate_limited(app_model.tenant_id): raise InvokeRateLimitError( "Rate limit exceeded, please upgrade your plan " diff --git a/api/services/billing_service.py b/api/services/billing_service.py index a6851d2638..1650bad0f5 100644 --- a/api/services/billing_service.py +++ b/api/services/billing_service.py @@ -4,6 +4,7 @@ from typing import Literal import httpx from tenacity import retry, retry_if_exception_type, stop_before_delay, wait_fixed +from enums.cloud_plan import CloudPlan from extensions.ext_database import db from extensions.ext_redis import redis_client from libs.helper import RateLimiter @@ -31,7 +32,7 @@ class BillingService: return { "limit": knowledge_rate_limit.get("limit", 10), - "subscription_plan": knowledge_rate_limit.get("subscription_plan", "sandbox"), + "subscription_plan": knowledge_rate_limit.get("subscription_plan", CloudPlan.SANDBOX), } @classmethod diff --git a/api/services/clear_free_plan_tenant_expired_logs.py b/api/services/clear_free_plan_tenant_expired_logs.py index f8f89d7428..aefc34fcae 100644 --- a/api/services/clear_free_plan_tenant_expired_logs.py +++ b/api/services/clear_free_plan_tenant_expired_logs.py @@ -11,6 +11,7 @@ from sqlalchemy.orm import Session, sessionmaker from configs import dify_config from core.model_runtime.utils.encoders import jsonable_encoder +from enums.cloud_plan import CloudPlan from extensions.ext_database import db from extensions.ext_storage import storage from models.account import Tenant @@ -358,7 +359,7 @@ class ClearFreePlanTenantExpiredLogs: try: if ( not dify_config.BILLING_ENABLED - or BillingService.get_info(tenant_id)["subscription"]["plan"] == "sandbox" + or BillingService.get_info(tenant_id)["subscription"]["plan"] == CloudPlan.SANDBOX ): # only process sandbox tenant cls.process_tenant(flask_app, tenant_id, days, batch) diff --git a/api/services/dataset_service.py b/api/services/dataset_service.py index a3e62544c6..2e255c0a9b 100644 --- a/api/services/dataset_service.py +++ b/api/services/dataset_service.py @@ -22,6 +22,7 @@ from core.model_runtime.entities.model_entities import ModelType from core.rag.index_processor.constant.built_in_field import BuiltInField from core.rag.index_processor.constant.index_type import IndexType from core.rag.retrieval.retrieval_methods import RetrievalMethod +from enums.cloud_plan import CloudPlan from events.dataset_event import dataset_was_deleted from events.document_event import document_was_deleted from extensions.ext_database import db @@ -1042,7 +1043,7 @@ class DatasetService: assert isinstance(current_user, Account) assert current_user.current_tenant_id is not None features = FeatureService.get_features(current_user.current_tenant_id) - if not features.billing.enabled or features.billing.subscription.plan == "sandbox": + if not features.billing.enabled or features.billing.subscription.plan == CloudPlan.SANDBOX: return { "document_ids": [], "count": 0, @@ -1438,7 +1439,7 @@ class DocumentService: count = len(website_info.urls) batch_upload_limit = int(dify_config.BATCH_UPLOAD_LIMIT) - if features.billing.subscription.plan == "sandbox" and count > 1: + if features.billing.subscription.plan == CloudPlan.SANDBOX and count > 1: raise ValueError("Your current plan does not support batch upload, please upgrade your plan.") if count > batch_upload_limit: raise ValueError(f"You have reached the batch upload limit of {batch_upload_limit}.") @@ -1727,7 +1728,7 @@ class DocumentService: # count = len(website_info.urls) # type: ignore # batch_upload_limit = int(dify_config.BATCH_UPLOAD_LIMIT) - # if features.billing.subscription.plan == "sandbox" and count > 1: + # if features.billing.subscription.plan == CloudPlan.SANDBOX and count > 1: # raise ValueError("Your current plan does not support batch upload, please upgrade your plan.") # if count > batch_upload_limit: # raise ValueError(f"You have reached the batch upload limit of {batch_upload_limit}.") @@ -2196,7 +2197,7 @@ class DocumentService: website_info = knowledge_config.data_source.info_list.website_info_list if website_info: count = len(website_info.urls) - if features.billing.subscription.plan == "sandbox" and count > 1: + if features.billing.subscription.plan == CloudPlan.SANDBOX and count > 1: raise ValueError("Your current plan does not support batch upload, please upgrade your plan.") batch_upload_limit = int(dify_config.BATCH_UPLOAD_LIMIT) if count > batch_upload_limit: diff --git a/api/services/feature_service.py b/api/services/feature_service.py index 148442f76e..44bea57769 100644 --- a/api/services/feature_service.py +++ b/api/services/feature_service.py @@ -3,12 +3,13 @@ from enum import StrEnum from pydantic import BaseModel, ConfigDict, Field from configs import dify_config +from enums.cloud_plan import CloudPlan from services.billing_service import BillingService from services.enterprise.enterprise_service import EnterpriseService class SubscriptionModel(BaseModel): - plan: str = "sandbox" + plan: str = CloudPlan.SANDBOX interval: str = "" @@ -186,7 +187,7 @@ class FeatureService: knowledge_rate_limit.enabled = True limit_info = BillingService.get_knowledge_rate_limit(tenant_id) knowledge_rate_limit.limit = limit_info.get("limit", 10) - knowledge_rate_limit.subscription_plan = limit_info.get("subscription_plan", "sandbox") + knowledge_rate_limit.subscription_plan = limit_info.get("subscription_plan", CloudPlan.SANDBOX) return knowledge_rate_limit @classmethod @@ -240,7 +241,7 @@ class FeatureService: features.billing.subscription.interval = billing_info["subscription"]["interval"] features.education.activated = billing_info["subscription"].get("education", False) - if features.billing.subscription.plan != "sandbox": + if features.billing.subscription.plan != CloudPlan.SANDBOX: features.webapp_copyright_enabled = True else: features.is_allow_transfer_workspace = False diff --git a/api/tasks/document_indexing_task.py b/api/tasks/document_indexing_task.py index 012ae8f706..07f469de0e 100644 --- a/api/tasks/document_indexing_task.py +++ b/api/tasks/document_indexing_task.py @@ -6,6 +6,7 @@ from celery import shared_task from configs import dify_config from core.indexing_runner import DocumentIsPausedError, IndexingRunner +from enums.cloud_plan import CloudPlan from extensions.ext_database import db from libs.datetime_utils import naive_utc_now from models.dataset import Dataset, Document @@ -38,7 +39,7 @@ def document_indexing_task(dataset_id: str, document_ids: list): vector_space = features.vector_space count = len(document_ids) batch_upload_limit = int(dify_config.BATCH_UPLOAD_LIMIT) - if features.billing.subscription.plan == "sandbox" and count > 1: + if features.billing.subscription.plan == CloudPlan.SANDBOX and count > 1: raise ValueError("Your current plan does not support batch upload, please upgrade your plan.") if count > batch_upload_limit: raise ValueError(f"You have reached the batch upload limit of {batch_upload_limit}.") diff --git a/api/tasks/duplicate_document_indexing_task.py b/api/tasks/duplicate_document_indexing_task.py index 2020179cd9..6492e356a3 100644 --- a/api/tasks/duplicate_document_indexing_task.py +++ b/api/tasks/duplicate_document_indexing_task.py @@ -8,6 +8,7 @@ from sqlalchemy import select from configs import dify_config from core.indexing_runner import DocumentIsPausedError, IndexingRunner from core.rag.index_processor.index_processor_factory import IndexProcessorFactory +from enums.cloud_plan import CloudPlan from extensions.ext_database import db from libs.datetime_utils import naive_utc_now from models.dataset import Dataset, Document, DocumentSegment @@ -41,7 +42,7 @@ def duplicate_document_indexing_task(dataset_id: str, document_ids: list): if features.billing.enabled: vector_space = features.vector_space count = len(document_ids) - if features.billing.subscription.plan == "sandbox" and count > 1: + if features.billing.subscription.plan == CloudPlan.SANDBOX and count > 1: raise ValueError("Your current plan does not support batch upload, please upgrade your plan.") batch_upload_limit = int(dify_config.BATCH_UPLOAD_LIMIT) if count > batch_upload_limit: diff --git a/api/tests/test_containers_integration_tests/services/test_app_generate_service.py b/api/tests/test_containers_integration_tests/services/test_app_generate_service.py index 9386687a04..8b8739d557 100644 --- a/api/tests/test_containers_integration_tests/services/test_app_generate_service.py +++ b/api/tests/test_containers_integration_tests/services/test_app_generate_service.py @@ -5,6 +5,7 @@ import pytest from faker import Faker from core.app.entities.app_invoke_entities import InvokeFrom +from enums.cloud_plan import CloudPlan from models.model import EndUser from models.workflow import Workflow from services.app_generate_service import AppGenerateService @@ -32,7 +33,7 @@ class TestAppGenerateService: patch("services.app_generate_service.dify_config") as mock_dify_config, ): # Setup default mock returns for billing service - mock_billing_service.get_info.return_value = {"subscription": {"plan": "sandbox"}} + mock_billing_service.get_info.return_value = {"subscription": {"plan": CloudPlan.SANDBOX}} # Setup default mock returns for workflow service mock_workflow_service_instance = mock_workflow_service.return_value @@ -430,7 +431,7 @@ class TestAppGenerateService: # Setup billing service mock for sandbox plan mock_external_service_dependencies["billing_service"].get_info.return_value = { - "subscription": {"plan": "sandbox"} + "subscription": {"plan": CloudPlan.SANDBOX} } # Set BILLING_ENABLED to True for this test @@ -461,7 +462,7 @@ class TestAppGenerateService: # Setup billing service mock for sandbox plan mock_external_service_dependencies["billing_service"].get_info.return_value = { - "subscription": {"plan": "sandbox"} + "subscription": {"plan": CloudPlan.SANDBOX} } # Set BILLING_ENABLED to True for this test diff --git a/api/tests/test_containers_integration_tests/services/test_feature_service.py b/api/tests/test_containers_integration_tests/services/test_feature_service.py index 8bd5440411..40380b09d2 100644 --- a/api/tests/test_containers_integration_tests/services/test_feature_service.py +++ b/api/tests/test_containers_integration_tests/services/test_feature_service.py @@ -3,6 +3,7 @@ from unittest.mock import patch import pytest from faker import Faker +from enums.cloud_plan import CloudPlan from services.feature_service import FeatureModel, FeatureService, KnowledgeRateLimitModel, SystemFeatureModel @@ -173,7 +174,7 @@ class TestFeatureService: # Set mock return value inside the patch context mock_external_service_dependencies["billing_service"].get_info.return_value = { "enabled": True, - "subscription": {"plan": "sandbox", "interval": "monthly", "education": False}, + "subscription": {"plan": CloudPlan.SANDBOX, "interval": "monthly", "education": False}, "members": {"size": 1, "limit": 3}, "apps": {"size": 1, "limit": 5}, "vector_space": {"size": 1, "limit": 2}, @@ -189,7 +190,7 @@ class TestFeatureService: result = FeatureService.get_features(tenant_id) # Assert: Verify sandbox-specific limitations - assert result.billing.subscription.plan == "sandbox" + assert result.billing.subscription.plan == CloudPlan.SANDBOX assert result.education.activated is False # Verify sandbox limitations diff --git a/api/tests/test_containers_integration_tests/tasks/test_document_indexing_task.py b/api/tests/test_containers_integration_tests/tasks/test_document_indexing_task.py index 448f6da5ec..1329bba082 100644 --- a/api/tests/test_containers_integration_tests/tasks/test_document_indexing_task.py +++ b/api/tests/test_containers_integration_tests/tasks/test_document_indexing_task.py @@ -3,6 +3,7 @@ from unittest.mock import MagicMock, patch import pytest from faker import Faker +from enums.cloud_plan import CloudPlan from extensions.ext_database import db from models import Account, Tenant, TenantAccountJoin, TenantAccountRole from models.dataset import Dataset, Document @@ -197,7 +198,7 @@ class TestDocumentIndexingTask: # Configure billing features mock_external_service_dependencies["features"].billing.enabled = billing_enabled if billing_enabled: - mock_external_service_dependencies["features"].billing.subscription.plan = "sandbox" + mock_external_service_dependencies["features"].billing.subscription.plan = CloudPlan.SANDBOX mock_external_service_dependencies["features"].vector_space.limit = 100 mock_external_service_dependencies["features"].vector_space.size = 50 @@ -442,7 +443,7 @@ class TestDocumentIndexingTask: ) # Configure sandbox plan with batch limit - mock_external_service_dependencies["features"].billing.subscription.plan = "sandbox" + mock_external_service_dependencies["features"].billing.subscription.plan = CloudPlan.SANDBOX # Create more documents than sandbox plan allows (limit is 1) fake = Faker()