Upload 697 files

.gitattributes

@@ -5,3 +5,4 @@
 # them.
 
 *.sh      text eol=lf
+api/tests/integration_tests/model_runtime/assets/audio.mp3 filter=lfs diff=lfs merge=lfs -text

api/docker/entrypoint.sh

@@ -0,0 +1,40 @@
+#!/bin/bash
+
+set -e
+
+if [[ "${MIGRATION_ENABLED}" == "true" ]]; then
+  echo "Running migrations"
+  flask upgrade-db
+fi
+
+if [[ "${MODE}" == "worker" ]]; then
+
+  # Get the number of available CPU cores
+  if [ "${CELERY_AUTO_SCALE,,}" = "true" ]; then
+    # Set MAX_WORKERS to the number of available cores if not specified
+    AVAILABLE_CORES=$(nproc)
+    MAX_WORKERS=${CELERY_MAX_WORKERS:-$AVAILABLE_CORES}
+    MIN_WORKERS=${CELERY_MIN_WORKERS:-1}
+    CONCURRENCY_OPTION="--autoscale=${MAX_WORKERS},${MIN_WORKERS}"
+  else
+    CONCURRENCY_OPTION="-c ${CELERY_WORKER_AMOUNT:-1}"
+  fi
+
+  exec celery -A app.celery worker -P ${CELERY_WORKER_CLASS:-gevent} $CONCURRENCY_OPTION --loglevel ${LOG_LEVEL:-INFO} \
+    -Q ${CELERY_QUEUES:-dataset,mail,ops_trace,app_deletion}
+
+elif [[ "${MODE}" == "beat" ]]; then
+  exec celery -A app.celery beat --loglevel ${LOG_LEVEL:-INFO}
+else
+  if [[ "${DEBUG}" == "true" ]]; then
+    exec flask run --host=${DIFY_BIND_ADDRESS:-0.0.0.0} --port=${DIFY_PORT:-5001} --debug
+  else
+    exec gunicorn \
+      --bind "${DIFY_BIND_ADDRESS:-0.0.0.0}:${DIFY_PORT:-5001}" \
+      --workers ${SERVER_WORKER_AMOUNT:-1} \
+      --worker-class ${SERVER_WORKER_CLASS:-gevent} \
+      --worker-connections ${SERVER_WORKER_CONNECTIONS:-10} \
+      --timeout ${GUNICORN_TIMEOUT:-200} \
+      app:app
+  fi
+fi

api/events/app_event.py

@@ -0,0 +1,13 @@
+from blinker import signal
+
+# sender: app
+app_was_created = signal("app-was-created")
+
+# sender: app, kwargs: app_model_config
+app_model_config_was_updated = signal("app-model-config-was-updated")
+
+# sender: app, kwargs: published_workflow
+app_published_workflow_was_updated = signal("app-published-workflow-was-updated")
+
+# sender: app, kwargs: synced_draft_workflow
+app_draft_workflow_was_synced = signal("app-draft-workflow-was-synced")

api/events/dataset_event.py

@@ -0,0 +1,4 @@
+from blinker import signal
+
+# sender: dataset
+dataset_was_deleted = signal("dataset-was-deleted")

api/events/document_event.py

@@ -0,0 +1,4 @@
+from blinker import signal
+
+# sender: document
+document_was_deleted = signal("document-was-deleted")

api/events/event_handlers/__init__.py

@@ -0,0 +1,10 @@
+from .clean_when_dataset_deleted import handle
+from .clean_when_document_deleted import handle
+from .create_document_index import handle
+from .create_installed_app_when_app_created import handle
+from .create_site_record_when_app_created import handle
+from .deduct_quota_when_message_created import handle
+from .delete_tool_parameters_cache_when_sync_draft_workflow import handle
+from .update_app_dataset_join_when_app_model_config_updated import handle
+from .update_app_dataset_join_when_app_published_workflow_updated import handle
+from .update_provider_last_used_at_when_message_created import handle

api/events/event_handlers/clean_when_dataset_deleted.py

@@ -0,0 +1,15 @@
+from events.dataset_event import dataset_was_deleted
+from tasks.clean_dataset_task import clean_dataset_task
+
+
+@dataset_was_deleted.connect
+def handle(sender, **kwargs):
+    dataset = sender
+    clean_dataset_task.delay(
+        dataset.id,
+        dataset.tenant_id,
+        dataset.indexing_technique,
+        dataset.index_struct,
+        dataset.collection_binding_id,
+        dataset.doc_form,
+    )

api/events/event_handlers/clean_when_document_deleted.py

@@ -0,0 +1,11 @@
+from events.document_event import document_was_deleted
+from tasks.clean_document_task import clean_document_task
+
+
+@document_was_deleted.connect
+def handle(sender, **kwargs):
+    document_id = sender
+    dataset_id = kwargs.get("dataset_id")
+    doc_form = kwargs.get("doc_form")
+    file_id = kwargs.get("file_id")
+    clean_document_task.delay(document_id, dataset_id, doc_form, file_id)

api/events/event_handlers/create_document_index.py

@@ -0,0 +1,49 @@
+import datetime
+import logging
+import time
+
+import click
+from werkzeug.exceptions import NotFound
+
+from core.indexing_runner import DocumentIsPausedError, IndexingRunner
+from events.event_handlers.document_index_event import document_index_created
+from extensions.ext_database import db
+from models.dataset import Document
+
+
+@document_index_created.connect
+def handle(sender, **kwargs):
+    dataset_id = sender
+    document_ids = kwargs.get("document_ids", [])
+    documents = []
+    start_at = time.perf_counter()
+    for document_id in document_ids:
+        logging.info(click.style("Start process document: {}".format(document_id), fg="green"))
+
+        document = (
+            db.session.query(Document)
+            .filter(
+                Document.id == document_id,
+                Document.dataset_id == dataset_id,
+            )
+            .first()
+        )
+
+        if not document:
+            raise NotFound("Document not found")
+
+        document.indexing_status = "parsing"
+        document.processing_started_at = datetime.datetime.now(datetime.UTC).replace(tzinfo=None)
+        documents.append(document)
+        db.session.add(document)
+    db.session.commit()
+
+    try:
+        indexing_runner = IndexingRunner()
+        indexing_runner.run(documents)
+        end_at = time.perf_counter()
+        logging.info(click.style("Processed dataset: {} latency: {}".format(dataset_id, end_at - start_at), fg="green"))
+    except DocumentIsPausedError as ex:
+        logging.info(click.style(str(ex), fg="yellow"))
+    except Exception:
+        pass

api/events/event_handlers/create_installed_app_when_app_created.py

@@ -0,0 +1,16 @@
+from events.app_event import app_was_created
+from extensions.ext_database import db
+from models.model import InstalledApp
+
+
+@app_was_created.connect
+def handle(sender, **kwargs):
+    """Create an installed app when an app is created."""
+    app = sender
+    installed_app = InstalledApp(
+        tenant_id=app.tenant_id,
+        app_id=app.id,
+        app_owner_tenant_id=app.tenant_id,
+    )
+    db.session.add(installed_app)
+    db.session.commit()

api/events/event_handlers/create_site_record_when_app_created.py

@@ -0,0 +1,26 @@
+from events.app_event import app_was_created
+from extensions.ext_database import db
+from models.model import Site
+
+
+@app_was_created.connect
+def handle(sender, **kwargs):
+    """Create site record when an app is created."""
+    app = sender
+    account = kwargs.get("account")
+    if account is not None:
+        site = Site(
+            app_id=app.id,
+            title=app.name,
+            icon_type=app.icon_type,
+            icon=app.icon,
+            icon_background=app.icon_background,
+            default_language=account.interface_language,
+            customize_token_strategy="not_allow",
+            code=Site.generate_code(16),
+            created_by=app.created_by,
+            updated_by=app.updated_by,
+        )
+
+        db.session.add(site)
+        db.session.commit()

api/events/event_handlers/deduct_quota_when_message_created.py

