jprdonnelly/dify
1
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2025-12-19 17:27:16 -05:00
Code Issues Packages Projects Releases Wiki Activity

refactor: replace hardcoded user plan strings with CloudPlan enum (#27675)

Browse Source
This commit is contained in:
-LAN-
2025-11-03 11:51:09 +08:00
committed by GitHub
parent b2b2816ade
commit 2abbc14703
20 changed files with 66 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
api/controllers/console/billing/billing.py
View File

@@ -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()

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

@@ -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,
}

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

@@ -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.",

3
api/controllers/service_api/wraps.py
View File

@@ -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."
)

3
api/core/app/apps/pipeline/pipeline_generator.py
View File

@@ -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}"

0
api/enums/__init__.py Normal file
View File

15
api/enums/cloud_plan.py Normal file
View File

@@ -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()

3
api/schedule/clean_messages.py
View File

@@ -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

3
api/schedule/clean_unused_datasets_task.py
View File

@@ -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,
},
]

3
api/schedule/mail_clean_document_notify_task.py
View File

@@ -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()

3
api/services/app_generate_service.py
View File

@@ -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 "

3
api/services/billing_service.py
View File

@@ -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

3
api/services/clear_free_plan_tenant_expired_logs.py
View File

@@ -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)

9
api/services/dataset_service.py
View File

@@ -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:

7
api/services/feature_service.py
View File

@@ -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

3
api/tasks/document_indexing_task.py
View File

@@ -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}.")

3
api/tasks/duplicate_document_indexing_task.py
View File

@@ -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:

7
api/tests/test_containers_integration_tests/services/test_app_generate_service.py
View File

@@ -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

5
api/tests/test_containers_integration_tests/services/test_feature_service.py
View File

@@ -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

5
api/tests/test_containers_integration_tests/tasks/test_document_indexing_task.py
View File

@@ -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()