diff --git a/workflow/agentworkflow/message_merger.go b/workflow/agentworkflow/message_merger.go index df7bb5ce..29fe5f1c 100644 --- a/workflow/agentworkflow/message_merger.go +++ b/workflow/agentworkflow/message_merger.go @@ -88,6 +88,37 @@ func (m *messageMerger) ComputeMerged(primaryResponseID string, primaryAgentID s return response } +type collectedResponseMergeState struct { + allUpdates *messageMerger + terminalWorkflowOutput *messageMerger + hasTerminalOutput bool +} + +func newCollectedResponseMergeState() *collectedResponseMergeState { + return &collectedResponseMergeState{ + allUpdates: newMessageMerger(), + terminalWorkflowOutput: newMessageMerger(), + } +} + +func (s *collectedResponseMergeState) AddUpdate(update *agent.ResponseUpdate, terminalWorkflowOutput bool) { + if s == nil { + return + } + s.allUpdates.AddUpdate(update) + if terminalWorkflowOutput { + s.terminalWorkflowOutput.AddUpdate(update) + s.hasTerminalOutput = true + } +} + +func (s *collectedResponseMergeState) ComputeMerged(primaryResponseID string, primaryAgentID string, primaryAgentName string) *agent.Response { + if s != nil && s.hasTerminalOutput { + return s.terminalWorkflowOutput.ComputeMerged(primaryResponseID, primaryAgentID, primaryAgentName) + } + return s.allUpdates.ComputeMerged(primaryResponseID, primaryAgentID, primaryAgentName) +} + type responseMergeState struct { responseID string updatesByMessageID map[string][]*agent.ResponseUpdate diff --git a/workflow/agentworkflow/workflow.go b/workflow/agentworkflow/workflow.go index c55b53db..8d9f0e8e 100644 --- a/workflow/agentworkflow/workflow.go +++ b/workflow/agentworkflow/workflow.go @@ -91,17 +91,37 @@ func NewAgent(wf *workflow.Workflow, cfg AgentConfig) (*agent.Agent, error) { env = inproc.OffThread } } + terminalOutputExecutors := make(map[string]struct{}) + for executorID, tags := range wf.OutputExecutors() { + if !slices.Contains(tags, workflow.OutputTagIntermediate) { + terminalOutputExecutors[executorID] = struct{}{} + } + } + isTerminalWorkflowOutput := func(evt workflow.OutputEvent) bool { + if evt.IsIntermediate() { + return false + } + if _, ok := terminalOutputExecutors[evt.ExecutorID]; !ok { + return false + } + switch evt.Output.(type) { + case *agent.ResponseUpdate, *agent.Response: + return false + default: + return true + } + } runFn := func(ctx context.Context, messages []*message.Message, options ...agent.Option) iter.Seq2[*agent.ResponseUpdate, error] { return func(yield func(*agent.ResponseUpdate, error) bool) { responseID := uuid.NewString() stream, _ := agent.GetOption(options, agent.Stream) - merger := newMessageMerger() - emitUpdate := func(update *agent.ResponseUpdate) bool { + mergeState := newCollectedResponseMergeState() + emitUpdate := func(update *agent.ResponseUpdate, terminalWorkflowOutput bool) bool { if update == nil { return true } - merger.AddUpdate(update) + mergeState.AddUpdate(update, terminalWorkflowOutput) if !stream { return true } @@ -171,13 +191,13 @@ func NewAgent(wf *workflow.Workflow, cfg AgentConfig) (*agent.Agent, error) { if e.CompletionInfo != nil && e.CompletionInfo.CheckpointInfo != nil { state.lastCheckpoint = e.CompletionInfo.CheckpointInfo } - if !emitUpdate(newUpdate(responseID, e)) { + if !emitUpdate(newUpdate(responseID, e), false) { return } case workflow.OutputEvent: switch out := e.Output.(type) { case *agent.ResponseUpdate: - if !emitUpdate(stampUpdate(out, responseID, e)) { + if !emitUpdate(stampUpdate(out, responseID, e), false) { return } case *agent.Response: @@ -188,7 +208,7 @@ func NewAgent(wf *workflow.Workflow, cfg AgentConfig) (*agent.Agent, error) { continue } for _, msg := range out.Messages { - if !emitUpdate(messageToUpdate(msg, responseID, e)) { + if !emitUpdate(messageToUpdate(msg, responseID, e), false) { return } } @@ -196,15 +216,16 @@ func NewAgent(wf *workflow.Workflow, cfg AgentConfig) (*agent.Agent, error) { if !cfg.IncludeOutputsInResponse { continue } - if !emitUpdate(messageToUpdate(out, responseID, e)) { + if !emitUpdate(messageToUpdate(out, responseID, e), isTerminalWorkflowOutput(e)) { return } case []*message.Message: if !cfg.IncludeOutputsInResponse { continue } + terminalWorkflowOutput := isTerminalWorkflowOutput(e) for _, msg := range out { - if !emitUpdate(messageToUpdate(msg, responseID, e)) { + if !emitUpdate(messageToUpdate(msg, responseID, e), terminalWorkflowOutput) { return } } @@ -213,39 +234,39 @@ func NewAgent(wf *workflow.Workflow, cfg AgentConfig) (*agent.Agent, error) { update, contentID, pending, ok, err := requestToUpdate(e.Request, responseID, e) if err != nil { update := newUpdate(responseID, e, workflowErrorContent(err, cfg.IncludeErrorDetails)) - if !emitUpdate(update) { + if !emitUpdate(update, false) { return } continue } if !ok { - if !emitUpdate(newUpdate(responseID, e)) { + if !emitUpdate(newUpdate(responseID, e), false) { return } continue } state.pending[contentID] = pending - if !emitUpdate(update) { + if !emitUpdate(update, false) { return } case workflow.ErrorEvent: update := newUpdate(responseID, e, workflowErrorContent(e.Error, cfg.IncludeErrorDetails)) - if !emitUpdate(update) { + if !emitUpdate(update, false) { return } case workflow.ExecutorFailedEvent: update := newUpdate(responseID, e, workflowErrorContent(e.Error, cfg.IncludeErrorDetails)) - if !emitUpdate(update) { + if !emitUpdate(update, false) { return } default: - if !emitUpdate(newUpdate(responseID, e)) { + if !emitUpdate(newUpdate(responseID, e), false) { return } } } if !stream { - for _, update := range merger.ComputeMerged(responseID, agentID, agentName).ToUpdates() { + for _, update := range mergeState.ComputeMerged(responseID, agentID, agentName).ToUpdates() { if !yield(update, nil) { return } diff --git a/workflow/agentworkflow/workflow_test.go b/workflow/agentworkflow/workflow_test.go index 8fcc608f..13e253c5 100644 --- a/workflow/agentworkflow/workflow_test.go +++ b/workflow/agentworkflow/workflow_test.go @@ -81,6 +81,65 @@ func echoExecutorBinding(id string) workflow.ExecutorBinding { return binding } +func fixedTextAgent(id, name, text string) *agent.Agent { + run := func(context.Context, []*message.Message, ...agent.Option) iter.Seq2[*agent.ResponseUpdate, error] { + return func(yield func(*agent.ResponseUpdate, error) bool) { + yield(&agent.ResponseUpdate{ + Role: message.RoleAssistant, + AgentID: id, + AuthorName: name, + MessageID: id + "-message", + Contents: []message.Content{ + &message.TextContent{Text: text}, + }, + }, nil) + } + } + return agent.New( + agent.ProviderConfig{ProviderName: "fixed-text", Run: run}, + agent.Config{ + ID: id, + Name: name, + DisableFuncAutoCall: true, + }, + ) +} + +func uppercaseLatestTextBinding(id string) workflow.ExecutorBinding { + binding := workflow.ExecutorBinding{ + ID: id, + ImplementationID: "*workflow.Executor", + RawValue: struct{}{}, + } + binding.NewExecutorFunc = func(_ string) (*workflow.Executor, error) { + return newMessageExecutor(id, &messageworkflow.Options{ + StateKey: id + "_msgs", + TakeTurnHandler: func(ctx *workflow.Context, _ workflow.TurnToken, messages []*message.Message) error { + var latest string + for _, current := range messages { + if current == nil { + continue + } + if text := strings.TrimSpace(current.Contents.Text()); text != "" { + latest = text + } + } + if latest == "" { + return nil + } + return ctx.YieldOutput(&message.Message{ + Role: message.RoleAssistant, + AuthorName: id, + Contents: []message.Content{ + &message.TextContent{Text: strings.ToUpper(latest)}, + }, + }) + }, + }), nil + } + return binding +} + func TestNew_StreamsResponseUpdates(t *testing.T) { binding := echoExecutorBinding("echo") wf, err := workflow.NewBuilder(binding). @@ -297,10 +356,10 @@ func TestNew_IncludeOutputsInResponse(t *testing.T) { } var updates int - for range ag.RunText(t.Context(), "ping") { + for range ag.RunText(t.Context(), "ping", agent.Stream(true)) { updates++ } - // One update from *agent.ResponseUpdate output, one from message output translation. + // Streaming keeps both the hosted update and the translated workflow output visible. if updates < 2 { t.Errorf("expected at least 2 updates with IncludeOutputsInResponse, got %d", updates) } @@ -401,6 +460,50 @@ func TestNew_GatesHostedAgentResponseOutputsByDefault(t *testing.T) { } } +func TestNew_CollectPrefersTerminalWorkflowOutputOverIntermediateHostedAgentUpdates(t *testing.T) { + first := agentworkflow.New( + fixedTextAgent("first-agent", "First Agent", "first answer"), + agentworkflow.Config{ + DisableForwardIncomingMessages: true, + EmitUpdateEvents: true, + }, + ) + second := agentworkflow.New( + fixedTextAgent("second-agent", "Second Agent", "second answer"), + agentworkflow.Config{ + DisableForwardIncomingMessages: true, + EmitUpdateEvents: true, + }, + ) + uppercase := uppercaseLatestTextBinding("uppercase") + + wf, err := workflow.NewBuilder(first). + AddEdge(first, second). + AddEdge(second, uppercase). + WithIntermediateOutputFrom(first, second). + WithOutputFrom(uppercase). + Build() + if err != nil { + t.Fatalf("build workflow: %v", err) + } + + ag, err := agentworkflow.NewAgent(wf, agentworkflow.AgentConfig{ + IncludeOutputsInResponse: true, + }) + if err != nil { + t.Fatalf("New: %v", err) + } + + resp, err := ag.RunText(t.Context(), "hello").Collect() + if err != nil { + t.Fatalf("RunText: %v", err) + } + + if got, want := responseTexts(resp), []string{"SECOND ANSWER"}; !reflect.DeepEqual(got, want) { + t.Fatalf("response texts = %v, want %v; response = %+v", got, want, resp) + } +} + func TestNew_ConvertsResponseOutputAsMessageUpdates(t *testing.T) { createdAt := time.Date(2024, 11, 10, 9, 20, 0, 0, time.UTC) binding := workflow.ExecutorBinding{