@@ -0,0 +1,53 @@
+from configs import dify_config
+from core.app.entities.app_invoke_entities import AgentChatAppGenerateEntity, ChatAppGenerateEntity
+from core.entities.provider_entities import QuotaUnit
+from events.message_event import message_was_created
+from extensions.ext_database import db
+from models.provider import Provider, ProviderType
+
+
+@message_was_created.connect
+def handle(sender, **kwargs):
+    message = sender
+    application_generate_entity = kwargs.get("application_generate_entity")
+
+    if not isinstance(application_generate_entity, ChatAppGenerateEntity | AgentChatAppGenerateEntity):
+        return
+
+    model_config = application_generate_entity.model_conf
+    provider_model_bundle = model_config.provider_model_bundle
+    provider_configuration = provider_model_bundle.configuration
+
+    if provider_configuration.using_provider_type != ProviderType.SYSTEM:
+        return
+
+    system_configuration = provider_configuration.system_configuration
+
+    quota_unit = None
+    for quota_configuration in system_configuration.quota_configurations:
+        if quota_configuration.quota_type == system_configuration.current_quota_type:
+            quota_unit = quota_configuration.quota_unit
+
+            if quota_configuration.quota_limit == -1:
+                return
+
+            break
+
+    used_quota = None
+    if quota_unit:
+        if quota_unit == QuotaUnit.TOKENS:
+            used_quota = message.message_tokens + message.answer_tokens
+        elif quota_unit == QuotaUnit.CREDITS:
+            used_quota = dify_config.get_model_credits(model_config.model)
+        else:
+            used_quota = 1
+
+    if used_quota is not None and system_configuration.current_quota_type is not None:
+        db.session.query(Provider).filter(
+            Provider.tenant_id == application_generate_entity.app_config.tenant_id,
+            Provider.provider_name == model_config.provider,
+            Provider.provider_type == ProviderType.SYSTEM.value,
+            Provider.quota_type == system_configuration.current_quota_type.value,
+            Provider.quota_limit > Provider.quota_used,
+        ).update({"quota_used": Provider.quota_used + used_quota})
+        db.session.commit()

api/events/event_handlers/delete_tool_parameters_cache_when_sync_draft_workflow.py

@@ -0,0 +1,34 @@
+from core.tools.tool_manager import ToolManager
+from core.tools.utils.configuration import ToolParameterConfigurationManager
+from core.workflow.nodes import NodeType
+from core.workflow.nodes.tool.entities import ToolEntity
+from events.app_event import app_draft_workflow_was_synced
+
+
+@app_draft_workflow_was_synced.connect
+def handle(sender, **kwargs):
+    app = sender
+    synced_draft_workflow = kwargs.get("synced_draft_workflow")
+    if synced_draft_workflow is None:
+        return
+    for node_data in synced_draft_workflow.graph_dict.get("nodes", []):
+        if node_data.get("data", {}).get("type") == NodeType.TOOL.value:
+            try:
+                tool_entity = ToolEntity(**node_data["data"])
+                tool_runtime = ToolManager.get_tool_runtime(
+                    provider_type=tool_entity.provider_type,
+                    provider_id=tool_entity.provider_id,
+                    tool_name=tool_entity.tool_name,
+                    tenant_id=app.tenant_id,
+                )
+                manager = ToolParameterConfigurationManager(
+                    tenant_id=app.tenant_id,
+                    tool_runtime=tool_runtime,
+                    provider_name=tool_entity.provider_name,
+                    provider_type=tool_entity.provider_type,
+                    identity_id=f"WORKFLOW.{app.id}.{node_data.get('id')}",
+                )
+                manager.delete_tool_parameters_cache()
+            except:
+                # tool dose not exist
+                pass

api/events/event_handlers/document_index_event.py

@@ -0,0 +1,4 @@
+from blinker import signal
+
+# sender: document
+document_index_created = signal("document-index-created")

api/events/event_handlers/update_app_dataset_join_when_app_model_config_updated.py

@@ -0,0 +1,68 @@
+from events.app_event import app_model_config_was_updated
+from extensions.ext_database import db
+from models.dataset import AppDatasetJoin
+from models.model import AppModelConfig
+
+
+@app_model_config_was_updated.connect
+def handle(sender, **kwargs):
+    app = sender
+    app_model_config = kwargs.get("app_model_config")
+    if app_model_config is None:
+        return
+
+    dataset_ids = get_dataset_ids_from_model_config(app_model_config)
+
+    app_dataset_joins = db.session.query(AppDatasetJoin).filter(AppDatasetJoin.app_id == app.id).all()
+
+    removed_dataset_ids: set[str] = set()
+    if not app_dataset_joins:
+        added_dataset_ids = dataset_ids
+    else:
+        old_dataset_ids: set[str] = set()
+        old_dataset_ids.update(app_dataset_join.dataset_id for app_dataset_join in app_dataset_joins)
+
+        added_dataset_ids = dataset_ids - old_dataset_ids
+        removed_dataset_ids = old_dataset_ids - dataset_ids
+
+    if removed_dataset_ids:
+        for dataset_id in removed_dataset_ids:
+            db.session.query(AppDatasetJoin).filter(
+                AppDatasetJoin.app_id == app.id, AppDatasetJoin.dataset_id == dataset_id
+            ).delete()
+
+    if added_dataset_ids:
+        for dataset_id in added_dataset_ids:
+            app_dataset_join = AppDatasetJoin(app_id=app.id, dataset_id=dataset_id)
+            db.session.add(app_dataset_join)
+
+    db.session.commit()
+
+
+def get_dataset_ids_from_model_config(app_model_config: AppModelConfig) -> set[str]:
+    dataset_ids: set[str] = set()
+    if not app_model_config:
+        return dataset_ids
+
+    agent_mode = app_model_config.agent_mode_dict
+
+    tools = agent_mode.get("tools", []) or []
+    for tool in tools:
+        if len(list(tool.keys())) != 1:
+            continue
+
+        tool_type = list(tool.keys())[0]
+        tool_config = list(tool.values())[0]
+        if tool_type == "dataset":
+            dataset_ids.add(tool_config.get("id"))
+
+    # get dataset from dataset_configs
+    dataset_configs = app_model_config.dataset_configs_dict
+    datasets = dataset_configs.get("datasets", {}) or {}
+    for dataset in datasets.get("datasets", []) or []:
+        keys = list(dataset.keys())
+        if len(keys) == 1 and keys[0] == "dataset":
+            if dataset["dataset"].get("id"):
+                dataset_ids.add(dataset["dataset"].get("id"))
+
+    return dataset_ids

api/events/event_handlers/update_app_dataset_join_when_app_published_workflow_updated.py

@@ -0,0 +1,67 @@
+from typing import cast
+
+from core.workflow.nodes import NodeType
+from core.workflow.nodes.knowledge_retrieval.entities import KnowledgeRetrievalNodeData
+from events.app_event import app_published_workflow_was_updated
+from extensions.ext_database import db
+from models.dataset import AppDatasetJoin
+from models.workflow import Workflow
+
+
+@app_published_workflow_was_updated.connect
+def handle(sender, **kwargs):
+    app = sender
+    published_workflow = kwargs.get("published_workflow")
+    published_workflow = cast(Workflow, published_workflow)
+
+    dataset_ids = get_dataset_ids_from_workflow(published_workflow)
+    app_dataset_joins = db.session.query(AppDatasetJoin).filter(AppDatasetJoin.app_id == app.id).all()
+
+    removed_dataset_ids: set[str] = set()
+    if not app_dataset_joins:
+        added_dataset_ids = dataset_ids
+    else:
+        old_dataset_ids: set[str] = set()
+        old_dataset_ids.update(app_dataset_join.dataset_id for app_dataset_join in app_dataset_joins)
+
+        added_dataset_ids = dataset_ids - old_dataset_ids
+        removed_dataset_ids = old_dataset_ids - dataset_ids
+
+    if removed_dataset_ids:
+        for dataset_id in removed_dataset_ids:
+            db.session.query(AppDatasetJoin).filter(
+                AppDatasetJoin.app_id == app.id, AppDatasetJoin.dataset_id == dataset_id
+            ).delete()
+
+    if added_dataset_ids:
+        for dataset_id in added_dataset_ids:
+            app_dataset_join = AppDatasetJoin(app_id=app.id, dataset_id=dataset_id)
+            db.session.add(app_dataset_join)
+
+    db.session.commit()
+
+
+def get_dataset_ids_from_workflow(published_workflow: Workflow) -> set[str]:
+    dataset_ids: set[str] = set()
+    graph = published_workflow.graph_dict
+    if not graph:
+        return dataset_ids
+
+    nodes = graph.get("nodes", [])
+
+    # fetch all knowledge retrieval nodes
+    knowledge_retrieval_nodes = [
+        node for node in nodes if node.get("data", {}).get("type") == NodeType.KNOWLEDGE_RETRIEVAL.value
+    ]
+
+    if not knowledge_retrieval_nodes:
+        return dataset_ids
+
+    for node in knowledge_retrieval_nodes:
+        try:
+            node_data = KnowledgeRetrievalNodeData(**node.get("data", {}))
+            dataset_ids.update(dataset_id for dataset_id in node_data.dataset_ids)
+        except Exception as e:
+            continue
+
+    return dataset_ids

api/events/event_handlers/update_provider_last_used_at_when_message_created.py

