feat: AWEL flow support branch operator

dbgpt/core/awel/flow/flow_factory.py

@@ -4,12 +4,12 @@ import logging
 import uuid
 from contextlib import suppress
 from enum import Enum
-from typing import Any, Dict, List, Optional, Type, Union, cast
+from typing import Any, Dict, List, Optional, Tuple, Type, Union, cast
 
 from dbgpt._private.pydantic import BaseModel, Field, root_validator, validator
 from dbgpt.core.awel.dag.base import DAG, DAGNode
 
-from .base import ResourceMetadata, ViewMetadata, _get_operator_class
+from .base import OperatorType, ResourceMetadata, ViewMetadata, _get_operator_class
 
 logger = logging.getLogger(__name__)
 
@@ -76,6 +76,10 @@ class FlowEdgeData(BaseModel):
         description="Source node data id",
         examples=["resource_dbgpt.model.proxy.llms.chatgpt.OpenAILLMClient_0"],
     )
+    source_order: int = Field(
+        description="The order of the source node in the source node's output",
+        examples=[0, 1],
+    )
     target: str = Field(
         ...,
         description="Target node data id",
@@ -105,6 +109,13 @@ class FlowEdgeData(BaseModel):
     @root_validator(pre=True)
     def pre_fill(cls, values: Dict[str, Any]) -> Dict[str, Any]:
         """Pre fill the metadata."""
+        if (
+            "source_order" not in values
+            and "source_handle" in values
+            and values["source_handle"] is not None
+        ):
+            with suppress(Exception):
+                values["source_order"] = int(values["source_handle"].split("|")[-1])
         if (
             "target_order" not in values
             and "target_handle" in values
@@ -231,8 +242,8 @@ class FlowFactory:
         key_to_operator_nodes: Dict[str, FlowNodeData] = {}
         key_to_resource_nodes: Dict[str, FlowNodeData] = {}
         key_to_resource: Dict[str, ResourceMetadata] = {}
-        key_to_downstream: Dict[str, List[str]] = {}
-        key_to_upstream: Dict[str, List[str]] = {}
+        key_to_downstream: Dict[str, List[Tuple[str, int, int]]] = {}
+        key_to_upstream: Dict[str, List[Tuple[str, int, int]]] = {}
         for node in flow_data.nodes:
             key = node.id
             if key in key_to_operator_nodes or key in key_to_resource_nodes:
@@ -263,11 +274,11 @@ class FlowFactory:
             if source_node.data.is_operator and target_node.data.is_operator:
                 # Operator to operator.
                 downstream = key_to_downstream.get(source_key, [])
-                downstream.append(target_key)
+                downstream.append((target_key, edge.source_order, edge.target_order))
                 key_to_downstream[source_key] = downstream
 
                 upstream = key_to_upstream.get(target_key, [])
-                upstream.append(source_key)
+                upstream.append((source_key, edge.source_order, edge.target_order))
                 key_to_upstream[target_key] = upstream
             elif not source_node.data.is_operator and target_node.data.is_operator:
                 # Resource to operator.
@@ -296,6 +307,14 @@ class FlowFactory:
             else:
                 raise ValueError("Unable to connect resource to resource.")
 
+        # Sort the keys by the order of the nodes.
+        for key, value in key_to_downstream.items():
+            # Sort by source_order.
+            key_to_downstream[key] = sorted(value, key=lambda x: x[1])
+        for key, value in key_to_upstream.items():
+            # Sort by target_order.
+            key_to_upstream[key] = sorted(value, key=lambda x: x[2])
+
         key_to_tasks: Dict[str, DAGNode] = {}
         for operator_key, node in key_to_operator_nodes.items():
             if not isinstance(node.data, ViewMetadata):
@@ -307,6 +326,21 @@ class FlowFactory:
             metadata = operator_cls.metadata
             if not metadata:
                 raise ValueError("Metadata is not set.")
+            if metadata.operator_type == OperatorType.BRANCH:
+                # Branch operator, we suppose than the task_name of downstream is the
+                # parameter value of the branch operator.
+                downstream = key_to_downstream.get(operator_key, [])
+                if not downstream:
+                    raise ValueError("Branch operator should have downstream.")
+                if len(downstream) != len(view_metadata.parameters):
+                    raise ValueError(
+                        "Branch operator should have the same number of downstream as "
+                        "parameters."
+                    )
+                for i, param in enumerate(view_metadata.parameters):
+                    downstream_key, _, _ = downstream[i]
+                    param.value = key_to_operator_nodes[downstream_key].data.name
+
             runnable_params = metadata.get_runnable_parameters(
                 view_metadata.parameters, key_to_resource
             )
@@ -326,8 +360,8 @@ class FlowFactory:
         self,
         flow_panel: FlowPanel,
         key_to_tasks: Dict[str, DAGNode],
-        key_to_downstream: Dict[str, List[str]],
-        key_to_upstream: Dict[str, List[str]],
+        key_to_downstream: Dict[str, List[Tuple[str, int, int]]],
+        key_to_upstream: Dict[str, List[Tuple[str, int, int]]],
         dag_id: Optional[str] = None,
     ) -> DAG:
         """Build the DAG."""
@@ -345,7 +379,7 @@ class FlowFactory:
                     # A single task.
                     dag._append_node(task)
                     continue
-                for downstream_key in downstream:
+                for downstream_key, _, _ in downstream:
                     # Just one direction.
                     downstream_task = key_to_tasks.get(downstream_key)
                     if not downstream_task:

dbgpt/core/awel/trigger/http_trigger.py

@@ -497,6 +497,9 @@ class HttpTrigger(Trigger):
             streaming_response = self._streaming_response
             if self._streaming_predict_func:
                 streaming_response = self._streaming_predict_func(body)
+            elif isinstance(body, BaseHttpBody):
+                # BaseHttpBody, read streaming flag from body
+                streaming_response = _default_streaming_predict_func(body)
             dag = self.dag
             if not dag:
                 raise AWELHttpError("DAG is not set")

dbgpt/core/interface/operators/llm_operator.py

@@ -16,7 +16,13 @@ from dbgpt.core.awel import (
     MapOperator,
     StreamifyAbsOperator,
 )
-from dbgpt.core.awel.flow import IOField, OperatorCategory, Parameter, ViewMetadata
+from dbgpt.core.awel.flow import (
+    IOField,
+    OperatorCategory,
+    OperatorType,
+    Parameter,
+    ViewMetadata,
+)
 from dbgpt.core.interface.llm import (
     LLMClient,
     ModelOutput,
@@ -274,6 +280,54 @@ class LLMBranchOperator(BranchOperator[ModelRequest, ModelRequest]):
     the stream flag of the request.
     """
 
+    metadata = ViewMetadata(
+        label="LLM Branch Operator",
+        name="llm_branch_operator",
+        category=OperatorCategory.LLM,
+        operator_type=OperatorType.BRANCH,
+        description="Branch the workflow based on the stream flag of the request.",
+        parameters=[
+            Parameter.build_from(
+                "Streaming Task Name",
+                "stream_task_name",
+                str,
+                optional=True,
+                default="streaming_llm_task",
+                description="The name of the streaming task.",
+            ),
+            Parameter.build_from(
+                "Non-Streaming Task Name",
+                "no_stream_task_name",
+                str,
+                optional=True,
+                default="llm_task",
+                description="The name of the non-streaming task.",
+            ),
+        ],
+        inputs=[
+            IOField.build_from(
+                "Model Request",
+                "input_value",
+                ModelRequest,
+                description="The input value of the operator.",
+            ),
+        ],
+        outputs=[
+            IOField.build_from(
+                "Streaming Model Request",
+                "streaming_request",
+                ModelRequest,
+                description="The streaming request, to streaming Operator.",
+            ),
+            IOField.build_from(
+                "Non-Streaming Model Request",
+                "no_streaming_request",
+                ModelRequest,
+                description="The non-streaming request, to non-streaming Operator.",
+            ),
+        ],
+    )
+
     def __init__(self, stream_task_name: str, no_stream_task_name: str, **kwargs):
         """Create a new LLM branch operator.
 

dbgpt/model/utils/chatgpt_utils.py

@@ -176,7 +176,7 @@ class OpenAIStreamingOutputOperator(TransformStreamAbsOperator[ModelOutput, str]
             IOField.build_from(
                 "Model Output",
                 "model_output",
-                ModelOutput,
+                str,
                 is_list=True,
                 description="The model output after transform to openai stream format",
             )

dbgpt/serve/flow/service/service.py

@@ -99,6 +99,9 @@ class Service(BaseService[ServeEntity, ServeRequest, ServerResponse]):
         Returns:
             ServerResponse: The response
         """
+        # Try to build the dag from the request
+        self._flow_factory.build(request)
+
         # Build the query request from the request
         query_request = {"uid": request.uid}
         inst = self.get(query_request)

examples/awel/data_analyst_assistant.py

@@ -354,6 +354,8 @@ with DAG("dbgpt_awel_data_analyst_assistant") as dag:
             lambda req: ModelRequestContext(
                 conv_uid=req.context.conv_uid,
                 stream=req.stream,
+                user_name=req.context.user_name,
+                sys_code=req.context.sys_code,
                 chat_mode=req.context.chat_mode,
             )
         )

web/pages/flow/canvas/index.tsx

@@ -44,6 +44,8 @@ const Canvas: React.FC<Props> = () => {
       const flowData = mapUnderlineToHump(data.flow_data);
       setName(data.name);
       setDescription(data.description);
+      setNodes(flowData.nodes.map((node) => ({ ...node, type: 'customNode' })));
+      setEdges(flowData.edges.map((edge) => ({ ...edge, type: 'buttonedge' })));
     }
     setLoading(false);
   }