import logging import time from typing import cast from sqlalchemy import select from core.app.apps.base_app_queue_manager import AppQueueManager from core.app.apps.pipeline.pipeline_config_manager import PipelineConfig from core.app.apps.workflow_app_runner import WorkflowBasedAppRunner from core.app.entities.app_invoke_entities import ( InvokeFrom, RagPipelineGenerateEntity, UserFrom, build_dify_run_context, ) from core.app.workflow.layers.persistence import PersistenceWorkflowInfo, WorkflowPersistenceLayer from core.repositories.factory import WorkflowExecutionRepository, WorkflowNodeExecutionRepository from core.workflow.node_factory import DifyGraphInitContext, DifyNodeFactory, get_default_root_node_id 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 core.workflow.workflow_entry import WorkflowEntry from extensions.ext_database import db from graphon.enums import WorkflowType from graphon.graph import Graph from graphon.graph_events import GraphEngineEvent, GraphRunFailedEvent from graphon.runtime import GraphRuntimeState, VariablePool from graphon.variable_loader import VariableLoader from graphon.variables.variables import RAGPipelineVariable, RAGPipelineVariableInput from models.dataset import Document, Pipeline from models.model import EndUser from models.workflow import Workflow logger = logging.getLogger(__name__) 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, workflow_execution_repository: WorkflowExecutionRepository, workflow_node_execution_repository: WorkflowNodeExecutionRepository, workflow_thread_pool_id: str | 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, 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._workflow_execution_repository = workflow_execution_repository self._workflow_node_execution_repository = workflow_node_execution_repository def _get_app_id(self) -> str: return self.application_generate_entity.app_config.app_id def run(self) -> None: """ Run application """ 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 if invoke_from in {InvokeFrom.WEB_APP, InvokeFrom.SERVICE_API}: end_user = db.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 = db.session.get(Pipeline, app_config.app_id) if not pipeline: raise ValueError("Pipeline not found") workflow = self.get_workflow(pipeline=pipeline, workflow_id=app_config.workflow_id) if not workflow: raise ValueError("Workflow not initialized") db.session.close() # 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 = GraphRuntimeState(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( tenant_id=workflow.tenant_id, app_id=workflow.app_id, workflow_id=workflow.id, graph=graph, graph_config=workflow.graph_dict, user_id=self.application_generate_entity.user_id, user_from=user_from, invoke_from=invoke_from, call_depth=self.application_generate_entity.call_depth, graph_runtime_state=graph_runtime_state, variable_pool=variable_pool, ) self._queue_manager.graph_runtime_state = graph_runtime_state persistence_layer = WorkflowPersistenceLayer( application_generate_entity=self.application_generate_entity, workflow_info=PersistenceWorkflowInfo( workflow_id=workflow.id, workflow_type=WorkflowType(workflow.type), version=workflow.version, graph_data=workflow.graph_dict, ), workflow_execution_repository=self._workflow_execution_repository, workflow_node_execution_repository=self._workflow_node_execution_repository, trace_manager=self.application_generate_entity.trace_manager, ) workflow_entry.graph_engine.layer(persistence_layer) generator = workflow_entry.run() for event in generator: self._update_document_status( event, self.application_generate_entity.document_id, self.application_generate_entity.dataset_id ) self._handle_event(workflow_entry, event) def get_workflow(self, pipeline: Pipeline, workflow_id: str) -> Workflow | None: """ Get workflow """ # fetch workflow by workflow_id workflow = db.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: GraphRuntimeState, 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, ) graph_init_context = DifyGraphInitContext( 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, graph_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: GraphEngineEvent, document_id: str | None, dataset_id: str | None) -> None: """ Update document status """ if isinstance(event, GraphRunFailedEvent): if document_id and dataset_id: document = db.session.scalar( select(Document).where(Document.id == document_id, Document.dataset_id == dataset_id).limit(1) ) if document: document.indexing_status = "error" document.error = event.error or "Unknown error" db.session.add(document) db.session.commit()