@@ -0,0 +1,21 @@
+from datetime import UTC, datetime
+
+from core.app.entities.app_invoke_entities import AgentChatAppGenerateEntity, ChatAppGenerateEntity
+from events.message_event import message_was_created
+from extensions.ext_database import db
+from models.provider import Provider
+
+
+@message_was_created.connect
+def handle(sender, **kwargs):
+    message = sender
+    application_generate_entity = kwargs.get("application_generate_entity")
+
+    if not isinstance(application_generate_entity, ChatAppGenerateEntity | AgentChatAppGenerateEntity):
+        return
+
+    db.session.query(Provider).filter(
+        Provider.tenant_id == application_generate_entity.app_config.tenant_id,
+        Provider.provider_name == application_generate_entity.model_conf.provider,
+    ).update({"last_used": datetime.now(UTC).replace(tzinfo=None)})
+    db.session.commit()

api/events/message_event.py

@@ -0,0 +1,4 @@
+from blinker import signal
+
+# sender: message, kwargs: conversation
+message_was_created = signal("message-was-created")

api/events/tenant_event.py

@@ -0,0 +1,7 @@
+from blinker import signal
+
+# sender: tenant
+tenant_was_created = signal("tenant-was-created")
+
+# sender: tenant
+tenant_was_updated = signal("tenant-was-updated")

api/extensions/ext_app_metrics.py

@@ -0,0 +1,67 @@
+import json
+import os
+import threading
+
+from flask import Response
+
+from configs import dify_config
+from dify_app import DifyApp
+
+
+def init_app(app: DifyApp):
+    @app.after_request
+    def after_request(response):
+        """Add Version headers to the response."""
+        response.headers.add("X-Version", dify_config.CURRENT_VERSION)
+        response.headers.add("X-Env", dify_config.DEPLOY_ENV)
+        return response
+
+    @app.route("/health")
+    def health():
+        return Response(
+            json.dumps({"pid": os.getpid(), "status": "ok", "version": dify_config.CURRENT_VERSION}),
+            status=200,
+            content_type="application/json",
+        )
+
+    @app.route("/threads")
+    def threads():
+        num_threads = threading.active_count()
+        threads = threading.enumerate()
+
+        thread_list = []
+        for thread in threads:
+            thread_name = thread.name
+            thread_id = thread.ident
+            is_alive = thread.is_alive()
+
+            thread_list.append(
+                {
+                    "name": thread_name,
+                    "id": thread_id,
+                    "is_alive": is_alive,
+                }
+            )
+
+        return {
+            "pid": os.getpid(),
+            "thread_num": num_threads,
+            "threads": thread_list,
+        }
+
+    @app.route("/db-pool-stat")
+    def pool_stat():
+        from extensions.ext_database import db
+
+        engine = db.engine
+        # TODO: Fix the type error
+        # FIXME maybe its sqlalchemy issue
+        return {
+            "pid": os.getpid(),
+            "pool_size": engine.pool.size(),  # type: ignore
+            "checked_in_connections": engine.pool.checkedin(),  # type: ignore
+            "checked_out_connections": engine.pool.checkedout(),  # type: ignore
+            "overflow_connections": engine.pool.overflow(),  # type: ignore
+            "connection_timeout": engine.pool.timeout(),  # type: ignore
+            "recycle_time": db.engine.pool._recycle,  # type: ignore
+        }

api/extensions/ext_blueprints.py

@@ -0,0 +1,48 @@
+from configs import dify_config
+from dify_app import DifyApp
+
+
+def init_app(app: DifyApp):
+    # register blueprint routers
+
+    from flask_cors import CORS  # type: ignore
+
+    from controllers.console import bp as console_app_bp
+    from controllers.files import bp as files_bp
+    from controllers.inner_api import bp as inner_api_bp
+    from controllers.service_api import bp as service_api_bp
+    from controllers.web import bp as web_bp
+
+    CORS(
+        service_api_bp,
+        allow_headers=["Content-Type", "Authorization", "X-App-Code"],
+        methods=["GET", "PUT", "POST", "DELETE", "OPTIONS", "PATCH"],
+    )
+    app.register_blueprint(service_api_bp)
+
+    CORS(
+        web_bp,
+        resources={r"/*": {"origins": dify_config.WEB_API_CORS_ALLOW_ORIGINS}},
+        supports_credentials=True,
+        allow_headers=["Content-Type", "Authorization", "X-App-Code"],
+        methods=["GET", "PUT", "POST", "DELETE", "OPTIONS", "PATCH"],
+        expose_headers=["X-Version", "X-Env"],
+    )
+
+    app.register_blueprint(web_bp)
+
+    CORS(
+        console_app_bp,
+        resources={r"/*": {"origins": dify_config.CONSOLE_CORS_ALLOW_ORIGINS}},
+        supports_credentials=True,
+        allow_headers=["Content-Type", "Authorization"],
+        methods=["GET", "PUT", "POST", "DELETE", "OPTIONS", "PATCH"],
+        expose_headers=["X-Version", "X-Env"],
+    )
+
+    app.register_blueprint(console_app_bp)
+
+    CORS(files_bp, allow_headers=["Content-Type"], methods=["GET", "PUT", "POST", "DELETE", "OPTIONS", "PATCH"])
+    app.register_blueprint(files_bp)
+
+    app.register_blueprint(inner_api_bp)

api/extensions/ext_celery.py

@@ -0,0 +1,104 @@
+from datetime import timedelta
+
+import pytz
+from celery import Celery, Task  # type: ignore
+from celery.schedules import crontab  # type: ignore
+
+from configs import dify_config
+from dify_app import DifyApp
+
+
+def init_app(app: DifyApp) -> Celery:
+    class FlaskTask(Task):
+        def __call__(self, *args: object, **kwargs: object) -> object:
+            with app.app_context():
+                return self.run(*args, **kwargs)
+
+    broker_transport_options = {}
+
+    if dify_config.CELERY_USE_SENTINEL:
+        broker_transport_options = {
+            "master_name": dify_config.CELERY_SENTINEL_MASTER_NAME,
+            "sentinel_kwargs": {
+                "socket_timeout": dify_config.CELERY_SENTINEL_SOCKET_TIMEOUT,
+            },
+        }
+
+    celery_app = Celery(
+        app.name,
+        task_cls=FlaskTask,
+        broker=dify_config.CELERY_BROKER_URL,
+        backend=dify_config.CELERY_BACKEND,
+        task_ignore_result=True,
+    )
+
+    # Add SSL options to the Celery configuration
+    ssl_options = {
+        "ssl_cert_reqs": None,
+        "ssl_ca_certs": None,
+        "ssl_certfile": None,
+        "ssl_keyfile": None,
+    }
+
+    celery_app.conf.update(
+        result_backend=dify_config.CELERY_RESULT_BACKEND,
+        broker_transport_options=broker_transport_options,
+        broker_connection_retry_on_startup=True,
+        worker_log_format=dify_config.LOG_FORMAT,
+        worker_task_log_format=dify_config.LOG_FORMAT,
+        worker_hijack_root_logger=False,
+        timezone=pytz.timezone(dify_config.LOG_TZ or "UTC"),
+    )
+
+    if dify_config.BROKER_USE_SSL:
+        celery_app.conf.update(
+            broker_use_ssl=ssl_options,  # Add the SSL options to the broker configuration
+        )
+
+    if dify_config.LOG_FILE:
+        celery_app.conf.update(
+            worker_logfile=dify_config.LOG_FILE,
+        )
+
+    celery_app.set_default()
+    app.extensions["celery"] = celery_app
+
+    imports = [
+        "schedule.clean_embedding_cache_task",
+        "schedule.clean_unused_datasets_task",
+        "schedule.create_tidb_serverless_task",
+        "schedule.update_tidb_serverless_status_task",
+        "schedule.clean_messages",
+        "schedule.mail_clean_document_notify_task",
+    ]
+    day = dify_config.CELERY_BEAT_SCHEDULER_TIME
+    beat_schedule = {
+        "clean_embedding_cache_task": {
+            "task": "schedule.clean_embedding_cache_task.clean_embedding_cache_task",
+            "schedule": timedelta(days=day),
+        },
+        "clean_unused_datasets_task": {
+            "task": "schedule.clean_unused_datasets_task.clean_unused_datasets_task",
+            "schedule": timedelta(days=day),
+        },
+        "create_tidb_serverless_task": {
+            "task": "schedule.create_tidb_serverless_task.create_tidb_serverless_task",
+            "schedule": crontab(minute="0", hour="*"),
+        },
+        "update_tidb_serverless_status_task": {
+            "task": "schedule.update_tidb_serverless_status_task.update_tidb_serverless_status_task",
+            "schedule": timedelta(minutes=10),
+        },
+        "clean_messages": {
+            "task": "schedule.clean_messages.clean_messages",
+            "schedule": timedelta(days=day),
+        },
+        # every Monday
+        "mail_clean_document_notify_task": {
+            "task": "schedule.mail_clean_document_notify_task.mail_clean_document_notify_task",
+            "schedule": crontab(minute="0", hour="10", day_of_week="1"),
+        },
+    }
+    celery_app.conf.update(beat_schedule=beat_schedule, imports=imports)
+
+    return celery_app

