fix(core): Fix AWEL branch bug (#1640)

examples/awel/data_analyst_assistant.py

@@ -36,6 +36,7 @@
             }'
 
 """
+
 import logging
 import os
 from functools import cache
@@ -53,12 +54,17 @@ from dbgpt.core import (
     PromptTemplate,
     SystemPromptTemplate,
 )
-from dbgpt.core.awel import DAG, HttpTrigger, JoinOperator, MapOperator
+from dbgpt.core.awel import (
+    DAG,
+    BranchJoinOperator,
+    HttpTrigger,
+    JoinOperator,
+    MapOperator,
+)
 from dbgpt.core.operators import (
     BufferedConversationMapperOperator,
     HistoryDynamicPromptBuilderOperator,
     LLMBranchOperator,
-    RequestBuilderOperator,
 )
 from dbgpt.model.operators import (
     LLMOperator,
@@ -318,7 +324,6 @@ with DAG("dbgpt_awel_data_analyst_assistant") as dag:
     )
 
     prompt_template_load_task = PromptTemplateBuilderOperator()
-    request_handle_task = RequestBuilderOperator()
 
     # Load and store chat history
     chat_history_load_task = ServePreChatHistoryLoadOperator()
@@ -343,9 +348,7 @@ with DAG("dbgpt_awel_data_analyst_assistant") as dag:
     )
     model_parse_task = MapOperator(lambda out: out.to_dict())
     openai_format_stream_task = OpenAIStreamingOutputOperator()
-    result_join_task = JoinOperator(
-        combine_function=lambda not_stream_out, stream_out: not_stream_out or stream_out
-    )
+    result_join_task = BranchJoinOperator()
     trigger >> prompt_template_load_task >> history_prompt_build_task
 
     (

examples/awel/simple_chat_history_example.py

@@ -55,6 +55,7 @@
 
 
 """
+
 import logging
 from typing import Dict, List, Optional, Union
 
@@ -69,7 +70,7 @@ from dbgpt.core import (
     ModelRequestContext,
     SystemPromptTemplate,
 )
-from dbgpt.core.awel import DAG, HttpTrigger, JoinOperator, MapOperator
+from dbgpt.core.awel import DAG, BranchJoinOperator, HttpTrigger, MapOperator
 from dbgpt.core.operators import (
     ChatComposerInput,
     ChatHistoryPromptComposerOperator,
@@ -150,9 +151,7 @@ with DAG("dbgpt_awel_simple_chat_history") as multi_round_dag:
     )
     model_parse_task = MapOperator(lambda out: out.to_dict())
     openai_format_stream_task = OpenAIStreamingOutputOperator()
-    result_join_task = JoinOperator(
-        combine_function=lambda not_stream_out, stream_out: not_stream_out or stream_out
-    )
+    result_join_task = BranchJoinOperator()
 
     req_handle_task = MapOperator(
         lambda req: ChatComposerInput(

examples/awel/simple_llm_client_example.py

@@ -35,12 +35,13 @@
             }'
 
 """
+
 import logging
 from typing import Any, Dict, List, Optional, Union
 
 from dbgpt._private.pydantic import BaseModel, Field
 from dbgpt.core import LLMClient
-from dbgpt.core.awel import DAG, HttpTrigger, JoinOperator, MapOperator
+from dbgpt.core.awel import DAG, BranchJoinOperator, HttpTrigger, MapOperator
 from dbgpt.core.operators import LLMBranchOperator, RequestBuilderOperator
 from dbgpt.model.operators import (
     LLMOperator,
@@ -94,9 +95,7 @@ with DAG("dbgpt_awel_simple_llm_client_generate") as client_generate_dag:
     )
     model_parse_task = MapOperator(lambda out: out.to_dict())
     openai_format_stream_task = OpenAIStreamingOutputOperator()
-    result_join_task = JoinOperator(
-        combine_function=lambda not_stream_out, stream_out: not_stream_out or stream_out
-    )
+    result_join_task = BranchJoinOperator()
 
     trigger >> request_handle_task >> branch_task
     branch_task >> llm_task >> model_parse_task >> result_join_task