dify/api/services/app_generate_service.py

from __future__ import annotations

import logging
import threading
import uuid
from collections.abc import Callable, Generator, Mapping
from typing import TYPE_CHECKING, Any, Literal

from sqlalchemy.orm import Session

from configs import dify_config
from core.app.apps.advanced_chat.app_generator import AdvancedChatAppGenerator
from core.app.apps.agent_app.app_generator import AgentAppGenerator
from core.app.apps.agent_chat.app_generator import AgentChatAppGenerator
from core.app.apps.chat.app_generator import ChatAppGenerator
from core.app.apps.completion.app_generator import CompletionAppGenerator
from core.app.apps.message_based_app_generator import MessageBasedAppGenerator
from core.app.apps.workflow.app_generator import WorkflowAppGenerator
from core.app.entities.app_invoke_entities import InvokeFrom
from core.app.features.rate_limiting import RateLimit
from core.app.features.rate_limiting.rate_limit import rate_limit_context
from core.app.layers.pause_state_persist_layer import PauseStateLayerConfig
from core.db import session_factory
from core.trigger.constants import is_trigger_node_type
from enums import DeploymentEdition, QuotaType
from extensions.otel import AppGenerateHandler, trace_span
from models.model import Account, App, AppMode, EndUser
from models.workflow import Workflow, WorkflowRun
from services.errors.app import (
    QuotaExceededError,
    TriggerWorkflowServiceModeUnavailableError,
    WorkflowIdFormatError,
    WorkflowNotFoundError,
)
from services.errors.llm import InvokeRateLimitError
from services.quota_service import QuotaService, unlimited
from services.workflow_service import WorkflowService
from tasks.app_generate.workflow_execute_task import AppExecutionParams, workflow_based_app_execution_task

logger = logging.getLogger(__name__)

SSE_TASK_START_FALLBACK_MS = 200
_MANUAL_WORKFLOW_INVOKE_SOURCES = frozenset(
    {
        InvokeFrom.OPENAPI,
        InvokeFrom.SERVICE_API,
        InvokeFrom.WEB_APP,
    }
)

if TYPE_CHECKING:
    from controllers.console.app.workflow import LoopNodeRunPayload


