feat: Parallel Execution of Nodes in Workflows (#8192)

Co-authored-by: StyleZhang <jasonapring2015@outlook.com>
Co-authored-by: Yi <yxiaoisme@gmail.com>
Co-authored-by: -LAN- <laipz8200@outlook.com>

api/core/workflow/nodes/answer/answer_node.py

@@ -1,9 +1,8 @@
-from typing import cast
+from collections.abc import Mapping, Sequence
+from typing import Any, cast
 
-from core.prompt.utils.prompt_template_parser import PromptTemplateParser
-from core.workflow.entities.base_node_data_entities import BaseNodeData
 from core.workflow.entities.node_entities import NodeRunResult, NodeType
-from core.workflow.entities.variable_pool import VariablePool
+from core.workflow.nodes.answer.answer_stream_generate_router import AnswerStreamGeneratorRouter
 from core.workflow.nodes.answer.entities import (
     AnswerNodeData,
     GenerateRouteChunk,
@@ -19,24 +18,26 @@ class AnswerNode(BaseNode):
     _node_data_cls = AnswerNodeData
     _node_type: NodeType = NodeType.ANSWER
 
-    def _run(self, variable_pool: VariablePool) -> NodeRunResult:
+    def _run(self) -> NodeRunResult:
         """
         Run node
-        :param variable_pool: variable pool
         :return:
         """
         node_data = self.node_data
         node_data = cast(AnswerNodeData, node_data)
 
         # generate routes
-        generate_routes = self.extract_generate_route_from_node_data(node_data)
+        generate_routes = AnswerStreamGeneratorRouter.extract_generate_route_from_node_data(node_data)
 
         answer = ''
         for part in generate_routes:
-            if part.type == "var":
+            if part.type == GenerateRouteChunk.ChunkType.VAR:
                 part = cast(VarGenerateRouteChunk, part)
                 value_selector = part.value_selector
-                value = variable_pool.get(value_selector)
+                value = self.graph_runtime_state.variable_pool.get(
+                    value_selector
+                )
+
                 if value:
                     answer += value.markdown
             else:
@@ -51,70 +52,16 @@ class AnswerNode(BaseNode):
         )
 
     @classmethod
-    def extract_generate_route_selectors(cls, config: dict) -> list[GenerateRouteChunk]:
-        """
-        Extract generate route selectors
-        :param config: node config
-        :return:
-        """
-        node_data = cls._node_data_cls(**config.get("data", {}))
-        node_data = cast(AnswerNodeData, node_data)
-
-        return cls.extract_generate_route_from_node_data(node_data)
-
-    @classmethod
-    def extract_generate_route_from_node_data(cls, node_data: AnswerNodeData) -> list[GenerateRouteChunk]:
-        """
-        Extract generate route from node data
-        :param node_data: node data object
-        :return:
-        """
-        variable_template_parser = VariableTemplateParser(template=node_data.answer)
-        variable_selectors = variable_template_parser.extract_variable_selectors()
-
-        value_selector_mapping = {
-            variable_selector.variable: variable_selector.value_selector
-            for variable_selector in variable_selectors
-        }
-
-        variable_keys = list(value_selector_mapping.keys())
-
-        # format answer template
-        template_parser = PromptTemplateParser(template=node_data.answer, with_variable_tmpl=True)
-        template_variable_keys = template_parser.variable_keys
-
-        # Take the intersection of variable_keys and template_variable_keys
-        variable_keys = list(set(variable_keys) & set(template_variable_keys))
-
-        template = node_data.answer
-        for var in variable_keys:
-            template = template.replace(f'{{{{{var}}}}}', f'Ω{{{{{var}}}}}Ω')
-
-        generate_routes = []
-        for part in template.split('Ω'):
-            if part:
-                if cls._is_variable(part, variable_keys):
-                    var_key = part.replace('Ω', '').replace('{{', '').replace('}}', '')
-                    value_selector = value_selector_mapping[var_key]
-                    generate_routes.append(VarGenerateRouteChunk(
-                        value_selector=value_selector
-                    ))
-                else:
-                    generate_routes.append(TextGenerateRouteChunk(
-                        text=part
-                    ))
-
-        return generate_routes
-
-    @classmethod
-    def _is_variable(cls, part, variable_keys):
-        cleaned_part = part.replace('{{', '').replace('}}', '')
-        return part.startswith('{{') and cleaned_part in variable_keys
-
-    @classmethod
-    def _extract_variable_selector_to_variable_mapping(cls, node_data: BaseNodeData) -> dict[str, list[str]]:
+    def _extract_variable_selector_to_variable_mapping(
+        cls, 
+        graph_config: Mapping[str, Any], 
+        node_id: str,
+        node_data: AnswerNodeData
+    ) -> Mapping[str, Sequence[str]]:
         """
         Extract variable selector to variable mapping
+        :param graph_config: graph config
+        :param node_id: node id
         :param node_data: node data
         :return:
         """
@@ -126,6 +73,6 @@ class AnswerNode(BaseNode):
 
         variable_mapping = {}
         for variable_selector in variable_selectors:
-            variable_mapping[variable_selector.variable] = variable_selector.value_selector
+            variable_mapping[node_id + '.' + variable_selector.variable] = variable_selector.value_selector
 
         return variable_mapping

api/core/workflow/nodes/answer/answer_stream_generate_router.py

@@ -0,0 +1,169 @@
+
+from core.prompt.utils.prompt_template_parser import PromptTemplateParser
+from core.workflow.entities.node_entities import NodeType
+from core.workflow.nodes.answer.entities import (
+    AnswerNodeData,
+    AnswerStreamGenerateRoute,
+    GenerateRouteChunk,
+    TextGenerateRouteChunk,
+    VarGenerateRouteChunk,
+)
+from core.workflow.utils.variable_template_parser import VariableTemplateParser
+
+
+class AnswerStreamGeneratorRouter:
+
+    @classmethod
+    def init(cls,
+             node_id_config_mapping: dict[str, dict],
+             reverse_edge_mapping: dict[str, list["GraphEdge"]]  # type: ignore[name-defined]
+             ) -> AnswerStreamGenerateRoute:
+        """
+        Get stream generate routes.
+        :return:
+        """
+        # parse stream output node value selectors of answer nodes
+        answer_generate_route: dict[str, list[GenerateRouteChunk]] = {}
+        for answer_node_id, node_config in node_id_config_mapping.items():
+            if not node_config.get('data', {}).get('type') == NodeType.ANSWER.value:
+                continue
+
+            # get generate route for stream output
+            generate_route = cls._extract_generate_route_selectors(node_config)
+            answer_generate_route[answer_node_id] = generate_route
+
+        # fetch answer dependencies
+        answer_node_ids = list(answer_generate_route.keys())
+        answer_dependencies = cls._fetch_answers_dependencies(
+            answer_node_ids=answer_node_ids,
+            reverse_edge_mapping=reverse_edge_mapping,
+            node_id_config_mapping=node_id_config_mapping
+        )
+
+        return AnswerStreamGenerateRoute(
+            answer_generate_route=answer_generate_route,
+            answer_dependencies=answer_dependencies
+        )
+
+    @classmethod
+    def extract_generate_route_from_node_data(cls, node_data: AnswerNodeData) -> list[GenerateRouteChunk]:
+        """
+        Extract generate route from node data
+        :param node_data: node data object
+        :return:
+        """
+        variable_template_parser = VariableTemplateParser(template=node_data.answer)
+        variable_selectors = variable_template_parser.extract_variable_selectors()
+
+        value_selector_mapping = {
+            variable_selector.variable: variable_selector.value_selector
+            for variable_selector in variable_selectors
+        }
+
+        variable_keys = list(value_selector_mapping.keys())
+
+        # format answer template
+        template_parser = PromptTemplateParser(template=node_data.answer, with_variable_tmpl=True)
+        template_variable_keys = template_parser.variable_keys
+
+        # Take the intersection of variable_keys and template_variable_keys
+        variable_keys = list(set(variable_keys) & set(template_variable_keys))
+
+        template = node_data.answer
+        for var in variable_keys:
+            template = template.replace(f'{{{{{var}}}}}', f'Ω{{{{{var}}}}}Ω')
+
+        generate_routes: list[GenerateRouteChunk] = []
+        for part in template.split('Ω'):
+            if part:
+                if cls._is_variable(part, variable_keys):
+                    var_key = part.replace('Ω', '').replace('{{', '').replace('}}', '')
+                    value_selector = value_selector_mapping[var_key]
+                    generate_routes.append(VarGenerateRouteChunk(
+                        value_selector=value_selector
+                    ))
+                else:
+                    generate_routes.append(TextGenerateRouteChunk(
+                        text=part
+                    ))
+
+        return generate_routes
+
+    @classmethod
+    def _extract_generate_route_selectors(cls, config: dict) -> list[GenerateRouteChunk]:
+        """
+        Extract generate route selectors
+        :param config: node config
+        :return:
+        """
+        node_data = AnswerNodeData(**config.get("data", {}))
+        return cls.extract_generate_route_from_node_data(node_data)
+
+    @classmethod
+    def _is_variable(cls, part, variable_keys):
+        cleaned_part = part.replace('{{', '').replace('}}', '')
+        return part.startswith('{{') and cleaned_part in variable_keys
+
+    @classmethod
+    def _fetch_answers_dependencies(cls,
+                                    answer_node_ids: list[str],
+                                    reverse_edge_mapping: dict[str, list["GraphEdge"]],  # type: ignore[name-defined]
+                                    node_id_config_mapping: dict[str, dict]
+                                    ) -> dict[str, list[str]]:
+        """
+        Fetch answer dependencies
+        :param answer_node_ids: answer node ids
+        :param reverse_edge_mapping: reverse edge mapping
+        :param node_id_config_mapping: node id config mapping
+        :return:
+        """
+        answer_dependencies: dict[str, list[str]] = {}
+        for answer_node_id in answer_node_ids:
+            if answer_dependencies.get(answer_node_id) is None:
+                answer_dependencies[answer_node_id] = []
+
+            cls._recursive_fetch_answer_dependencies(
+                current_node_id=answer_node_id,
+                answer_node_id=answer_node_id,
+                node_id_config_mapping=node_id_config_mapping,
+                reverse_edge_mapping=reverse_edge_mapping,
+                answer_dependencies=answer_dependencies
+            )
+
+        return answer_dependencies
+
+    @classmethod
+    def _recursive_fetch_answer_dependencies(cls,
+                                             current_node_id: str,
+                                             answer_node_id: str,
+                                             node_id_config_mapping: dict[str, dict],
+                                             reverse_edge_mapping: dict[str, list["GraphEdge"]],  # type: ignore[name-defined]
+                                             answer_dependencies: dict[str, list[str]]
+                                             ) -> None:
+        """
+        Recursive fetch answer dependencies
+        :param current_node_id: current node id
+        :param answer_node_id: answer node id
+        :param node_id_config_mapping: node id config mapping
+        :param reverse_edge_mapping: reverse edge mapping
+        :param answer_dependencies: answer dependencies
+        :return:
+        """
+        reverse_edges = reverse_edge_mapping.get(current_node_id, [])
+        for edge in reverse_edges:
+            source_node_id = edge.source_node_id
+            source_node_type = node_id_config_mapping[source_node_id].get('data', {}).get('type')
+            if source_node_type in (
+                    NodeType.ANSWER.value,
+                    NodeType.IF_ELSE.value,
+                    NodeType.QUESTION_CLASSIFIER,
+            ):
+                answer_dependencies[answer_node_id].append(source_node_id)
+            else:
+                cls._recursive_fetch_answer_dependencies(
+                    current_node_id=source_node_id,
+                    answer_node_id=answer_node_id,
+                    node_id_config_mapping=node_id_config_mapping,
+                    reverse_edge_mapping=reverse_edge_mapping,
+                    answer_dependencies=answer_dependencies
+                )

api/core/workflow/nodes/answer/answer_stream_processor.py

@@ -0,0 +1,221 @@
+import logging
+from collections.abc import Generator
+from typing import Optional, cast
+
+from core.file.file_obj import FileVar
+from core.workflow.entities.variable_pool import VariablePool
+from core.workflow.graph_engine.entities.event import (
+    GraphEngineEvent,
+    NodeRunStartedEvent,
+    NodeRunStreamChunkEvent,
+    NodeRunSucceededEvent,
+)
+from core.workflow.graph_engine.entities.graph import Graph
+from core.workflow.nodes.answer.base_stream_processor import StreamProcessor
+from core.workflow.nodes.answer.entities import GenerateRouteChunk, TextGenerateRouteChunk, VarGenerateRouteChunk
+
+logger = logging.getLogger(__name__)
+
+
+class AnswerStreamProcessor(StreamProcessor):
+
+    def __init__(self, graph: Graph, variable_pool: VariablePool) -> None:
+        super().__init__(graph, variable_pool)
+        self.generate_routes = graph.answer_stream_generate_routes
+        self.route_position = {}
+        for answer_node_id, route_chunks in self.generate_routes.answer_generate_route.items():
+            self.route_position[answer_node_id] = 0
+        self.current_stream_chunk_generating_node_ids: dict[str, list[str]] = {}
+
+    def process(self,
+                generator: Generator[GraphEngineEvent, None, None]
+                ) -> Generator[GraphEngineEvent, None, None]:
+        for event in generator:
+            if isinstance(event, NodeRunStartedEvent):
+                if event.route_node_state.node_id == self.graph.root_node_id and not self.rest_node_ids:
+                    self.reset()
+
+                yield event
+            elif isinstance(event, NodeRunStreamChunkEvent):
+                if event.in_iteration_id:
+                    yield event
+                    continue
+
+                if event.route_node_state.node_id in self.current_stream_chunk_generating_node_ids:
+                    stream_out_answer_node_ids = self.current_stream_chunk_generating_node_ids[
+                        event.route_node_state.node_id
+                    ]
+                else:
+                    stream_out_answer_node_ids = self._get_stream_out_answer_node_ids(event)
+                    self.current_stream_chunk_generating_node_ids[
+                        event.route_node_state.node_id
+                    ] = stream_out_answer_node_ids
+
+                for _ in stream_out_answer_node_ids:
+                    yield event
+            elif isinstance(event, NodeRunSucceededEvent):
+                yield event
+                if event.route_node_state.node_id in self.current_stream_chunk_generating_node_ids:
+                    # update self.route_position after all stream event finished
+                    for answer_node_id in self.current_stream_chunk_generating_node_ids[event.route_node_state.node_id]:
+                        self.route_position[answer_node_id] += 1
+
+                    del self.current_stream_chunk_generating_node_ids[event.route_node_state.node_id]
+
+                # remove unreachable nodes
+                self._remove_unreachable_nodes(event)
+
+                # generate stream outputs
+                yield from self._generate_stream_outputs_when_node_finished(event)
+            else:
+                yield event
+
+    def reset(self) -> None:
+        self.route_position = {}
+        for answer_node_id, route_chunks in self.generate_routes.answer_generate_route.items():
+            self.route_position[answer_node_id] = 0
+        self.rest_node_ids = self.graph.node_ids.copy()
+        self.current_stream_chunk_generating_node_ids = {}
+
+    def _generate_stream_outputs_when_node_finished(self,
+                                                    event: NodeRunSucceededEvent
+                                                    ) -> Generator[GraphEngineEvent, None, None]:
+        """
+        Generate stream outputs.
+        :param event: node run succeeded event
+        :return:
+        """
+        for answer_node_id, position in self.route_position.items():
+            # all depends on answer node id not in rest node ids
+            if (event.route_node_state.node_id != answer_node_id
+                    and (answer_node_id not in self.rest_node_ids
+                         or not all(dep_id not in self.rest_node_ids
+                                    for dep_id in self.generate_routes.answer_dependencies[answer_node_id]))):
+                continue
+
+            route_position = self.route_position[answer_node_id]
+            route_chunks = self.generate_routes.answer_generate_route[answer_node_id][route_position:]
+
+            for route_chunk in route_chunks:
+                if route_chunk.type == GenerateRouteChunk.ChunkType.TEXT:
+                    route_chunk = cast(TextGenerateRouteChunk, route_chunk)
+                    yield NodeRunStreamChunkEvent(
+                        id=event.id,
+                        node_id=event.node_id,
+                        node_type=event.node_type,
+                        node_data=event.node_data,
+                        chunk_content=route_chunk.text,
+                        route_node_state=event.route_node_state,
+                        parallel_id=event.parallel_id,
+                        parallel_start_node_id=event.parallel_start_node_id,
+                    )
+                else:
+                    route_chunk = cast(VarGenerateRouteChunk, route_chunk)
+                    value_selector = route_chunk.value_selector
+                    if not value_selector:
+                        break
+
+                    value = self.variable_pool.get(
+                        value_selector
+                    )
+
+                    if value is None:
+                        break
+
+                    text = value.markdown
+
+                    if text:
+                        yield NodeRunStreamChunkEvent(
+                            id=event.id,
+                            node_id=event.node_id,
+                            node_type=event.node_type,
+                            node_data=event.node_data,
+                            chunk_content=text,
+                            from_variable_selector=value_selector,
+                            route_node_state=event.route_node_state,
+                            parallel_id=event.parallel_id,
+                            parallel_start_node_id=event.parallel_start_node_id,
+                        )
+
+                self.route_position[answer_node_id] += 1
+
+    def _get_stream_out_answer_node_ids(self, event: NodeRunStreamChunkEvent) -> list[str]:
+        """
+        Is stream out support
+        :param event: queue text chunk event
+        :return:
+        """
+        if not event.from_variable_selector:
+            return []
+
+        stream_output_value_selector = event.from_variable_selector
+        if not stream_output_value_selector:
+            return []
+
+        stream_out_answer_node_ids = []
+        for answer_node_id, route_position in self.route_position.items():
+            if answer_node_id not in self.rest_node_ids:
+                continue
+
+            # all depends on answer node id not in rest node ids
+            if all(dep_id not in self.rest_node_ids
+                   for dep_id in self.generate_routes.answer_dependencies[answer_node_id]):
+                if route_position >= len(self.generate_routes.answer_generate_route[answer_node_id]):
+                    continue
+
+                route_chunk = self.generate_routes.answer_generate_route[answer_node_id][route_position]
+
+                if route_chunk.type != GenerateRouteChunk.ChunkType.VAR:
+                    continue
+
+                route_chunk = cast(VarGenerateRouteChunk, route_chunk)
+                value_selector = route_chunk.value_selector
+
+                # check chunk node id is before current node id or equal to current node id
+                if value_selector != stream_output_value_selector:
+                    continue
+
+                stream_out_answer_node_ids.append(answer_node_id)
+
+        return stream_out_answer_node_ids
+
+    @classmethod
+    def _fetch_files_from_variable_value(cls, value: dict | list) -> list[dict]:
+        """
+        Fetch files from variable value
+        :param value: variable value
+        :return:
+        """
+        if not value:
+            return []
+
+        files = []
+        if isinstance(value, list):
+            for item in value:
+                file_var = cls._get_file_var_from_value(item)
+                if file_var:
+                    files.append(file_var)
+        elif isinstance(value, dict):
+            file_var = cls._get_file_var_from_value(value)
+            if file_var:
+                files.append(file_var)
+
+        return files
+
+    @classmethod
+    def _get_file_var_from_value(cls, value: dict | list) -> Optional[dict]:
+        """
+        Get file var from value
+        :param value: variable value
+        :return:
+        """
+        if not value:
+            return None
+
+        if isinstance(value, dict):
+            if '__variant' in value and value['__variant'] == FileVar.__name__:
+                return value
+        elif isinstance(value, FileVar):
+            return value.to_dict()
+
+        return None

api/core/workflow/nodes/answer/base_stream_processor.py

@@ -0,0 +1,71 @@
+from abc import ABC, abstractmethod
+from collections.abc import Generator
+
+from core.workflow.entities.variable_pool import VariablePool
+from core.workflow.graph_engine.entities.event import GraphEngineEvent, NodeRunSucceededEvent
+from core.workflow.graph_engine.entities.graph import Graph
+
+
+class StreamProcessor(ABC):
+
+    def __init__(self, graph: Graph, variable_pool: VariablePool) -> None:
+        self.graph = graph
+        self.variable_pool = variable_pool
+        self.rest_node_ids = graph.node_ids.copy()
+
+    @abstractmethod
+    def process(self,
+                generator: Generator[GraphEngineEvent, None, None]
+                ) -> Generator[GraphEngineEvent, None, None]:
+        raise NotImplementedError
+
+    def _remove_unreachable_nodes(self, event: NodeRunSucceededEvent) -> None:
+        finished_node_id = event.route_node_state.node_id
+        if finished_node_id not in self.rest_node_ids:
+            return
+
+        # remove finished node id
+        self.rest_node_ids.remove(finished_node_id)
+
+        run_result = event.route_node_state.node_run_result
+        if not run_result:
+            return
+
+        if run_result.edge_source_handle:
+            reachable_node_ids = []
+            unreachable_first_node_ids = []
+            for edge in self.graph.edge_mapping[finished_node_id]:
+                if (edge.run_condition
+                        and edge.run_condition.branch_identify
+                        and run_result.edge_source_handle == edge.run_condition.branch_identify):
+                    reachable_node_ids.extend(self._fetch_node_ids_in_reachable_branch(edge.target_node_id))
+                    continue
+                else:
+                    unreachable_first_node_ids.append(edge.target_node_id)
+
+            for node_id in unreachable_first_node_ids:
+                self._remove_node_ids_in_unreachable_branch(node_id, reachable_node_ids)
+
+    def _fetch_node_ids_in_reachable_branch(self, node_id: str) -> list[str]:
+        node_ids = []
+        for edge in self.graph.edge_mapping.get(node_id, []):
+            if edge.target_node_id == self.graph.root_node_id:
+                continue
+
+            node_ids.append(edge.target_node_id)
+            node_ids.extend(self._fetch_node_ids_in_reachable_branch(edge.target_node_id))
+        return node_ids
+
+    def _remove_node_ids_in_unreachable_branch(self, node_id: str, reachable_node_ids: list[str]) -> None:
+        """
+        remove target node ids until merge
+        """
+        if node_id not in self.rest_node_ids:
+            return
+
+        self.rest_node_ids.remove(node_id)
+        for edge in self.graph.edge_mapping.get(node_id, []):
+            if edge.target_node_id in reachable_node_ids:
+                continue
+
+            self._remove_node_ids_in_unreachable_branch(edge.target_node_id, reachable_node_ids)

api/core/workflow/nodes/answer/entities.py

@@ -1,5 +1,6 @@
+from enum import Enum
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 from core.workflow.entities.base_node_data_entities import BaseNodeData
 
@@ -8,27 +9,54 @@ class AnswerNodeData(BaseNodeData):
     """
     Answer Node Data.
     """
-    answer: str
+    answer: str = Field(..., description="answer template string")
 
 
 class GenerateRouteChunk(BaseModel):
     """
     Generate Route Chunk.
     """
-    type: str
+
+    class ChunkType(Enum):
+        VAR = "var"
+        TEXT = "text"
+
+    type: ChunkType = Field(..., description="generate route chunk type")
 
 
 class VarGenerateRouteChunk(GenerateRouteChunk):
     """
     Var Generate Route Chunk.
     """
-    type: str = "var"
-    value_selector: list[str]
+    type: GenerateRouteChunk.ChunkType = GenerateRouteChunk.ChunkType.VAR
+    """generate route chunk type"""
+    value_selector: list[str] = Field(..., description="value selector")
 
 
 class TextGenerateRouteChunk(GenerateRouteChunk):
     """
     Text Generate Route Chunk.
     """
-    type: str = "text"
-    text: str
+    type: GenerateRouteChunk.ChunkType = GenerateRouteChunk.ChunkType.TEXT
+    """generate route chunk type"""
+    text: str = Field(..., description="text")
+
+
+class AnswerNodeDoubleLink(BaseModel):
+    node_id: str = Field(..., description="node id")
+    source_node_ids: list[str] = Field(..., description="source node ids")
+    target_node_ids: list[str] = Field(..., description="target node ids")
+
+
+class AnswerStreamGenerateRoute(BaseModel):
+    """
+    AnswerStreamGenerateRoute entity
+    """
+    answer_dependencies: dict[str, list[str]] = Field(
+        ...,
+        description="answer dependencies (answer node id -> dependent answer node ids)"
+    )
+    answer_generate_route: dict[str, list[GenerateRouteChunk]] = Field(
+        ...,
+        description="answer generate route (answer node id -> generate route chunks)"
+    )