api/extensions/ext_code_based_extension.py

@@ -0,0 +1,9 @@
+from core.extension.extension import Extension
+from dify_app import DifyApp
+
+
+def init_app(app: DifyApp):
+    code_based_extension.init()
+
+
+code_based_extension = Extension()

api/extensions/ext_commands.py

@@ -0,0 +1,29 @@
+from dify_app import DifyApp
+
+
+def init_app(app: DifyApp):
+    from commands import (
+        add_qdrant_doc_id_index,
+        convert_to_agent_apps,
+        create_tenant,
+        fix_app_site_missing,
+        reset_email,
+        reset_encrypt_key_pair,
+        reset_password,
+        upgrade_db,
+        vdb_migrate,
+    )
+
+    cmds_to_register = [
+        reset_password,
+        reset_email,
+        reset_encrypt_key_pair,
+        vdb_migrate,
+        convert_to_agent_apps,
+        add_qdrant_doc_id_index,
+        create_tenant,
+        upgrade_db,
+        fix_app_site_missing,
+    ]
+    for cmd in cmds_to_register:
+        app.cli.add_command(cmd)

api/extensions/ext_compress.py

@@ -0,0 +1,13 @@
+from configs import dify_config
+from dify_app import DifyApp
+
+
+def is_enabled() -> bool:
+    return dify_config.API_COMPRESSION_ENABLED
+
+
+def init_app(app: DifyApp):
+    from flask_compress import Compress  # type: ignore
+
+    compress = Compress()
+    compress.init_app(app)

api/extensions/ext_database.py

@@ -0,0 +1,6 @@
+from dify_app import DifyApp
+from models import db
+
+
+def init_app(app: DifyApp):
+    db.init_app(app)

api/extensions/ext_hosting_provider.py

@@ -0,0 +1,10 @@
+from core.hosting_configuration import HostingConfiguration
+
+hosting_configuration = HostingConfiguration()
+
+
+from dify_app import DifyApp
+
+
+def init_app(app: DifyApp):
+    hosting_configuration.init_app(app)

api/extensions/ext_import_modules.py

@@ -0,0 +1,5 @@
+from dify_app import DifyApp
+
+
+def init_app(app: DifyApp):
+    from events import event_handlers  # noqa: F401

api/extensions/ext_logging.py

@@ -0,0 +1,71 @@
+import logging
+import os
+import sys
+import uuid
+from logging.handlers import RotatingFileHandler
+
+import flask
+
+from configs import dify_config
+from dify_app import DifyApp
+
+
+def init_app(app: DifyApp):
+    log_handlers: list[logging.Handler] = []
+    log_file = dify_config.LOG_FILE
+    if log_file:
+        log_dir = os.path.dirname(log_file)
+        os.makedirs(log_dir, exist_ok=True)
+        log_handlers.append(
+            RotatingFileHandler(
+                filename=log_file,
+                maxBytes=dify_config.LOG_FILE_MAX_SIZE * 1024 * 1024,
+                backupCount=dify_config.LOG_FILE_BACKUP_COUNT,
+            )
+        )
+
+    # Always add StreamHandler to log to console
+    sh = logging.StreamHandler(sys.stdout)
+    sh.addFilter(RequestIdFilter())
+    log_handlers.append(sh)
+
+    logging.basicConfig(
+        level=dify_config.LOG_LEVEL,
+        format=dify_config.LOG_FORMAT,
+        datefmt=dify_config.LOG_DATEFORMAT,
+        handlers=log_handlers,
+        force=True,
+    )
+    log_tz = dify_config.LOG_TZ
+    if log_tz:
+        from datetime import datetime
+
+        import pytz
+
+        timezone = pytz.timezone(log_tz)
+
+        def time_converter(seconds):
+            return datetime.fromtimestamp(seconds, tz=timezone).timetuple()
+
+        for handler in logging.root.handlers:
+            if handler.formatter:
+                handler.formatter.converter = time_converter
+
+
+def get_request_id():
+    if getattr(flask.g, "request_id", None):
+        return flask.g.request_id
+
+    new_uuid = uuid.uuid4().hex[:10]
+    flask.g.request_id = new_uuid
+
+    return new_uuid
+
+
+class RequestIdFilter(logging.Filter):
+    # This is a logging filter that makes the request ID available for use in
+    # the logging format. Note that we're checking if we're in a request
+    # context, as we may want to log things before Flask is fully loaded.
+    def filter(self, record):
+        record.req_id = get_request_id() if flask.has_request_context() else ""
+        return True

api/extensions/ext_login.py

@@ -0,0 +1,62 @@
+import json
+
+import flask_login  # type: ignore
+from flask import Response, request
+from flask_login import user_loaded_from_request, user_logged_in
+from werkzeug.exceptions import Unauthorized
+
+import contexts
+from dify_app import DifyApp
+from libs.passport import PassportService
+from services.account_service import AccountService
+
+login_manager = flask_login.LoginManager()
+
+
+# Flask-Login configuration
+@login_manager.request_loader
+def load_user_from_request(request_from_flask_login):
+    """Load user based on the request."""
+    if request.blueprint not in {"console", "inner_api"}:
+        return None
+    # Check if the user_id contains a dot, indicating the old format
+    auth_header = request.headers.get("Authorization", "")
+    if not auth_header:
+        auth_token = request.args.get("_token")
+        if not auth_token:
+            raise Unauthorized("Invalid Authorization token.")
+    else:
+        if " " not in auth_header:
+            raise Unauthorized("Invalid Authorization header format. Expected 'Bearer <api-key>' format.")
+        auth_scheme, auth_token = auth_header.split(None, 1)
+        auth_scheme = auth_scheme.lower()
+        if auth_scheme != "bearer":
+            raise Unauthorized("Invalid Authorization header format. Expected 'Bearer <api-key>' format.")
+
+    decoded = PassportService().verify(auth_token)
+    user_id = decoded.get("user_id")
+
+    logged_in_account = AccountService.load_logged_in_account(account_id=user_id)
+    return logged_in_account
+
+
+@user_logged_in.connect
+@user_loaded_from_request.connect
+def on_user_logged_in(_sender, user):
+    """Called when a user logged in."""
+    if user:
+        contexts.tenant_id.set(user.current_tenant_id)
+
+
+@login_manager.unauthorized_handler
+def unauthorized_handler():
+    """Handle unauthorized requests."""
+    return Response(
+        json.dumps({"code": "unauthorized", "message": "Unauthorized."}),
+        status=401,
+        content_type="application/json",
+    )
+
+
+def init_app(app: DifyApp):
+    login_manager.init_app(app)

api/extensions/ext_mail.py

@@ -0,0 +1,97 @@
+import logging
+from typing import Optional
+
+from flask import Flask
+
+from configs import dify_config
+from dify_app import DifyApp
+
+
+class Mail:
+    def __init__(self):
+        self._client = None
+        self._default_send_from = None
+
+    def is_inited(self) -> bool:
+        return self._client is not None
+
+    def init_app(self, app: Flask):
+        mail_type = dify_config.MAIL_TYPE
+        if not mail_type:
+            logging.warning("MAIL_TYPE is not set")
+            return
+
+        if dify_config.MAIL_DEFAULT_SEND_FROM:
+            self._default_send_from = dify_config.MAIL_DEFAULT_SEND_FROM
+
+        match mail_type:
+            case "resend":
+                import resend  # type: ignore
+
+                api_key = dify_config.RESEND_API_KEY
+                if not api_key:
+                    raise ValueError("RESEND_API_KEY is not set")
+
+                api_url = dify_config.RESEND_API_URL
+                if api_url:
+                    resend.api_url = api_url
+
+                resend.api_key = api_key
+                self._client = resend.Emails
+            case "smtp":
+                from libs.smtp import SMTPClient
+
+                if not dify_config.SMTP_SERVER or not dify_config.SMTP_PORT:
+                    raise ValueError("SMTP_SERVER and SMTP_PORT are required for smtp mail type")
+                if not dify_config.SMTP_USE_TLS and dify_config.SMTP_OPPORTUNISTIC_TLS:
+                    raise ValueError("SMTP_OPPORTUNISTIC_TLS is not supported without enabling SMTP_USE_TLS")
+                self._client = SMTPClient(
+                    server=dify_config.SMTP_SERVER,
+                    port=dify_config.SMTP_PORT,
+                    username=dify_config.SMTP_USERNAME or "",
+                    password=dify_config.SMTP_PASSWORD or "",
+                    _from=dify_config.MAIL_DEFAULT_SEND_FROM or "",
+                    use_tls=dify_config.SMTP_USE_TLS,
+                    opportunistic_tls=dify_config.SMTP_OPPORTUNISTIC_TLS,
+                )
+            case _:
+                raise ValueError("Unsupported mail type {}".format(mail_type))
+
+    def send(self, to: str, subject: str, html: str, from_: Optional[str] = None):
+        if not self._client:
+            raise ValueError("Mail client is not initialized")
+
+        if not from_ and self._default_send_from:
+            from_ = self._default_send_from
+
+        if not from_:
+            raise ValueError("mail from is not set")
+
+        if not to:
+            raise ValueError("mail to is not set")
+
+        if not subject:
+            raise ValueError("mail subject is not set")
+
+        if not html:
+            raise ValueError("mail html is not set")
+
+        self._client.send(
+            {
+                "from": from_,
+                "to": to,
+                "subject": subject,
+                "html": html,
+            }
+        )
+
+
+def is_enabled() -> bool:
+    return dify_config.MAIL_TYPE is not None and dify_config.MAIL_TYPE != ""
+
+
+def init_app(app: DifyApp):
+    mail.init_app(app)
+
+
+mail = Mail()