class AppGenerateService:
    @staticmethod
    def _build_streaming_task_on_subscribe(
        start_task: Callable[[], None],
    ) -> Callable[[], None]:
        """
        Build a subscription callback that starts the background task on first subscribe.

        Pub/Sub transports also use a short fallback timer so the task still runs if the
        client never connects. Streams rely on their prepared delivery boundary instead.
        """
        started = False
        lock = threading.Lock()

        def _try_start() -> bool:
            nonlocal started
            with lock:
                if started:
                    return True
                try:
                    start_task()
                except Exception:
                    logger.exception("Failed to enqueue streaming task")
                    return False
                started = True
                return True

        if dify_config.PUBSUB_REDIS_CHANNEL_TYPE == "streams":

            def _on_subscribe_streams() -> None:
                _try_start()

            return _on_subscribe_streams

        timer = threading.Timer(SSE_TASK_START_FALLBACK_MS / 1000.0, _try_start)
        timer.daemon = True
        timer.start()

        def _on_subscribe() -> None:
            if _try_start():
                timer.cancel()

        return _on_subscribe

    @classmethod
    @trace_span(AppGenerateHandler)
    def generate(
        cls,
        app_model: App,
        user: Account | EndUser,
        args: Mapping[str, Any],
        invoke_from: InvokeFrom,
        *,
        session: Session,
        streaming: bool = True,
        root_node_id: str | None = None,
        workflow_execution_mode: Literal["celery", "in_process"] = "celery",
    ):
        """
        App Content Generate
        :param app_model: app model
        :param user: user
        :param args: args
        :param invoke_from: invoke from
        :param streaming: streaming
        :param workflow_execution_mode: Execute workflow streams locally when already
            inside a worker task, avoiding a synchronous wait for another Celery slot.
        :return:
        """
        return cls._run_with_guardrails(
            app_model=app_model,
            streaming=streaming,
            action=lambda rate_limit, request_id: cls._dispatch_generate(
                app_model=app_model,
                user=user,
                args=args,
                invoke_from=invoke_from,
                streaming=streaming,
                root_node_id=root_node_id,
                session=session,
                rate_limit=rate_limit,
                request_id=request_id,
                workflow_execution_mode=workflow_execution_mode,
            ),
        )

    @classmethod
    def _run_with_guardrails(
        cls,
        *,
        app_model: App,
        streaming: bool,
        action: Callable[[RateLimit, str], Any],
    ):
        quota_charge = unlimited()
        if dify_config.DEPLOYMENT_EDITION == DeploymentEdition.CLOUD:
            try:
                quota_charge = QuotaService.reserve(QuotaType.WORKFLOW, app_model.tenant_id)
            except QuotaExceededError:
                raise InvokeRateLimitError(f"Workflow execution quota limit reached for tenant {app_model.tenant_id}")

        # app level rate limiter
        max_active_request = cls._get_max_active_requests(app_model)
        rate_limit = RateLimit(app_model.id, max_active_request)
        request_id = RateLimit.gen_request_key()
        try:
            request_id = rate_limit.enter(request_id)
            quota_charge.commit()
            return action(rate_limit, request_id)
        except Exception:
            quota_charge.refund()
            if streaming:
                rate_limit.exit(request_id)
            raise
        finally:
            if not streaming:
                rate_limit.exit(request_id)

    @classmethod
    def _dispatch_generate(
        cls,
        *,
        app_model: App,
        user: Account | EndUser,
        args: Mapping[str, Any],
        invoke_from: InvokeFrom,
        streaming: bool,
        root_node_id: str | None,
        session: Session,
        rate_limit: RateLimit,
        request_id: str,
        workflow_execution_mode: Literal["celery", "in_process"] = "celery",
    ):
        effective_mode = (
            AppMode.AGENT_CHAT
            if app_model.is_agent_with_session(session=session) and app_model.mode != AppMode.AGENT_CHAT
            else app_model.mode
        )
        match effective_mode:
            case AppMode.COMPLETION:
                return rate_limit.generate(
                    CompletionAppGenerator.convert_to_event_stream(
                        CompletionAppGenerator().generate(
                            session=session,
                            app_model=app_model,
                            user=user,
                            args=args,
                            invoke_from=invoke_from,
                            streaming=streaming,
                        ),
                    ),
                    request_id=request_id,
                )
            case AppMode.AGENT_CHAT:
                return rate_limit.generate(
                    AgentChatAppGenerator.convert_to_event_stream(
                        AgentChatAppGenerator().generate(
                            session=session,
                            app_model=app_model,
                            user=user,
                            args=args,
                            invoke_from=invoke_from,
                            streaming=streaming,
                        ),
                    ),
                    request_id,
                )
            case AppMode.AGENT:
                return rate_limit.generate(
                    AgentAppGenerator.convert_to_event_stream(
                        AgentAppGenerator().generate(
                            app_model=app_model,
                            user=user,
                            args=args,
                            invoke_from=invoke_from,
                            session=session,
                            streaming=streaming,
                        ),
                    ),
                    request_id,
                )
            case AppMode.CHAT:
                return rate_limit.generate(
                    ChatAppGenerator.convert_to_event_stream(
                        ChatAppGenerator().generate(
                            session=session,
                            app_model=app_model,
                            user=user,
                            args=args,
                            invoke_from=invoke_from,
                            streaming=streaming,
                        ),
                    ),
                    request_id=request_id,
                )
            case AppMode.ADVANCED_CHAT:
                workflow_id = args.get("workflow_id")
                workflow = cls._get_workflow(app_model, invoke_from, workflow_id, session=session)

                if streaming and workflow_execution_mode == "celery":
                    # Streaming mode: subscribe to SSE and enqueue the execution on first subscriber
                    with rate_limit_context(rate_limit, request_id):
                        payload = AppExecutionParams.new(
                            app_model=app_model,
                            workflow=workflow,
                            user=user,
                            args=args,
                            invoke_from=invoke_from,
                            streaming=True,
                            call_depth=0,
                            workflow_run_id=str(uuid.uuid4()),
                        )
                        payload_json = payload.model_dump_json()

                    def on_subscribe():
                        workflow_based_app_execution_task.delay(payload_json)

                    on_subscribe = cls._build_streaming_task_on_subscribe(on_subscribe)
                    generator = AdvancedChatAppGenerator()
                    return rate_limit.generate(
                        generator.convert_to_event_stream(
                            generator.retrieve_events(
                                AppMode.ADVANCED_CHAT,
                                payload.workflow_run_id,
                                on_subscribe=on_subscribe,
                            ),
                        ),
                        request_id=request_id,
                    )

                # Reuse the native generator for blocking requests and streams
                # owned by an existing worker task. Both retain pause persistence.
                pause_config = PauseStateLayerConfig(
                    session_factory=session_factory.get_session_maker(),
                    state_owner_user_id=workflow.created_by,
                )
                advanced_generator = AdvancedChatAppGenerator()
                return rate_limit.generate(
                    advanced_generator.convert_to_event_stream(
                        advanced_generator.generate(
                            app_model=app_model,
                            workflow=workflow,
                            user=user,
                            args=args,
                            invoke_from=invoke_from,
                            workflow_run_id=str(uuid.uuid4()),
                            streaming=streaming,
                            pause_state_config=pause_config,
                            session=session,
                        )
                    ),
                    request_id=request_id,
                )
            case AppMode.WORKFLOW:
                workflow_id = args.get("workflow_id")
                workflow = cls._get_workflow(app_model, invoke_from, workflow_id, session=session)
                cls._ensure_workflow_service_mode_available(workflow=workflow, invoke_from=invoke_from)
                if streaming and workflow_execution_mode == "celery":
                    with rate_limit_context(rate_limit, request_id):
                        payload = AppExecutionParams.new(
                            app_model=app_model,
                            workflow=workflow,
                            user=user,
                            args=args,
                            invoke_from=invoke_from,
                            streaming=True,
                            call_depth=0,
                            root_node_id=root_node_id,
                            workflow_run_id=str(uuid.uuid4()),
                        )
                        payload_json = payload.model_dump_json()

                    def on_subscribe():
                        workflow_based_app_execution_task.delay(payload_json)

                    on_subscribe = cls._build_streaming_task_on_subscribe(on_subscribe)
                    return rate_limit.generate(
                        WorkflowAppGenerator.convert_to_event_stream(
                            MessageBasedAppGenerator.retrieve_events(
                                AppMode.WORKFLOW,
                                payload.workflow_run_id,
                                on_subscribe=on_subscribe,
                            ),
                        ),
                        request_id,
                    )

                pause_config = PauseStateLayerConfig(
                    session_factory=session_factory.get_session_maker(),
                    state_owner_user_id=workflow.created_by,
                )
                return rate_limit.generate(
                    WorkflowAppGenerator.convert_to_event_stream(
                        WorkflowAppGenerator().generate(
                            app_model=app_model,
                            workflow=workflow,
                            user=user,
                            args=args,
                            invoke_from=invoke_from,
                            streaming=streaming,
                            root_node_id=root_node_id,
                            call_depth=0,
                            pause_state_config=pause_config,
                        ),
                    ),
                    request_id,
                )
            case _:
                raise ValueError(f"Invalid app mode {app_model.mode}")

    @staticmethod
    def _ensure_workflow_service_mode_available(*, workflow: Workflow, invoke_from: InvokeFrom) -> None:
        if invoke_from not in _MANUAL_WORKFLOW_INVOKE_SOURCES:
            return

        for _, node_data in workflow.walk_nodes():
            node_type = node_data.get("type")
            if isinstance(node_type, str) and is_trigger_node_type(node_type):
                raise TriggerWorkflowServiceModeUnavailableError()

    @staticmethod
    def _get_max_active_requests(app: App) -> int:
        """
        Get the maximum number of active requests allowed for an app.

        Returns the smaller value between app's custom limit and global config limit.
        A value of 0 means infinite (no limit).

        Args:
            app: The App model instance

        Returns:
            The maximum number of active requests allowed
        """
        app_limit = app.max_active_requests or dify_config.APP_DEFAULT_ACTIVE_REQUESTS
        config_limit = dify_config.APP_MAX_ACTIVE_REQUESTS

        # Filter out infinite (0) values and return the minimum, or 0 if both are infinite
        limits = [limit for limit in [app_limit, config_limit] if limit > 0]
        return min(limits) if limits else 0

    @classmethod
    def generate_single_iteration(
        cls,
        app_model: App,
        user: Account,
        node_id: str,
        args: Any,
        *,
        session: Session,
        streaming: bool = True,
    ):
        match app_model.mode:
            case AppMode.COMPLETION | AppMode.CHAT | AppMode.AGENT_CHAT:
                raise ValueError(f"Invalid app mode {app_model.mode}")
            case AppMode.ADVANCED_CHAT:
                workflow = cls._get_workflow(app_model, InvokeFrom.DEBUGGER, session=session)
                return AdvancedChatAppGenerator.convert_to_event_stream(
                    AdvancedChatAppGenerator().single_iteration_generate(
                        app_model=app_model,
                        workflow=workflow,
                        node_id=node_id,
                        user=user,
                        args=args,
                        streaming=streaming,
                        session=session,
                    )
                )
            case AppMode.WORKFLOW:
                workflow = cls._get_workflow(app_model, InvokeFrom.DEBUGGER, session=session)
                return AdvancedChatAppGenerator.convert_to_event_stream(
                    WorkflowAppGenerator().single_iteration_generate(
                        app_model=app_model,
                        workflow=workflow,
                        node_id=node_id,
                        user=user,
                        args=args,
                        streaming=streaming,
                        session=session,
                    )
                )
            case AppMode.CHANNEL | AppMode.RAG_PIPELINE:
                raise ValueError(f"Invalid app mode {app_model.mode}")
            case _:
                raise ValueError(f"Invalid app mode {app_model.mode}")

    @classmethod
    def generate_single_loop(
        cls,
        app_model: App,
        user: Account,
        node_id: str,
        args: LoopNodeRunPayload,
        *,
        session: Session,
        streaming: bool = True,
    ):
        match app_model.mode:
            case AppMode.COMPLETION | AppMode.CHAT | AppMode.AGENT_CHAT:
                raise ValueError(f"Invalid app mode {app_model.mode}")
            case AppMode.ADVANCED_CHAT:
                workflow = cls._get_workflow(app_model, InvokeFrom.DEBUGGER, session=session)
                return AdvancedChatAppGenerator.convert_to_event_stream(
                    AdvancedChatAppGenerator().single_loop_generate(
                        app_model=app_model,
                        workflow=workflow,
                        node_id=node_id,
                        user=user,
                        args=args,
                        streaming=streaming,
                        session=session,
                    )
                )
            case AppMode.WORKFLOW:
                workflow = cls._get_workflow(app_model, InvokeFrom.DEBUGGER, session=session)
                return AdvancedChatAppGenerator.convert_to_event_stream(
                    WorkflowAppGenerator().single_loop_generate(
                        app_model=app_model,
                        workflow=workflow,
                        node_id=node_id,
                        user=user,
                        args=args,
                        streaming=streaming,
                        session=session,
                    )
                )
            case AppMode.CHANNEL | AppMode.RAG_PIPELINE:
                raise ValueError(f"Invalid app mode {app_model.mode}")
            case _:
                raise ValueError(f"Invalid app mode {app_model.mode}")

    @classmethod
    def generate_more_like_this(
        cls,
        app_model: App,
        user: Account | EndUser,
        message_id: str,
        invoke_from: InvokeFrom,
        *,
        session: Session,
        streaming: bool = True,
    ) -> Mapping | Generator:
        """
        Generate more like this
        :param app_model: app model
        :param user: user
        :param message_id: message id
        :param invoke_from: invoke from
        :param streaming: streaming
        :return:
        """
        return CompletionAppGenerator().generate_more_like_this(
            session=session,
            app_model=app_model,
            message_id=message_id,
            user=user,
            invoke_from=invoke_from,
            stream=streaming,
        )

    @classmethod
    def _get_workflow(
        cls,
        app_model: App,
        invoke_from: InvokeFrom,
        workflow_id: str | None = None,
        *,
        session: Session,
    ) -> Workflow:
        """
        Get workflow
        :param app_model: app model
        :param invoke_from: invoke from
        :param workflow_id: optional workflow id to specify a specific version
        :return:
        """
        workflow_service = WorkflowService()

        # If workflow_id is specified, get the specific workflow version
        if workflow_id:
            try:
                _ = uuid.UUID(workflow_id)
            except ValueError:
                raise WorkflowIdFormatError(f"Invalid workflow_id format: '{workflow_id}'. ")
            workflow = workflow_service.get_published_workflow_by_id(
                app_model=app_model, workflow_id=workflow_id, session=session
            )
            if not workflow:
                raise WorkflowNotFoundError(f"Workflow not found with id: {workflow_id}")
            return workflow

        if invoke_from == InvokeFrom.DEBUGGER:
            # fetch draft workflow by app_model
            workflow = workflow_service.get_draft_workflow(app_model=app_model, session=session)

            if not workflow:
                raise ValueError("Workflow not initialized")
        else:
            # fetch published workflow by app_model
            workflow = workflow_service.get_published_workflow(app_model=app_model, session=session)

            if not workflow:
                raise ValueError("Workflow not published")

        return workflow

    @classmethod
    def get_response_generator(
        cls,
        app_model: App,
        workflow_run: WorkflowRun,
    ):
        if workflow_run.status.is_ended():
            # TODO(QuantumGhost): handled the ended scenario.
            pass

        generator = AdvancedChatAppGenerator()

        return generator.convert_to_event_stream(
            generator.retrieve_events(AppMode(app_model.mode), workflow_run.id),
        )