dify/api/services/workflow_runtime/pipeline_runner.py

import logging
import time
from typing import Protocol, cast, override

from sqlalchemy import select
from sqlalchemy.orm import Session

from core.app.apps.base_app_queue_manager import AppQueueManager
from core.app.apps.pipeline.pipeline_config_manager import PipelineConfig
from core.app.entities.app_invoke_entities import (
    InvokeFrom,
    RagPipelineGenerateEntity,
    UserFrom,
    build_dify_run_context,
)
from core.credit_usage import CreditUsageAppType
from core.db.session_factory import create_session
from core.workflow.system_variables import build_bootstrap_variables, build_system_variables
from core.workflow.variable_pool_initializer import add_node_inputs_to_pool, add_variables_to_pool
from graphon.engine_events import EngineEvent, GraphRunFailedEvent
from graphon.enums import WorkflowType
from graphon.graph import Graph
from graphon.runtime import RuntimeState, VariablePool
from graphon.variable_loader import VariableLoader
from graphon.variables.variables import RAGPipelineVariable, RAGPipelineVariableInput
from models.dataset import Dataset, Pipeline
from models.model import EndUser
from models.workflow import Workflow
from services.workflow_runtime.base_runner import PreparedWorkflowRun, WorkflowBasedAppRunner, WorkflowRunDriver
from services.workflow_runtime.entry import WorkflowEntry
from services.workflow_runtime.node_factory import DifyGraphInitContext, DifyNodeFactory, get_default_root_node_id

logger = logging.getLogger(__name__)


class PipelineDocumentStore(Protocol):
    """Knowledge state required by a pipeline run, with explicit session ownership."""

    def get_pipeline_dataset(self, pipeline: Pipeline, *, session: Session) -> Dataset | None: ...

    def exists(self, *, workspace_id: str, dataset_id: str, document_id: str) -> bool: ...

    def mark_failed(self, *, workspace_id: str, dataset_id: str, document_id: str, error: str) -> None: ...