api/extensions/ext_migrate.py

@@ -0,0 +1,9 @@
+from dify_app import DifyApp
+
+
+def init_app(app: DifyApp):
+    import flask_migrate  # type: ignore
+
+    from extensions.ext_database import db
+
+    flask_migrate.Migrate(app, db)

api/extensions/ext_proxy_fix.py

@@ -0,0 +1,9 @@
+from configs import dify_config
+from dify_app import DifyApp
+
+
+def init_app(app: DifyApp):
+    if dify_config.RESPECT_XFORWARD_HEADERS_ENABLED:
+        from werkzeug.middleware.proxy_fix import ProxyFix
+
+        app.wsgi_app = ProxyFix(app.wsgi_app, x_port=1)  # type: ignore

api/extensions/ext_redis.py

@@ -0,0 +1,98 @@
+from typing import Any, Union
+
+import redis
+from redis.cluster import ClusterNode, RedisCluster
+from redis.connection import Connection, SSLConnection
+from redis.sentinel import Sentinel
+
+from configs import dify_config
+from dify_app import DifyApp
+
+
+class RedisClientWrapper:
+    """
+    A wrapper class for the Redis client that addresses the issue where the global
+    `redis_client` variable cannot be updated when a new Redis instance is returned
+    by Sentinel.
+
+    This class allows for deferred initialization of the Redis client, enabling the
+    client to be re-initialized with a new instance when necessary. This is particularly
+    useful in scenarios where the Redis instance may change dynamically, such as during
+    a failover in a Sentinel-managed Redis setup.
+
+    Attributes:
+        _client (redis.Redis): The actual Redis client instance. It remains None until
+                               initialized with the `initialize` method.
+
+    Methods:
+        initialize(client): Initializes the Redis client if it hasn't been initialized already.
+        __getattr__(item): Delegates attribute access to the Redis client, raising an error
+                           if the client is not initialized.
+    """
+
+    def __init__(self):
+        self._client = None
+
+    def initialize(self, client):
+        if self._client is None:
+            self._client = client
+
+    def __getattr__(self, item):
+        if self._client is None:
+            raise RuntimeError("Redis client is not initialized. Call init_app first.")
+        return getattr(self._client, item)
+
+
+redis_client = RedisClientWrapper()
+
+
+def init_app(app: DifyApp):
+    global redis_client
+    connection_class: type[Union[Connection, SSLConnection]] = Connection
+    if dify_config.REDIS_USE_SSL:
+        connection_class = SSLConnection
+
+    redis_params: dict[str, Any] = {
+        "username": dify_config.REDIS_USERNAME,
+        "password": dify_config.REDIS_PASSWORD,
+        "db": dify_config.REDIS_DB,
+        "encoding": "utf-8",
+        "encoding_errors": "strict",
+        "decode_responses": False,
+    }
+
+    if dify_config.REDIS_USE_SENTINEL:
+        assert dify_config.REDIS_SENTINELS is not None, "REDIS_SENTINELS must be set when REDIS_USE_SENTINEL is True"
+        sentinel_hosts = [
+            (node.split(":")[0], int(node.split(":")[1])) for node in dify_config.REDIS_SENTINELS.split(",")
+        ]
+        sentinel = Sentinel(
+            sentinel_hosts,
+            sentinel_kwargs={
+                "socket_timeout": dify_config.REDIS_SENTINEL_SOCKET_TIMEOUT,
+                "username": dify_config.REDIS_SENTINEL_USERNAME,
+                "password": dify_config.REDIS_SENTINEL_PASSWORD,
+            },
+        )
+        master = sentinel.master_for(dify_config.REDIS_SENTINEL_SERVICE_NAME, **redis_params)
+        redis_client.initialize(master)
+    elif dify_config.REDIS_USE_CLUSTERS:
+        assert dify_config.REDIS_CLUSTERS is not None, "REDIS_CLUSTERS must be set when REDIS_USE_CLUSTERS is True"
+        nodes = [
+            ClusterNode(host=node.split(":")[0], port=int(node.split(":")[1]))
+            for node in dify_config.REDIS_CLUSTERS.split(",")
+        ]
+        # FIXME: mypy error here, try to figure out how to fix it
+        redis_client.initialize(RedisCluster(startup_nodes=nodes, password=dify_config.REDIS_CLUSTERS_PASSWORD))  # type: ignore
+    else:
+        redis_params.update(
+            {
+                "host": dify_config.REDIS_HOST,
+                "port": dify_config.REDIS_PORT,
+                "connection_class": connection_class,
+            }
+        )
+        pool = redis.ConnectionPool(**redis_params)
+        redis_client.initialize(redis.Redis(connection_pool=pool))
+
+    app.extensions["redis"] = redis_client

api/extensions/ext_sentry.py

@@ -0,0 +1,40 @@
+from configs import dify_config
+from dify_app import DifyApp
+
+
+def init_app(app: DifyApp):
+    if dify_config.SENTRY_DSN:
+        import openai
+        import sentry_sdk
+        from langfuse import parse_error  # type: ignore
+        from sentry_sdk.integrations.celery import CeleryIntegration
+        from sentry_sdk.integrations.flask import FlaskIntegration
+        from werkzeug.exceptions import HTTPException
+
+        from core.model_runtime.errors.invoke import InvokeRateLimitError
+
+        def before_send(event, hint):
+            if "exc_info" in hint:
+                exc_type, exc_value, tb = hint["exc_info"]
+                if parse_error.defaultErrorResponse in str(exc_value):
+                    return None
+
+            return event
+
+        sentry_sdk.init(
+            dsn=dify_config.SENTRY_DSN,
+            integrations=[FlaskIntegration(), CeleryIntegration()],
+            ignore_errors=[
+                HTTPException,
+                ValueError,
+                FileNotFoundError,
+                openai.APIStatusError,
+                InvokeRateLimitError,
+                parse_error.defaultErrorResponse,
+            ],
+            traces_sample_rate=dify_config.SENTRY_TRACES_SAMPLE_RATE,
+            profiles_sample_rate=dify_config.SENTRY_PROFILES_SAMPLE_RATE,
+            environment=dify_config.DEPLOY_ENV,
+            release=f"dify-{dify_config.CURRENT_VERSION}-{dify_config.COMMIT_SHA}",
+            before_send=before_send,
+        )

api/extensions/ext_set_secretkey.py

@@ -0,0 +1,6 @@
+from configs import dify_config
+from dify_app import DifyApp
+
+
+def init_app(app: DifyApp):
+    app.secret_key = dify_config.SECRET_KEY

api/extensions/ext_storage.py

