feat: add branding support

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

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

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

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,

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