class PipelineRunner(WorkflowBasedAppRunner):
    """
    Pipeline Application Runner
    """

    def __init__(
        self,
        application_generate_entity: RagPipelineGenerateEntity,
        queue_manager: AppQueueManager,
        variable_loader: VariableLoader,
        workflow: Workflow,
        system_user_id: str,
        documents: PipelineDocumentStore,
        workflow_thread_pool_id: str | None = None,
        execution_driver: WorkflowRunDriver | None = None,
    ) -> None:
        """
        :param application_generate_entity: application generate entity
        :param queue_manager: application queue manager
        :param workflow_thread_pool_id: workflow thread pool id
        """
        super().__init__(
            queue_manager=queue_manager,
            execution_driver=execution_driver,
            variable_loader=variable_loader,
            app_id=application_generate_entity.app_config.app_id,
        )
        self.application_generate_entity = application_generate_entity
        self.workflow_thread_pool_id = workflow_thread_pool_id
        self._workflow = workflow
        self._sys_user_id = system_user_id
        self._documents = documents

    def _get_app_id(self) -> str:
        return self.application_generate_entity.app_config.app_id

    @override
    def prepare(self) -> PreparedWorkflowRun:
        """Prepare graph execution and passive application observers."""
        app_config = self.application_generate_entity.app_config
        app_config = cast(PipelineConfig, app_config)
        invoke_from = self.application_generate_entity.invoke_from

        if self.application_generate_entity.single_iteration_run or self.application_generate_entity.single_loop_run:
            invoke_from = InvokeFrom.DEBUGGER

        user_from = self._resolve_user_from(invoke_from)

        user_id = None
        with create_session() as session:
            if invoke_from in {InvokeFrom.WEB_APP, InvokeFrom.SERVICE_API}:
                end_user = session.get(EndUser, self.application_generate_entity.user_id)
                if end_user:
                    user_id = end_user.session_id
            else:
                user_id = self.application_generate_entity.user_id

            pipeline = session.get(Pipeline, app_config.app_id)
            if not pipeline or pipeline.tenant_id != app_config.tenant_id:
                raise ValueError("Pipeline not found")

            dataset = self._documents.get_pipeline_dataset(pipeline, session=session)
            if (
                not dataset
                or dataset.tenant_id != pipeline.tenant_id
                or dataset.id != self.application_generate_entity.dataset_id
            ):
                raise ValueError("Pipeline dataset not found")

            document_id = self.application_generate_entity.document_id
            original_document_id = self.application_generate_entity.original_document_id
            dataset_workspace_id = dataset.tenant_id
            dataset_id = dataset.id
            if document_id and not self._documents.exists(
                workspace_id=dataset_workspace_id,
                dataset_id=dataset_id,
                document_id=document_id,
            ):
                raise ValueError("Pipeline document not found")
            if original_document_id and original_document_id != document_id:
                if not self._documents.exists(
                    workspace_id=dataset_workspace_id,
                    dataset_id=dataset_id,
                    document_id=original_document_id,
                ):
                    raise ValueError("Pipeline original document not found")

            workflow = self.get_workflow(session=session, pipeline=pipeline, workflow_id=app_config.workflow_id)
            if not workflow:
                raise ValueError("Workflow not initialized")

            session.expunge(pipeline)
            session.expunge(workflow)

        # if only single iteration run is requested
        if self.application_generate_entity.single_iteration_run or self.application_generate_entity.single_loop_run:
            # Handle single iteration or single loop run
            graph, variable_pool, graph_runtime_state = self._prepare_single_node_execution(
                workflow=workflow,
                single_iteration_run=self.application_generate_entity.single_iteration_run,
                single_loop_run=self.application_generate_entity.single_loop_run,
                user_id=self.application_generate_entity.user_id,
            )
        else:
            inputs = self.application_generate_entity.inputs
            files = self.application_generate_entity.files

            # Create a variable pool.
            system_inputs = build_system_variables(
                files=files,
                user_id=user_id,
                app_id=app_config.app_id,
                workflow_id=app_config.workflow_id,
                workflow_execution_id=self.application_generate_entity.workflow_execution_id,
                document_id=self.application_generate_entity.document_id,
                original_document_id=self.application_generate_entity.original_document_id,
                batch=self.application_generate_entity.batch,
                dataset_id=self.application_generate_entity.dataset_id,
                datasource_type=self.application_generate_entity.datasource_type,
                datasource_info=self.application_generate_entity.datasource_info,
                invoke_from=invoke_from.value,
            )

            rag_pipeline_variables = []
            if workflow.rag_pipeline_variables:
                for v in workflow.rag_pipeline_variables:
                    rag_pipeline_variable = RAGPipelineVariable.model_validate(v)
                    if (
                        rag_pipeline_variable.belong_to_node_id
                        in (self.application_generate_entity.start_node_id, "shared")
                    ) and rag_pipeline_variable.variable in inputs:
                        rag_pipeline_variables.append(
                            RAGPipelineVariableInput(
                                variable=rag_pipeline_variable,
                                value=inputs[rag_pipeline_variable.variable],
                            )
                        )

            variable_pool = VariablePool()
            add_variables_to_pool(
                variable_pool,
                build_bootstrap_variables(
                    system_variables=system_inputs,
                    environment_variables=workflow.environment_variables,
                    rag_pipeline_variables=rag_pipeline_variables,
                ),
            )
            root_node_id = self.application_generate_entity.start_node_id or get_default_root_node_id(
                workflow.graph_dict
            )
            add_node_inputs_to_pool(variable_pool, node_id=root_node_id, inputs=inputs)
            graph_runtime_state = RuntimeState(
                workflow_id=workflow.id, variable_pool=variable_pool, start_at=time.perf_counter()
            )

            # init graph
            graph = self._init_rag_pipeline_graph(
                graph_runtime_state=graph_runtime_state,
                start_node_id=root_node_id,
                workflow=workflow,
                user_from=user_from,
                invoke_from=invoke_from,
            )

        # RUN WORKFLOW
        workflow_entry = WorkflowEntry(
            graph=graph,
            graph_config=workflow.graph_dict,
            call_depth=self.application_generate_entity.call_depth,
            graph_runtime_state=graph_runtime_state,
        )

        self._queue_manager.graph_runtime_state = graph_runtime_state

        self._dataset_workspace_id = dataset_workspace_id
        self._dataset_id = dataset_id
        self._document_id = document_id
        return PreparedWorkflowRun(
            entry=workflow_entry,
            application_layers=self._graph_engine_layers,
            generate_entity=self.application_generate_entity,
            workflow_type=WorkflowType(workflow.type),
            workflow_version=workflow.version,
        )

    @override
    def handle_event(self, workflow_entry: WorkflowEntry, event: EngineEvent, **kwargs) -> None:
        self._update_document_status(
            event,
            workspace_id=self._dataset_workspace_id,
            dataset_id=self._dataset_id,
            document_id=self._document_id,
        )
        super().handle_event(workflow_entry, event, **kwargs)

    def get_workflow(self, session: Session, pipeline: Pipeline, workflow_id: str) -> Workflow | None:
        """
        Get workflow
        """
        # fetch workflow by workflow_id
        workflow = session.scalar(
            select(Workflow)
            .where(Workflow.tenant_id == pipeline.tenant_id, Workflow.app_id == pipeline.id, Workflow.id == workflow_id)
            .limit(1)
        )

        # return workflow
        return workflow

    def _init_rag_pipeline_graph(
        self,
        workflow: Workflow,
        graph_runtime_state: RuntimeState,
        start_node_id: str | None = None,
        user_from: UserFrom = UserFrom.ACCOUNT,
        invoke_from: InvokeFrom = InvokeFrom.SERVICE_API,
    ) -> Graph:
        """
        Init pipeline graph
        """
        graph_config = workflow.graph_dict
        if "nodes" not in graph_config or "edges" not in graph_config:
            raise ValueError("nodes or edges not found in workflow graph")

        if not isinstance(graph_config.get("nodes"), list):
            raise ValueError("nodes in workflow graph must be a list")

        if not isinstance(graph_config.get("edges"), list):
            raise ValueError("edges in workflow graph must be a list")
        # nodes = graph_config.get("nodes", [])
        # edges = graph_config.get("edges", [])
        # real_run_nodes = []
        # real_edges = []
        # exclude_node_ids = []
        # for node in nodes:
        #     node_id = node.get("id")
        #     node_type = node.get("data", {}).get("type", "")
        #     if node_type == "datasource":
        #         if start_node_id != node_id:
        #             exclude_node_ids.append(node_id)
        #             continue
        #     real_run_nodes.append(node)

        # for edge in edges:
        #     if edge.get("source") in exclude_node_ids:
        #         continue
        #     real_edges.append(edge)
        # graph_config = dict(graph_config)
        # graph_config["nodes"] = real_run_nodes
        # graph_config["edges"] = real_edges
        # init graph
        # Create explicit graph init context for Graph.init.
        run_context = build_dify_run_context(
            tenant_id=workflow.tenant_id,
            app_id=self._app_id,
            user_id=self.application_generate_entity.user_id,
            user_from=user_from,
            invoke_from=invoke_from,
            app_type=CreditUsageAppType.RAG_PIPELINE,
        )
        graph_init_context = DifyGraphInitContext(
            execution_driver=self._execution_driver,
            workflow_id=workflow.id,
            graph_config=graph_config,
            run_context=run_context,
            call_depth=0,
        )

        node_factory = DifyNodeFactory.from_graph_init_context(
            graph_init_context=graph_init_context,
            runtime_state=graph_runtime_state,
        )
        if start_node_id is None:
            start_node_id = get_default_root_node_id(graph_config)
        graph = Graph.init(graph_config=graph_config, node_factory=node_factory, root_node_id=start_node_id)

        if not graph:
            raise ValueError("graph not found in workflow")

        return graph

    def _update_document_status(
        self,
        event: EngineEvent,
        *,
        workspace_id: str,
        dataset_id: str,
        document_id: str | None,
    ) -> None:
        """Set an owner-bound document to error after a failed graph run, if it exists."""
        if not isinstance(event, GraphRunFailedEvent) or document_id is None:
            return

        self._documents.mark_failed(
            workspace_id=workspace_id,
            dataset_id=dataset_id,
            document_id=document_id,
            error=event.error or "Unknown error",
        )