@@ -0,0 +1,138 @@
+import logging
+from collections.abc import Callable, Generator
+from typing import Literal, Union, overload
+
+from flask import Flask
+
+from configs import dify_config
+from dify_app import DifyApp
+from extensions.storage.base_storage import BaseStorage
+from extensions.storage.storage_type import StorageType
+
+logger = logging.getLogger(__name__)
+
+
+class Storage:
+    def init_app(self, app: Flask):
+        storage_factory = self.get_storage_factory(dify_config.STORAGE_TYPE)
+        with app.app_context():
+            self.storage_runner = storage_factory()
+
+    @staticmethod
+    def get_storage_factory(storage_type: str) -> Callable[[], BaseStorage]:
+        match storage_type:
+            case StorageType.S3:
+                from extensions.storage.aws_s3_storage import AwsS3Storage
+
+                return AwsS3Storage
+            case StorageType.OPENDAL:
+                from extensions.storage.opendal_storage import OpenDALStorage
+
+                return lambda: OpenDALStorage(dify_config.OPENDAL_SCHEME)
+            case StorageType.LOCAL:
+                from extensions.storage.opendal_storage import OpenDALStorage
+
+                return lambda: OpenDALStorage(scheme="fs", root=dify_config.STORAGE_LOCAL_PATH)
+            case StorageType.AZURE_BLOB:
+                from extensions.storage.azure_blob_storage import AzureBlobStorage
+
+                return AzureBlobStorage
+            case StorageType.ALIYUN_OSS:
+                from extensions.storage.aliyun_oss_storage import AliyunOssStorage
+
+                return AliyunOssStorage
+            case StorageType.GOOGLE_STORAGE:
+                from extensions.storage.google_cloud_storage import GoogleCloudStorage
+
+                return GoogleCloudStorage
+            case StorageType.TENCENT_COS:
+                from extensions.storage.tencent_cos_storage import TencentCosStorage
+
+                return TencentCosStorage
+            case StorageType.OCI_STORAGE:
+                from extensions.storage.oracle_oci_storage import OracleOCIStorage
+
+                return OracleOCIStorage
+            case StorageType.HUAWEI_OBS:
+                from extensions.storage.huawei_obs_storage import HuaweiObsStorage
+
+                return HuaweiObsStorage
+            case StorageType.BAIDU_OBS:
+                from extensions.storage.baidu_obs_storage import BaiduObsStorage
+
+                return BaiduObsStorage
+            case StorageType.VOLCENGINE_TOS:
+                from extensions.storage.volcengine_tos_storage import VolcengineTosStorage
+
+                return VolcengineTosStorage
+            case StorageType.SUPBASE:
+                from extensions.storage.supabase_storage import SupabaseStorage
+
+                return SupabaseStorage
+            case _:
+                raise ValueError(f"unsupported storage type {storage_type}")
+
+    def save(self, filename, data):
+        try:
+            self.storage_runner.save(filename, data)
+        except Exception as e:
+            logger.exception(f"Failed to save file {filename}")
+            raise e
+
+    @overload
+    def load(self, filename: str, /, *, stream: Literal[False] = False) -> bytes: ...
+
+    @overload
+    def load(self, filename: str, /, *, stream: Literal[True]) -> Generator: ...
+
+    def load(self, filename: str, /, *, stream: bool = False) -> Union[bytes, Generator]:
+        try:
+            if stream:
+                return self.load_stream(filename)
+            else:
+                return self.load_once(filename)
+        except Exception as e:
+            logger.exception(f"Failed to load file {filename}")
+            raise e
+
+    def load_once(self, filename: str) -> bytes:
+        try:
+            return self.storage_runner.load_once(filename)
+        except Exception as e:
+            logger.exception(f"Failed to load_once file {filename}")
+            raise e
+
+    def load_stream(self, filename: str) -> Generator:
+        try:
+            return self.storage_runner.load_stream(filename)
+        except Exception as e:
+            logger.exception(f"Failed to load_stream file {filename}")
+            raise e
+
+    def download(self, filename, target_filepath):
+        try:
+            self.storage_runner.download(filename, target_filepath)
+        except Exception as e:
+            logger.exception(f"Failed to download file {filename}")
+            raise e
+
+    def exists(self, filename):
+        try:
+            return self.storage_runner.exists(filename)
+        except Exception as e:
+            logger.exception(f"Failed to check file exists {filename}")
+            raise e
+
+    def delete(self, filename):
+        try:
+            return self.storage_runner.delete(filename)
+        except Exception as e:
+            logger.exception(f"Failed to delete file {filename}")
+            raise e
+
+
+storage = Storage()
+
+
+def init_app(app: DifyApp):
+    storage.init_app(app)

api/extensions/ext_timezone.py

@@ -0,0 +1,11 @@
+import os
+import time
+
+from dify_app import DifyApp
+
+
+def init_app(app: DifyApp):
+    os.environ["TZ"] = "UTC"
+    # windows platform not support tzset
+    if hasattr(time, "tzset"):
+        time.tzset()

api/extensions/ext_warnings.py

@@ -0,0 +1,7 @@
+from dify_app import DifyApp
+
+
+def init_app(app: DifyApp):
+    import warnings
+
+    warnings.simplefilter("ignore", ResourceWarning)

api/extensions/storage/aliyun_oss_storage.py

@@ -0,0 +1,54 @@
+import posixpath
+from collections.abc import Generator
+
+import oss2 as aliyun_s3  # type: ignore
+
+from configs import dify_config
+from extensions.storage.base_storage import BaseStorage
+
+
+class AliyunOssStorage(BaseStorage):
+    """Implementation for Aliyun OSS storage."""
+
+    def __init__(self):
+        super().__init__()
+        self.bucket_name = dify_config.ALIYUN_OSS_BUCKET_NAME
+        self.folder = dify_config.ALIYUN_OSS_PATH
+        oss_auth_method = aliyun_s3.Auth
+        region = None
+        if dify_config.ALIYUN_OSS_AUTH_VERSION == "v4":
+            oss_auth_method = aliyun_s3.AuthV4
+            region = dify_config.ALIYUN_OSS_REGION
+        oss_auth = oss_auth_method(dify_config.ALIYUN_OSS_ACCESS_KEY, dify_config.ALIYUN_OSS_SECRET_KEY)
+        self.client = aliyun_s3.Bucket(
+            oss_auth,
+            dify_config.ALIYUN_OSS_ENDPOINT,
+            self.bucket_name,
+            connect_timeout=30,
+            region=region,
+        )
+
+    def save(self, filename, data):
+        self.client.put_object(self.__wrapper_folder_filename(filename), data)
+
+    def load_once(self, filename: str) -> bytes:
+        obj = self.client.get_object(self.__wrapper_folder_filename(filename))
+        data: bytes = obj.read()
+        return data
+
+    def load_stream(self, filename: str) -> Generator:
+        obj = self.client.get_object(self.__wrapper_folder_filename(filename))
+        while chunk := obj.read(4096):
+            yield chunk
+
+    def download(self, filename: str, target_filepath):
+        self.client.get_object_to_file(self.__wrapper_folder_filename(filename), target_filepath)
+
+    def exists(self, filename: str):
+        return self.client.object_exists(self.__wrapper_folder_filename(filename))
+
+    def delete(self, filename: str):
+        self.client.delete_object(self.__wrapper_folder_filename(filename))
+
+    def __wrapper_folder_filename(self, filename: str) -> str:
+        return posixpath.join(self.folder, filename) if self.folder else filename

api/extensions/storage/aws_s3_storage.py

@@ -0,0 +1,91 @@
+import logging
+from collections.abc import Generator
+
+import boto3  # type: ignore
+from botocore.client import Config  # type: ignore
+from botocore.exceptions import ClientError  # type: ignore
+
+from configs import dify_config
+from extensions.storage.base_storage import BaseStorage
+
+logger = logging.getLogger(__name__)
+
+
+class AwsS3Storage(BaseStorage):
+    """Implementation for Amazon Web Services S3 storage."""
+
+    def __init__(self):
+        super().__init__()
+        self.bucket_name = dify_config.S3_BUCKET_NAME
+        if dify_config.S3_USE_AWS_MANAGED_IAM:
+            logger.info("Using AWS managed IAM role for S3")
+
+            session = boto3.Session()
+            region_name = dify_config.S3_REGION
+            self.client = session.client(service_name="s3", region_name=region_name)
+        else:
+            logger.info("Using ak and sk for S3")
+
+            self.client = boto3.client(
+                "s3",
+                aws_secret_access_key=dify_config.S3_SECRET_KEY,
+                aws_access_key_id=dify_config.S3_ACCESS_KEY,
+                endpoint_url=dify_config.S3_ENDPOINT,
+                region_name=dify_config.S3_REGION,
+                config=Config(
+                    s3={"addressing_style": dify_config.S3_ADDRESS_STYLE},
+                    request_checksum_calculation="when_required",
+                    response_checksum_validation="when_required",
+                ),
+            )
+        # create bucket
+        try:
+            self.client.head_bucket(Bucket=self.bucket_name)
+        except ClientError as e:
+            # if bucket not exists, create it
+            if e.response["Error"]["Code"] == "404":
+                self.client.create_bucket(Bucket=self.bucket_name)
+            # if bucket is not accessible, pass, maybe the bucket is existing but not accessible
+            elif e.response["Error"]["Code"] == "403":
+                pass
+            else:
+                # other error, raise exception
+                raise
+
+    def save(self, filename, data):
+        self.client.put_object(Bucket=self.bucket_name, Key=filename, Body=data)
+
+    def load_once(self, filename: str) -> bytes:
+        try:
+            data: bytes = self.client.get_object(Bucket=self.bucket_name, Key=filename)["Body"].read()
+        except ClientError as ex:
+            if ex.response["Error"]["Code"] == "NoSuchKey":
+                raise FileNotFoundError("File not found")
+            else:
+                raise
+        return data
+
+    def load_stream(self, filename: str) -> Generator:
+        try:
+            response = self.client.get_object(Bucket=self.bucket_name, Key=filename)
+            yield from response["Body"].iter_chunks()
+        except ClientError as ex:
+            if ex.response["Error"]["Code"] == "NoSuchKey":
+                raise FileNotFoundError("file not found")
+            elif "reached max retries" in str(ex):
+                raise ValueError("please do not request the same file too frequently")
+            else:
+                raise
+
+    def download(self, filename, target_filepath):
+        self.client.download_file(self.bucket_name, filename, target_filepath)
+
+    def exists(self, filename):
+        try:
+            self.client.head_object(Bucket=self.bucket_name, Key=filename)
+            return True
+        except:
+            return False
+
+    def delete(self, filename):
+        self.client.delete_object(Bucket=self.bucket_name, Key=filename)

