From e79b2351259da0880e9fe1335111e55bd888557e Mon Sep 17 00:00:00 2001 From: ruirui6946 <2733936092@qq.com> Date: Tue, 14 Jul 2026 21:03:52 +0800 Subject: [PATCH] fix(flow): detect tool calls after streamed text --- flow/agent/multiagent/host/compose.go | 5 +- .../host/stream_tool_call_checker_test.go | 60 +++++++++++++++++++ flow/agent/multiagent/host/types.go | 12 +--- flow/agent/react/react.go | 17 +----- .../react/stream_tool_call_checker_test.go | 60 +++++++++++++++++++ 5 files changed, 126 insertions(+), 28 deletions(-) create mode 100644 flow/agent/multiagent/host/stream_tool_call_checker_test.go create mode 100644 flow/agent/react/stream_tool_call_checker_test.go diff --git a/flow/agent/multiagent/host/compose.go b/flow/agent/multiagent/host/compose.go index 6a65cfb51..24439bbc5 100644 --- a/flow/agent/multiagent/host/compose.go +++ b/flow/agent/multiagent/host/compose.go @@ -43,9 +43,6 @@ type state struct { // NewMultiAgent creates a new host multi-agent system. // -// IMPORTANT!! For models that don't output tool calls in the first streaming chunk (e.g. Claude) -// the default StreamToolCallChecker may not work properly since it only checks the first chunk for tool calls. -// In such cases, you need to implement a custom StreamToolCallChecker that can properly detect tool calls. func NewMultiAgent(ctx context.Context, config *MultiAgentConfig) (*MultiAgent, error) { if err := config.validate(); err != nil { return nil, err @@ -71,7 +68,7 @@ func NewMultiAgent(ctx context.Context, config *MultiAgentConfig) (*MultiAgent, } if toolCallChecker == nil { - toolCallChecker = firstChunkStreamToolCallChecker + toolCallChecker = defaultStreamToolCallChecker } g := compose.NewGraph[[]*schema.Message, *schema.Message]( diff --git a/flow/agent/multiagent/host/stream_tool_call_checker_test.go b/flow/agent/multiagent/host/stream_tool_call_checker_test.go new file mode 100644 index 000000000..49f8cd760 --- /dev/null +++ b/flow/agent/multiagent/host/stream_tool_call_checker_test.go @@ -0,0 +1,60 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package host + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/cloudwego/eino/schema" +) + +func TestDefaultStreamToolCallChecker(t *testing.T) { + tests := []struct { + name string + chunks []*schema.Message + want bool + }{ + { + name: "tool call after text", + chunks: []*schema.Message{ + {Role: schema.Assistant, Content: "thinking"}, + {Role: schema.Assistant, ToolCalls: []schema.ToolCall{{ID: "call_1"}}}, + }, + want: true, + }, + { + name: "text only", + chunks: []*schema.Message{ + {Role: schema.Assistant, Content: "final answer"}, + }, + want: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + sr := schema.StreamReaderFromArray(tt.chunks) + got, err := defaultStreamToolCallChecker(context.Background(), sr) + require.NoError(t, err) + assert.Equal(t, tt.want, got) + }) + } +} diff --git a/flow/agent/multiagent/host/types.go b/flow/agent/multiagent/host/types.go index 19d417a49..8c6fbd63d 100644 --- a/flow/agent/multiagent/host/types.go +++ b/flow/agent/multiagent/host/types.go @@ -89,9 +89,7 @@ type MultiAgentConfig struct { // - false if no tool calls and agent should stop // Note: This field only needs to be configured when using streaming mode // Note: The handler MUST close the modelOutput stream before returning - // Optional. By default, it checks if the first chunk contains tool calls. - // Note: The default implementation does not work well with Claude, which typically outputs tool calls after text content. - // Note: If your ChatModel doesn't output tool calls first, you can try adding prompts to constrain the model from generating extra text during the tool call. + // Optional. By default, it scans the stream until it finds a tool call or reaches EOF. StreamToolCallChecker func(ctx context.Context, modelOutput *schema.StreamReader[*schema.Message]) (bool, error) // Summarizer is the summarizer agent that will summarize the outputs of all the chosen specialist agents. @@ -179,7 +177,7 @@ type Summarizer struct { SystemPrompt string } -func firstChunkStreamToolCallChecker(_ context.Context, sr *schema.StreamReader[*schema.Message]) (bool, error) { +func defaultStreamToolCallChecker(_ context.Context, sr *schema.StreamReader[*schema.Message]) (bool, error) { defer sr.Close() for { @@ -194,11 +192,5 @@ func firstChunkStreamToolCallChecker(_ context.Context, sr *schema.StreamReader[ if len(msg.ToolCalls) > 0 { return true, nil } - - if len(msg.Content) == 0 { // skip empty chunks at the front - continue - } - - return false, nil } } diff --git a/flow/agent/react/react.go b/flow/agent/react/react.go index 794a41fe1..d335cee20 100644 --- a/flow/agent/react/react.go +++ b/flow/agent/react/react.go @@ -173,9 +173,7 @@ type AgentConfig struct { // - false if no tool calls and agent should stop // Note: This field only needs to be configured when using streaming mode // Note: The handler MUST close the modelOutput stream before returning - // Optional. By default, it checks if the first chunk contains tool calls. - // Note: The default implementation does not work well with Claude, which typically outputs tool calls after text content. - // Note: If your ChatModel doesn't output tool calls first, you can try adding prompts to constrain the model from generating extra text during the tool call. + // Optional. By default, it scans the stream until it finds a tool call or reaches EOF. StreamToolCallChecker func(ctx context.Context, modelOutput *schema.StreamReader[*schema.Message]) (bool, error) // GraphName is the graph name of the ReAct Agent. @@ -215,7 +213,7 @@ func NewPersonaModifier(persona string) MessageModifier { } } -func firstChunkStreamToolCallChecker(_ context.Context, sr *schema.StreamReader[*schema.Message]) (bool, error) { +func defaultStreamToolCallChecker(_ context.Context, sr *schema.StreamReader[*schema.Message]) (bool, error) { defer sr.Close() for { @@ -230,12 +228,6 @@ func firstChunkStreamToolCallChecker(_ context.Context, sr *schema.StreamReader[ if len(msg.ToolCalls) > 0 { return true, nil } - - if len(msg.Content) == 0 { // skip empty chunks at the front - continue - } - - return false, nil } } @@ -278,9 +270,6 @@ type Agent struct { // NewAgent creates a ReAct agent that feeds tool response into next round of Chat Model generation. // -// IMPORTANT!! For models that don't output tool calls in the first streaming chunk (e.g. Claude) -// the default StreamToolCallChecker may not work properly since it only checks the first chunk for tool calls. -// In such cases, you need to implement a custom StreamToolCallChecker that can properly detect tool calls. func NewAgent(ctx context.Context, config *AgentConfig) (_ *Agent, err error) { var ( chatModel model.BaseChatModel @@ -306,7 +295,7 @@ func NewAgent(ctx context.Context, config *AgentConfig) (_ *Agent, err error) { } if toolCallChecker == nil { - toolCallChecker = firstChunkStreamToolCallChecker + toolCallChecker = defaultStreamToolCallChecker } if toolInfos, err = genToolInfos(ctx, config.ToolsConfig); err != nil { diff --git a/flow/agent/react/stream_tool_call_checker_test.go b/flow/agent/react/stream_tool_call_checker_test.go new file mode 100644 index 000000000..0a812263e --- /dev/null +++ b/flow/agent/react/stream_tool_call_checker_test.go @@ -0,0 +1,60 @@ +/* + * Copyright 2026 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package react + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/cloudwego/eino/schema" +) + +func TestDefaultStreamToolCallChecker(t *testing.T) { + tests := []struct { + name string + chunks []*schema.Message + want bool + }{ + { + name: "tool call after text", + chunks: []*schema.Message{ + {Role: schema.Assistant, Content: "thinking"}, + {Role: schema.Assistant, ToolCalls: []schema.ToolCall{{ID: "call_1"}}}, + }, + want: true, + }, + { + name: "text only", + chunks: []*schema.Message{ + {Role: schema.Assistant, Content: "final answer"}, + }, + want: false, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + sr := schema.StreamReaderFromArray(tt.chunks) + got, err := defaultStreamToolCallChecker(context.Background(), sr) + require.NoError(t, err) + assert.Equal(t, tt.want, got) + }) + } +}