From e157c76ecf25914806bc7c8c2cc02f0c7ffdb8e1 Mon Sep 17 00:00:00 2001 From: Fangyin Cheng Date: Thu, 1 Feb 2024 14:53:22 +0800 Subject: [PATCH] feat: AWEL flow support branch operator --- dbgpt/core/awel/flow/flow_factory.py | 52 ++++++++++++++--- dbgpt/core/awel/trigger/http_trigger.py | 3 + .../core/interface/operators/llm_operator.py | 56 ++++++++++++++++++- dbgpt/model/utils/chatgpt_utils.py | 2 +- dbgpt/serve/flow/service/service.py | 3 + examples/awel/data_analyst_assistant.py | 2 + web/pages/flow/canvas/index.tsx | 2 + 7 files changed, 109 insertions(+), 11 deletions(-) diff --git a/dbgpt/core/awel/flow/flow_factory.py b/dbgpt/core/awel/flow/flow_factory.py index 4e6ef84ef..77acf1294 100644 --- a/dbgpt/core/awel/flow/flow_factory.py +++ b/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: diff --git a/dbgpt/core/awel/trigger/http_trigger.py b/dbgpt/core/awel/trigger/http_trigger.py index f995eb384..d1458a69c 100644 --- a/dbgpt/core/awel/trigger/http_trigger.py +++ b/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") diff --git a/dbgpt/core/interface/operators/llm_operator.py b/dbgpt/core/interface/operators/llm_operator.py index 8c1f9bb34..f13a8c734 100644 --- a/dbgpt/core/interface/operators/llm_operator.py +++ b/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. diff --git a/dbgpt/model/utils/chatgpt_utils.py b/dbgpt/model/utils/chatgpt_utils.py index 24b85d7c3..0267b8959 100644 --- a/dbgpt/model/utils/chatgpt_utils.py +++ b/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", ) diff --git a/dbgpt/serve/flow/service/service.py b/dbgpt/serve/flow/service/service.py index c167f8f64..82d7cbea2 100644 --- a/dbgpt/serve/flow/service/service.py +++ b/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) diff --git a/examples/awel/data_analyst_assistant.py b/examples/awel/data_analyst_assistant.py index cdd7967fb..99ce6eebc 100644 --- a/examples/awel/data_analyst_assistant.py +++ b/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, ) ) diff --git a/web/pages/flow/canvas/index.tsx b/web/pages/flow/canvas/index.tsx index 50bfbbc4d..919af8ef8 100644 --- a/web/pages/flow/canvas/index.tsx +++ b/web/pages/flow/canvas/index.tsx @@ -44,6 +44,8 @@ const Canvas: React.FC = () => { 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); }