api/extensions/storage/azure_blob_storage.py

@@ -0,0 +1,84 @@
+from collections.abc import Generator
+from datetime import UTC, datetime, timedelta
+from typing import Optional
+
+from azure.identity import ChainedTokenCredential, DefaultAzureCredential
+from azure.storage.blob import AccountSasPermissions, BlobServiceClient, ResourceTypes, generate_account_sas
+
+from configs import dify_config
+from extensions.ext_redis import redis_client
+from extensions.storage.base_storage import BaseStorage
+
+
+class AzureBlobStorage(BaseStorage):
+    """Implementation for Azure Blob storage."""
+
+    def __init__(self):
+        super().__init__()
+        self.bucket_name = dify_config.AZURE_BLOB_CONTAINER_NAME
+        self.account_url = dify_config.AZURE_BLOB_ACCOUNT_URL
+        self.account_name = dify_config.AZURE_BLOB_ACCOUNT_NAME
+        self.account_key = dify_config.AZURE_BLOB_ACCOUNT_KEY
+
+        self.credential: Optional[ChainedTokenCredential] = None
+        if self.account_key == "managedidentity":
+            self.credential = DefaultAzureCredential()
+        else:
+            self.credential = None
+
+    def save(self, filename, data):
+        client = self._sync_client()
+        blob_container = client.get_container_client(container=self.bucket_name)
+        blob_container.upload_blob(filename, data)
+
+    def load_once(self, filename: str) -> bytes:
+        client = self._sync_client()
+        blob = client.get_container_client(container=self.bucket_name)
+        blob = blob.get_blob_client(blob=filename)
+        data: bytes = blob.download_blob().readall()
+        return data
+
+    def load_stream(self, filename: str) -> Generator:
+        client = self._sync_client()
+        blob = client.get_blob_client(container=self.bucket_name, blob=filename)
+        blob_data = blob.download_blob()
+        yield from blob_data.chunks()
+
+    def download(self, filename, target_filepath):
+        client = self._sync_client()
+
+        blob = client.get_blob_client(container=self.bucket_name, blob=filename)
+        with open(target_filepath, "wb") as my_blob:
+            blob_data = blob.download_blob()
+            blob_data.readinto(my_blob)
+
+    def exists(self, filename):
+        client = self._sync_client()
+
+        blob = client.get_blob_client(container=self.bucket_name, blob=filename)
+        return blob.exists()
+
+    def delete(self, filename):
+        client = self._sync_client()
+
+        blob_container = client.get_container_client(container=self.bucket_name)
+        blob_container.delete_blob(filename)
+
+    def _sync_client(self):
+        if self.account_key == "managedidentity":
+            return BlobServiceClient(account_url=self.account_url, credential=self.credential)  # type: ignore
+
+        cache_key = "azure_blob_sas_token_{}_{}".format(self.account_name, self.account_key)
+        cache_result = redis_client.get(cache_key)
+        if cache_result is not None:
+            sas_token = cache_result.decode("utf-8")
+        else:
+            sas_token = generate_account_sas(
+                account_name=self.account_name or "",
+                account_key=self.account_key or "",
+                resource_types=ResourceTypes(service=True, container=True, object=True),
+                permission=AccountSasPermissions(read=True, write=True, delete=True, list=True, add=True, create=True),
+                expiry=datetime.now(UTC).replace(tzinfo=None) + timedelta(hours=1),
+            )
+            redis_client.set(cache_key, sas_token, ex=3000)
+        return BlobServiceClient(account_url=self.account_url or "", credential=sas_token)

api/extensions/storage/baidu_obs_storage.py

@@ -0,0 +1,57 @@
+import base64
+import hashlib
+from collections.abc import Generator
+
+from baidubce.auth.bce_credentials import BceCredentials  # type: ignore
+from baidubce.bce_client_configuration import BceClientConfiguration  # type: ignore
+from baidubce.services.bos.bos_client import BosClient  # type: ignore
+
+from configs import dify_config
+from extensions.storage.base_storage import BaseStorage
+
+
+class BaiduObsStorage(BaseStorage):
+    """Implementation for Baidu OBS storage."""
+
+    def __init__(self):
+        super().__init__()
+        self.bucket_name = dify_config.BAIDU_OBS_BUCKET_NAME
+        client_config = BceClientConfiguration(
+            credentials=BceCredentials(
+                access_key_id=dify_config.BAIDU_OBS_ACCESS_KEY,
+                secret_access_key=dify_config.BAIDU_OBS_SECRET_KEY,
+            ),
+            endpoint=dify_config.BAIDU_OBS_ENDPOINT,
+        )
+
+        self.client = BosClient(config=client_config)
+
+    def save(self, filename, data):
+        md5 = hashlib.md5()
+        md5.update(data)
+        content_md5 = base64.standard_b64encode(md5.digest())
+        self.client.put_object(
+            bucket_name=self.bucket_name, key=filename, data=data, content_length=len(data), content_md5=content_md5
+        )
+
+    def load_once(self, filename: str) -> bytes:
+        response = self.client.get_object(bucket_name=self.bucket_name, key=filename)
+        data: bytes = response.data.read()
+        return data
+
+    def load_stream(self, filename: str) -> Generator:
+        response = self.client.get_object(bucket_name=self.bucket_name, key=filename).data
+        while chunk := response.read(4096):
+            yield chunk
+
+    def download(self, filename, target_filepath):
+        self.client.get_object_to_file(bucket_name=self.bucket_name, key=filename, file_name=target_filepath)
+
+    def exists(self, filename):
+        res = self.client.get_object_meta_data(bucket_name=self.bucket_name, key=filename)
+        if res is None:
+            return False
+        return True
+
+    def delete(self, filename):
+        self.client.delete_object(bucket_name=self.bucket_name, key=filename)

api/extensions/storage/base_storage.py

@@ -0,0 +1,32 @@
+"""Abstract interface for file storage implementations."""
+
+from abc import ABC, abstractmethod
+from collections.abc import Generator
+
+
+class BaseStorage(ABC):
+    """Interface for file storage."""
+
+    @abstractmethod
+    def save(self, filename, data):
+        raise NotImplementedError
+
+    @abstractmethod
+    def load_once(self, filename: str) -> bytes:
+        raise NotImplementedError
+
+    @abstractmethod
+    def load_stream(self, filename: str) -> Generator:
+        raise NotImplementedError
+
+    @abstractmethod
+    def download(self, filename, target_filepath):
+        raise NotImplementedError
+
+    @abstractmethod
+    def exists(self, filename):
+        raise NotImplementedError
+
+    @abstractmethod
+    def delete(self, filename):
+        raise NotImplementedError

api/extensions/storage/google_cloud_storage.py

