From fdb1e649d4495aa136fd4f498da777d5811e01b0 Mon Sep 17 00:00:00 2001 From: GareArc Date: Tue, 11 Mar 2025 07:14:52 -0400 Subject: [PATCH] feat: add branding support --- api/controllers/console/wraps.py | 11 +++++++++++ api/services/feature_service.py | 13 +++++++++++++ api/tasks/mail_email_code_login.py | 11 +++++++++-- api/tasks/mail_invite_member_task.py | 10 ++++++++-- api/tasks/mail_reset_password_task.py | 11 +++++++++-- 5 files changed, 50 insertions(+), 6 deletions(-) diff --git a/api/controllers/console/wraps.py b/api/controllers/console/wraps.py index 111db7ccf2..510ea242a9 100644 --- a/api/controllers/console/wraps.py +++ b/api/controllers/console/wraps.py @@ -39,6 +39,17 @@ def only_edition_cloud(view): return decorated +def only_enterprise_edition(view): + @wraps(view) + def decorated(*args, **kwargs): + if not dify_config.ENTERPRISE_ENABLED: + abort(404) + + return view(*args, **kwargs) + + return decorated + + def only_edition_self_hosted(view): @wraps(view) def decorated(*args, **kwargs): diff --git a/api/services/feature_service.py b/api/services/feature_service.py index b9261d19d7..4ef73bae06 100644 --- a/api/services/feature_service.py +++ b/api/services/feature_service.py @@ -36,6 +36,13 @@ class LicenseModel(BaseModel): expired_at: str = "" +class BrandingModel(BaseModel): + enabled: bool = False + login_page_logo: str = ("",) + workspace_logo: str = ("",) + favicon: str = ("",) + + class FeatureModel(BaseModel): billing: BillingModel = BillingModel() members: LimitationModel = LimitationModel(size=0, limit=1) @@ -65,6 +72,8 @@ class SystemFeatureModel(BaseModel): is_allow_create_workspace: bool = False is_email_setup: bool = False license: LicenseModel = LicenseModel() + is_enterprise: bool = False + branding: BrandingModel = BrandingModel() class FeatureService: @@ -87,6 +96,7 @@ class FeatureService: if dify_config.ENTERPRISE_ENABLED: system_features.enable_web_sso_switch_component = True + system_features.is_enterprise = True cls._fulfill_params_from_enterprise(system_features) @@ -172,6 +182,9 @@ class FeatureService: if "is_allow_create_workspace" in enterprise_info: features.is_allow_create_workspace = enterprise_info["is_allow_create_workspace"] + if "branding" in enterprise_info: + features.branding = enterprise_info["branding"] + if "license" in enterprise_info: license_info = enterprise_info["license"] diff --git a/api/tasks/mail_email_code_login.py b/api/tasks/mail_email_code_login.py index 5dc935548f..ac4b8f054f 100644 --- a/api/tasks/mail_email_code_login.py +++ b/api/tasks/mail_email_code_login.py @@ -5,6 +5,7 @@ import click from celery import shared_task # type: ignore from flask import render_template +from configs import dify_config from extensions.ext_mail import mail @@ -25,10 +26,16 @@ def send_email_code_login_mail_task(language: str, to: str, code: str): # send email code login mail using different languages try: if language == "zh-Hans": - html_content = render_template("email_code_login_mail_template_zh-CN.html", to=to, code=code) + template = "email_code_login_mail_template_zh-CN.html" + if dify_config.ENTERPRISE_ENABLED: + template = "without-brand/email_code_login_mail_template_zh-CN_enterprise.html" + html_content = render_template(template, to=to, code=code) mail.send(to=to, subject="邮箱验证码", html=html_content) else: - html_content = render_template("email_code_login_mail_template_en-US.html", to=to, code=code) + template = "email_code_login_mail_template_en-US.html" + if dify_config.ENTERPRISE_ENABLED: + template = "without-brand/email_code_login_mail_template_en-US_enterprise.html" + html_content = render_template(template, to=to, code=code) mail.send(to=to, subject="Email Code", html=html_content) end_at = time.perf_counter() diff --git a/api/tasks/mail_invite_member_task.py b/api/tasks/mail_invite_member_task.py index 3094527fd4..b4fcf7caef 100644 --- a/api/tasks/mail_invite_member_task.py +++ b/api/tasks/mail_invite_member_task.py @@ -33,8 +33,11 @@ def send_invite_member_mail_task(language: str, to: str, token: str, inviter_nam try: url = f"{dify_config.CONSOLE_WEB_URL}/activate?token={token}" if language == "zh-Hans": + template = "invite_member_mail_template_zh-CN.html" + if dify_config.ENTERPRISE_ENABLED: + template = "without-brand/invite_member_mail_template_zh-CN.html" html_content = render_template( - "invite_member_mail_template_zh-CN.html", + template, to=to, inviter_name=inviter_name, workspace_name=workspace_name, @@ -42,8 +45,11 @@ def send_invite_member_mail_task(language: str, to: str, token: str, inviter_nam ) mail.send(to=to, subject="立即加入 Dify 工作空间", html=html_content) else: + template = "invite_member_mail_template_en-US.html" + if dify_config.ENTERPRISE_ENABLED: + template = "without-brand/invite_member_mail_template_en-US.html" html_content = render_template( - "invite_member_mail_template_en-US.html", + template, to=to, inviter_name=inviter_name, workspace_name=workspace_name, diff --git a/api/tasks/mail_reset_password_task.py b/api/tasks/mail_reset_password_task.py index d5be94431b..e434bc30c0 100644 --- a/api/tasks/mail_reset_password_task.py +++ b/api/tasks/mail_reset_password_task.py @@ -5,6 +5,7 @@ import click from celery import shared_task # type: ignore from flask import render_template +from configs import dify_config from extensions.ext_mail import mail @@ -25,10 +26,16 @@ def send_reset_password_mail_task(language: str, to: str, code: str): # send reset password mail using different languages try: if language == "zh-Hans": - html_content = render_template("reset_password_mail_template_zh-CN.html", to=to, code=code) + template = "reset_password_mail_template_zh-CN.html" + if dify_config.ENTERPRISE_ENABLED: + template = "without-brand/reset_password_mail_template_zh-CN.html" + html_content = render_template(template, to=to, code=code) mail.send(to=to, subject="设置您的 Dify 密码", html=html_content) else: - html_content = render_template("reset_password_mail_template_en-US.html", to=to, code=code) + template = "reset_password_mail_template_en-US.html" + if dify_config.ENTERPRISE_ENABLED: + template = "without-brand/reset_password_mail_template_en-US.html" + html_content = render_template(template, to=to, code=code) mail.send(to=to, subject="Set Your Dify Password", html=html_content) end_at = time.perf_counter()