Underground-Digital/Workflow-Engine

1from datetime import timedelta2 3from celery import Celery, Task4from celery.schedules import crontab5from flask import Flask6 7from configs import dify_config8 9 10def init_app(app: Flask) -> Celery:11    class FlaskTask(Task):12        def __call__(self, *args: object, **kwargs: object) -> object:13            with app.app_context():14                return self.run(*args, **kwargs)15 16    broker_transport_options = {}17 18    if dify_config.CELERY_USE_SENTINEL:19        broker_transport_options = {20            "master_name": dify_config.CELERY_SENTINEL_MASTER_NAME,21            "sentinel_kwargs": {22                "socket_timeout": dify_config.CELERY_SENTINEL_SOCKET_TIMEOUT,23            },24        }25 26    celery_app = Celery(27        app.name,28        task_cls=FlaskTask,29        broker=dify_config.CELERY_BROKER_URL,30        backend=dify_config.CELERY_BACKEND,31        task_ignore_result=True,32    )33 34    # Add SSL options to the Celery configuration35    ssl_options = {36        "ssl_cert_reqs": None,37        "ssl_ca_certs": None,38        "ssl_certfile": None,39        "ssl_keyfile": None,40    }41 42    celery_app.conf.update(43        result_backend=dify_config.CELERY_RESULT_BACKEND,44        broker_transport_options=broker_transport_options,45        broker_connection_retry_on_startup=True,46    )47 48    if dify_config.BROKER_USE_SSL:49        celery_app.conf.update(50            broker_use_ssl=ssl_options,  # Add the SSL options to the broker configuration51        )52 53    celery_app.set_default()54    app.extensions["celery"] = celery_app55 56    imports = [57        "schedule.clean_embedding_cache_task",58        "schedule.clean_unused_datasets_task",59        "schedule.create_tidb_serverless_task",60        "schedule.update_tidb_serverless_status_task",61    ]62    day = dify_config.CELERY_BEAT_SCHEDULER_TIME63    beat_schedule = {64        "clean_embedding_cache_task": {65            "task": "schedule.clean_embedding_cache_task.clean_embedding_cache_task",66            "schedule": timedelta(days=day),67        },68        "clean_unused_datasets_task": {69            "task": "schedule.clean_unused_datasets_task.clean_unused_datasets_task",70            "schedule": timedelta(days=day),71        },72        "create_tidb_serverless_task": {73            "task": "schedule.create_tidb_serverless_task.create_tidb_serverless_task",74            "schedule": crontab(minute="0", hour="*"),75        },76        "update_tidb_serverless_status_task": {77            "task": "schedule.update_tidb_serverless_status_task.update_tidb_serverless_status_task",78            "schedule": crontab(minute="30", hour="*"),79        },80    }81    celery_app.conf.update(beat_schedule=beat_schedule, imports=imports)82 83    return celery_app84