@@ -0,0 +1,60 @@
+import base64
+import io
+import json
+from collections.abc import Generator
+
+from google.cloud import storage as google_cloud_storage  # type: ignore
+
+from configs import dify_config
+from extensions.storage.base_storage import BaseStorage
+
+
+class GoogleCloudStorage(BaseStorage):
+    """Implementation for Google Cloud storage."""
+
+    def __init__(self):
+        super().__init__()
+
+        self.bucket_name = dify_config.GOOGLE_STORAGE_BUCKET_NAME
+        service_account_json_str = dify_config.GOOGLE_STORAGE_SERVICE_ACCOUNT_JSON_BASE64
+        # if service_account_json_str is empty, use Application Default Credentials
+        if service_account_json_str:
+            service_account_json = base64.b64decode(service_account_json_str).decode("utf-8")
+            # convert str to object
+            service_account_obj = json.loads(service_account_json)
+            self.client = google_cloud_storage.Client.from_service_account_info(service_account_obj)
+        else:
+            self.client = google_cloud_storage.Client()
+
+    def save(self, filename, data):
+        bucket = self.client.get_bucket(self.bucket_name)
+        blob = bucket.blob(filename)
+        with io.BytesIO(data) as stream:
+            blob.upload_from_file(stream)
+
+    def load_once(self, filename: str) -> bytes:
+        bucket = self.client.get_bucket(self.bucket_name)
+        blob = bucket.get_blob(filename)
+        data: bytes = blob.download_as_bytes()
+        return data
+
+    def load_stream(self, filename: str) -> Generator:
+        bucket = self.client.get_bucket(self.bucket_name)
+        blob = bucket.get_blob(filename)
+        with blob.open(mode="rb") as blob_stream:
+            while chunk := blob_stream.read(4096):
+                yield chunk
+
+    def download(self, filename, target_filepath):
+        bucket = self.client.get_bucket(self.bucket_name)
+        blob = bucket.get_blob(filename)
+        blob.download_to_filename(target_filepath)
+
+    def exists(self, filename):
+        bucket = self.client.get_bucket(self.bucket_name)
+        blob = bucket.blob(filename)
+        return blob.exists()
+
+    def delete(self, filename):
+        bucket = self.client.get_bucket(self.bucket_name)
+        bucket.delete_blob(filename)

api/extensions/storage/huawei_obs_storage.py

@@ -0,0 +1,51 @@
+from collections.abc import Generator
+
+from obs import ObsClient  # type: ignore
+
+from configs import dify_config
+from extensions.storage.base_storage import BaseStorage
+
+
+class HuaweiObsStorage(BaseStorage):
+    """Implementation for Huawei OBS storage."""
+
+    def __init__(self):
+        super().__init__()
+
+        self.bucket_name = dify_config.HUAWEI_OBS_BUCKET_NAME
+        self.client = ObsClient(
+            access_key_id=dify_config.HUAWEI_OBS_ACCESS_KEY,
+            secret_access_key=dify_config.HUAWEI_OBS_SECRET_KEY,
+            server=dify_config.HUAWEI_OBS_SERVER,
+        )
+
+    def save(self, filename, data):
+        self.client.putObject(bucketName=self.bucket_name, objectKey=filename, content=data)
+
+    def load_once(self, filename: str) -> bytes:
+        data: bytes = self.client.getObject(bucketName=self.bucket_name, objectKey=filename)["body"].response.read()
+        return data
+
+    def load_stream(self, filename: str) -> Generator:
+        response = self.client.getObject(bucketName=self.bucket_name, objectKey=filename)["body"].response
+        while chunk := response.read(4096):
+            yield chunk
+
+    def download(self, filename, target_filepath):
+        self.client.getObject(bucketName=self.bucket_name, objectKey=filename, downloadPath=target_filepath)
+
+    def exists(self, filename):
+        res = self._get_meta(filename)
+        if res is None:
+            return False
+        return True
+
+    def delete(self, filename):
+        self.client.deleteObject(bucketName=self.bucket_name, objectKey=filename)
+
+    def _get_meta(self, filename):
+        res = self.client.getObjectMetadata(bucketName=self.bucket_name, objectKey=filename)
+        if res.status < 300:
+            return res
+        else:
+            return None

api/extensions/storage/opendal_storage.py

@@ -0,0 +1,89 @@
+import logging
+import os
+from collections.abc import Generator
+from pathlib import Path
+
+import opendal  # type: ignore[import]
+from dotenv import dotenv_values
+
+from extensions.storage.base_storage import BaseStorage
+
+logger = logging.getLogger(__name__)
+
+
+def _get_opendal_kwargs(*, scheme: str, env_file_path: str = ".env", prefix: str = "OPENDAL_"):
+    kwargs = {}
+    config_prefix = prefix + scheme.upper() + "_"
+    for key, value in os.environ.items():
+        if key.startswith(config_prefix):
+            kwargs[key[len(config_prefix) :].lower()] = value
+
+    file_env_vars: dict = dotenv_values(env_file_path) or {}
+    for key, value in file_env_vars.items():
+        if key.startswith(config_prefix) and key[len(config_prefix) :].lower() not in kwargs and value:
+            kwargs[key[len(config_prefix) :].lower()] = value
+
+    return kwargs
+
+
+class OpenDALStorage(BaseStorage):
+    def __init__(self, scheme: str, **kwargs):
+        kwargs = kwargs or _get_opendal_kwargs(scheme=scheme)
+
+        if scheme == "fs":
+            root = kwargs.get("root", "storage")
+            Path(root).mkdir(parents=True, exist_ok=True)
+
+        self.op = opendal.Operator(scheme=scheme, **kwargs)
+        logger.debug(f"opendal operator created with scheme {scheme}")
+        retry_layer = opendal.layers.RetryLayer(max_times=3, factor=2.0, jitter=True)
+        self.op = self.op.layer(retry_layer)
+        logger.debug("added retry layer to opendal operator")
+
+    def save(self, filename: str, data: bytes) -> None:
+        self.op.write(path=filename, bs=data)
+        logger.debug(f"file {filename} saved")
+
+    def load_once(self, filename: str) -> bytes:
+        if not self.exists(filename):
+            raise FileNotFoundError("File not found")
+
+        content: bytes = self.op.read(path=filename)
+        logger.debug(f"file {filename} loaded")
+        return content
+
+    def load_stream(self, filename: str) -> Generator:
+        if not self.exists(filename):
+            raise FileNotFoundError("File not found")
+
+        batch_size = 4096
+        file = self.op.open(path=filename, mode="rb")
+        while chunk := file.read(batch_size):
+            yield chunk
+        logger.debug(f"file {filename} loaded as stream")
+
+    def download(self, filename: str, target_filepath: str):
+        if not self.exists(filename):
+            raise FileNotFoundError("File not found")
+
+        with Path(target_filepath).open("wb") as f:
+            f.write(self.op.read(path=filename))
+        logger.debug(f"file {filename} downloaded to {target_filepath}")
+
+    def exists(self, filename: str) -> bool:
+        # FIXME this is a workaround for opendal python-binding do not have a exists method and no better
+        # error handler here when opendal python-binding has a exists method, we should use it
+        # more https://github.com/apache/opendal/blob/main/bindings/python/src/operator.rs
+        try:
+            res: bool = self.op.stat(path=filename).mode.is_file()
+            logger.debug(f"file {filename} checked")
+            return res
+        except Exception:
+            return False
+
+    def delete(self, filename: str):
+        if self.exists(filename):
+            self.op.delete(path=filename)
+            logger.debug(f"file {filename} deleted")
+            return
+        logger.debug(f"file {filename} not found, skip delete")

api/extensions/storage/oracle_oci_storage.py

@@ -0,0 +1,59 @@
+from collections.abc import Generator
+
+import boto3  # type: ignore
+from botocore.exceptions import ClientError  # type: ignore
+
+from configs import dify_config
+from extensions.storage.base_storage import BaseStorage
+
+
+class OracleOCIStorage(BaseStorage):
+    """Implementation for Oracle OCI storage."""
+
+    def __init__(self):
+        super().__init__()
+
+        self.bucket_name = dify_config.OCI_BUCKET_NAME
+        self.client = boto3.client(
+            "s3",
+            aws_secret_access_key=dify_config.OCI_SECRET_KEY,
+            aws_access_key_id=dify_config.OCI_ACCESS_KEY,
+            endpoint_url=dify_config.OCI_ENDPOINT,
+            region_name=dify_config.OCI_REGION,
+        )
+
+    def save(self, filename, data):
+        self.client.put_object(Bucket=self.bucket_name, Key=filename, Body=data)
+
+    def load_once(self, filename: str) -> bytes:
+        try:
+            data: bytes = self.client.get_object(Bucket=self.bucket_name, Key=filename)["Body"].read()
+        except ClientError as ex:
+            if ex.response["Error"]["Code"] == "NoSuchKey":
+                raise FileNotFoundError("File not found")
+            else:
+                raise
+        return data
+
+    def load_stream(self, filename: str) -> Generator:
+        try:
+            response = self.client.get_object(Bucket=self.bucket_name, Key=filename)
+            yield from response["Body"].iter_chunks()
+        except ClientError as ex:
+            if ex.response["Error"]["Code"] == "NoSuchKey":
+                raise FileNotFoundError("File not found")
+            else:
+                raise
+
+    def download(self, filename, target_filepath):
+        self.client.download_file(self.bucket_name, filename, target_filepath)
+
+    def exists(self, filename):
+        try:
+            self.client.head_object(Bucket=self.bucket_name, Key=filename)
+            return True
+        except:
+            return False
+
+    def delete(self, filename):
+        self.client.delete_object(Bucket=self.bucket_name, Key=filename)