feat: enhancement celery configuration (#32145)

api/configs/middleware/__init__.py

@@ -259,11 +259,20 @@ class CeleryConfig(DatabaseConfig):
         description="Password of the Redis Sentinel master.",
         default=None,
     )
+
     CELERY_SENTINEL_SOCKET_TIMEOUT: PositiveFloat | None = Field(
         description="Timeout for Redis Sentinel socket operations in seconds.",
         default=0.1,
     )
 
+    CELERY_TASK_ANNOTATIONS: dict[str, Any] | None = Field(
+        description=(
+            "Annotations for Celery tasks as a JSON mapping of task name -> options "
+            "(for example, rate limits or other task-specific settings)."
+        ),
+        default=None,
+    )
+
     @computed_field
     def CELERY_RESULT_BACKEND(self) -> str | None:
         if self.CELERY_BACKEND in ("database", "rabbitmq"):

api/extensions/ext_celery.py

@@ -80,8 +80,14 @@ def init_app(app: DifyApp) -> Celery:
         worker_hijack_root_logger=False,
         timezone=pytz.timezone(dify_config.LOG_TZ or "UTC"),
         task_ignore_result=True,
+        task_annotations=dify_config.CELERY_TASK_ANNOTATIONS,
     )
 
+    if dify_config.CELERY_BACKEND == "redis":
+        celery_app.conf.update(
+            result_backend_transport_options=broker_transport_options,
+        )
+
     # Apply SSL configuration if enabled
     ssl_options = _get_celery_ssl_options()
     if ssl_options:

docker/.env.example

@@ -387,6 +387,8 @@ CELERY_USE_SENTINEL=false
 CELERY_SENTINEL_MASTER_NAME=
 CELERY_SENTINEL_PASSWORD=
 CELERY_SENTINEL_SOCKET_TIMEOUT=0.1
+# e.g. {"tasks.add": {"rate_limit": "10/s"}}
+CELERY_TASK_ANNOTATIONS=null
 
 # ------------------------------
 # CORS Configuration

docker/docker-compose.yaml

@@ -106,6 +106,7 @@ x-shared-env: &shared-api-worker-env
   CELERY_SENTINEL_MASTER_NAME: ${CELERY_SENTINEL_MASTER_NAME:-}
   CELERY_SENTINEL_PASSWORD: ${CELERY_SENTINEL_PASSWORD:-}
   CELERY_SENTINEL_SOCKET_TIMEOUT: ${CELERY_SENTINEL_SOCKET_TIMEOUT:-0.1}
+  CELERY_TASK_ANNOTATIONS: ${CELERY_TASK_ANNOTATIONS:-null}
   WEB_API_CORS_ALLOW_ORIGINS: ${WEB_API_CORS_ALLOW_ORIGINS:-*}
   CONSOLE_CORS_ALLOW_ORIGINS: ${CONSOLE_CORS_ALLOW_ORIGINS:-*}
   COOKIE_DOMAIN: ${COOKIE_DOMAIN:-}