From f8026c1db8f5363ed64fb0af5b03783450c4ce72 Mon Sep 17 00:00:00 2001 From: zhangx1n Date: Thu, 24 Apr 2025 13:08:35 +0800 Subject: [PATCH] refactor: add `enabled` field to LicenseLimitationModel --- api/controllers/console/auth/login.py | 12 +---- api/controllers/console/workspace/members.py | 7 +-- api/services/account_service.py | 9 ++-- api/services/enterprise/base.py | 1 - api/services/feature_service.py | 53 +++++++++++++++----- 5 files changed, 48 insertions(+), 34 deletions(-) diff --git a/api/controllers/console/auth/login.py b/api/controllers/console/auth/login.py index 58de237fbf..fdcc47d680 100644 --- a/api/controllers/console/auth/login.py +++ b/api/controllers/console/auth/login.py @@ -91,11 +91,7 @@ class LoginApi(Resource): if len(tenants) == 0: system_features = FeatureService.get_system_features() - if ( - system_features.is_allow_create_workspace - and system_features.license.workspaces.limit != 0 - and system_features.license.workspaces.limit - system_features.license.workspaces.size <= 0 - ): + if system_features.is_allow_create_workspace and not system_features.license.workspaces.is_available(): raise WorkspacesLimitExceeded() else: return { @@ -209,11 +205,7 @@ class EmailCodeLoginApi(Resource): tenant = TenantService.get_join_tenants(account) if not tenant: workspaces = FeatureService.get_system_features().license.workspaces - if ( - FeatureService.get_system_features().license.product_id == "DIFY_ENTERPRISE_STANDARD" - and workspaces.limit != 0 # if limit == 0 means unlimited - and workspaces.limit - workspaces.size <= 0 - ): + if not workspaces.is_available(): raise WorkspacesLimitExceeded() if not FeatureService.get_system_features().is_allow_create_workspace: raise NotAllowedCreateWorkspace() diff --git a/api/controllers/console/workspace/members.py b/api/controllers/console/workspace/members.py index ab55a08161..3628aa1670 100644 --- a/api/controllers/console/workspace/members.py +++ b/api/controllers/console/workspace/members.py @@ -58,11 +58,8 @@ class MemberInviteEmailApi(Resource): console_web_url = dify_config.CONSOLE_WEB_URL workspace_members = FeatureService.get_features(tenant_id=inviter.current_tenant.id).workspace_members - if ( - FeatureService.get_system_features().license.product_id == "DIFY_ENTERPRISE_STANDARD" - and workspace_members.limit != 0 # if limit == 0, it means unlimited - and len(invitee_emails) > workspace_members.limit - workspace_members.size - ): + + if not workspace_members.is_available(len(invitee_emails)): raise WorkspaceMembersLimitExceeded() for invitee_email in invitee_emails: diff --git a/api/services/account_service.py b/api/services/account_service.py index c6b8d0e9b7..19885860d9 100644 --- a/api/services/account_service.py +++ b/api/services/account_service.py @@ -53,9 +53,10 @@ from services.errors.workspace import WorkSpaceNotAllowedCreateError, Workspaces from services.feature_service import FeatureService from tasks.delete_account_task import delete_account_task from tasks.mail_account_deletion_task import send_account_deletion_verification_code +from tasks.mail_email_code_login import send_email_code_login_mail_task from tasks.mail_invite_member_task import send_invite_member_mail_task from tasks.mail_reset_password_task import send_reset_password_mail_task -from tasks.mail_email_code_login import send_email_code_login_mail_task + class TokenPair(BaseModel): access_token: str @@ -599,11 +600,7 @@ class TenantService: raise WorkSpaceNotAllowedCreateError() workspaces = FeatureService.get_system_features().license.workspaces - if ( - FeatureService.get_system_features().license.product_id == "DIFY_ENTERPRISE_STANDARD" - and workspaces.limit != 0 # if limit == 0, it means unlimited - and workspaces.limit - workspaces.size <= 0 - ): + if not workspaces.is_available(): raise WorkspacesLimitExceededError() if name: diff --git a/api/services/enterprise/base.py b/api/services/enterprise/base.py index 6120de5351..3c3f970444 100644 --- a/api/services/enterprise/base.py +++ b/api/services/enterprise/base.py @@ -16,6 +16,5 @@ class EnterpriseRequest: def send_request(cls, method, endpoint, json=None, params=None): headers = {"Content-Type": "application/json", "Enterprise-Api-Secret-Key": cls.secret_key} url = f"{cls.base_url}{endpoint}" - print("url: ", url) response = requests.request(method, url, json=json, params=params, headers=headers, proxies=cls.proxies) return response.json() diff --git a/api/services/feature_service.py b/api/services/feature_service.py index c87b28e9dd..1bbf46dd71 100644 --- a/api/services/feature_service.py +++ b/api/services/feature_service.py @@ -1,6 +1,6 @@ from enum import StrEnum -from pydantic import BaseModel, ConfigDict +from pydantic import BaseModel, ConfigDict, Field from configs import dify_config from services.billing_service import BillingService @@ -22,6 +22,32 @@ class LimitationModel(BaseModel): limit: int = 0 +class LicenseLimitationModel(BaseModel): + """ + - enabled: whether this limit is enforced + - size: current usage count + - limit: maximum allowed count; 0 means unlimited + """ + + enabled: bool = Field(False, description="Whether this limit is currently active") + size: int = Field(0, description="Number of resources already consumed") + limit: int = Field(0, description="Maximum number of resources allowed; 0 means no limit") + + def is_available(self, required: int = 1) -> bool: + """ + Determine whether the requested amount can be allocated. + + Returns True if: + - this limit is not active, or + - the limit is zero (unlimited), or + - there is enough remaining quota. + """ + if not self.enabled or self.limit == 0: + return True + + return (self.limit - self.size) >= required + + class LicenseStatus(StrEnum): NONE = "none" INACTIVE = "inactive" @@ -35,6 +61,8 @@ class LicenseModel(BaseModel): status: LicenseStatus = LicenseStatus.NONE expired_at: str = "" product_id: str = "" + workspaces: LicenseLimitationModel = LicenseLimitationModel(enabled=False, size=0, limit=0) + class BrandingModel(BaseModel): enabled: bool = False @@ -68,7 +96,7 @@ class FeatureModel(BaseModel): model_load_balancing_enabled: bool = False dataset_operator_enabled: bool = False webapp_copyright_enabled: bool = False - workspace_members: LimitationModel = LimitationModel(size=0, limit=0) + workspace_members: LicenseLimitationModel = LicenseLimitationModel(enabled=False, size=0, limit=0) # pydantic configs model_config = ConfigDict(protected_namespaces=()) @@ -84,7 +112,6 @@ class SystemFeatureModel(BaseModel): is_allow_create_workspace: bool = False is_email_setup: bool = False license: LicenseModel = LicenseModel() - workspaces: LimitationModel = LimitationModel(size=0, limit=0) branding: BrandingModel = BrandingModel() webapp_auth: WebAppAuthModel = WebAppAuthModel() @@ -101,7 +128,7 @@ class FeatureService: if dify_config.ENTERPRISE_ENABLED: features.webapp_copyright_enabled = True - cls._fulfill_params_from_workspaces_quota(features, tenant_id) + cls._fulfill_params_from_workspace_info(features, tenant_id) return features @@ -134,10 +161,12 @@ class FeatureService: features.dataset_operator_enabled = dify_config.DATASET_OPERATOR_ENABLED @classmethod - def _fulfill_params_from_workspaces_quota(cls, features: FeatureModel, tenant_id: str): - workspace_members_quota = EnterpriseService.get_workspace_info(tenant_id)["WorkspaceMembersQuota"] - features.workspace_members.limit = workspace_members_quota["limit"] - features.workspace_members.size = workspace_members_quota["used"] + def _fulfill_params_from_workspace_info(cls, features: FeatureModel, tenant_id: str): + workspace_info = EnterpriseService.get_workspace_info(tenant_id) + if "WorkspaceMembers" in workspace_info: + features.workspace_members.size = workspace_info["WorkspaceMembers"]["used"] + features.workspace_members.limit = workspace_info["WorkspaceMembers"]["limit"] + features.workspace_members.enabled = workspace_info["WorkspaceMembers"]["enabled"] @classmethod def _fulfill_params_from_billing_api(cls, features: FeatureModel, tenant_id: str): @@ -229,7 +258,7 @@ class FeatureService: if "productId" in license_info: features.license.product_id = license_info["productId"] - if "WorkspacesQuota" in enterprise_info: - features.workspaces.limit =enterprise_info["WorkspacesQuota"]["limit"] - features.workspaces.size = enterprise_info["WorkspacesQuota"]["used"] - + if "workspaces" in license_info: + features.license.workspaces.enabled = license_info["workspaces"]["enabled"] + features.license.workspaces.limit = license_info["workspaces"]["limit"] + features.license.workspaces.size = license_info["workspaces"]["used"]