diff --git a/internal/multiagent/hitl_middleware.go b/internal/multiagent/hitl_middleware.go index b4cf23a3..c7963a1d 100644 --- a/internal/multiagent/hitl_middleware.go +++ b/internal/multiagent/hitl_middleware.go @@ -3,6 +3,7 @@ package multiagent import ( "context" "errors" + "io" "strings" "github.com/cloudwego/eino/adk" @@ -65,11 +66,51 @@ func hitlClearReturnDirectlyIfTransfer(ctx context.Context, toolName string) { }) } +func hitlEditedArgumentsNotice(original, edited string) string { + original = strings.TrimSpace(original) + edited = strings.TrimSpace(edited) + if edited == "" || edited == original { + return "" + } + return "[HITL] Human reviewer approved this tool call with edited arguments.\n" + + "Original arguments: " + original + "\n" + + "Executed arguments: " + edited + "\n\n" +} + +func hitlPrependEditedArgumentsNotice(result, original, edited string) string { + notice := hitlEditedArgumentsNotice(original, edited) + if notice == "" { + return result + } + return notice + result +} + +func hitlCollectStringStream(sr *schema.StreamReader[string]) (string, error) { + if sr == nil { + return "", nil + } + defer sr.Close() + var b strings.Builder + for { + chunk, err := sr.Recv() + if errors.Is(err, io.EOF) { + return b.String(), nil + } + if err != nil { + return b.String(), err + } + b.WriteString(chunk) + } +} + func hitlInvokableToolCallMiddleware() compose.InvokableToolMiddleware { return func(next compose.InvokableToolEndpoint) compose.InvokableToolEndpoint { return func(ctx context.Context, input *compose.ToolInput) (*compose.ToolOutput, error) { + originalArgs := "" + editedArgs := "" if input != nil { if fn, ok := ctx.Value(hitlInterceptorKey{}).(HITLToolInterceptor); ok && fn != nil { + originalArgs = input.Arguments edited, err := fn(ctx, input.Name, input.Arguments) if err != nil { if IsHumanRejectError(err) { @@ -85,11 +126,17 @@ func hitlInvokableToolCallMiddleware() compose.InvokableToolMiddleware { return nil, err } if edited != "" { + editedArgs = edited input.Arguments = edited } } } - return next(ctx, input) + out, err := next(ctx, input) + if err != nil || out == nil { + return out, err + } + out.Result = hitlPrependEditedArgumentsNotice(out.Result, originalArgs, editedArgs) + return out, nil } } } @@ -97,8 +144,11 @@ func hitlInvokableToolCallMiddleware() compose.InvokableToolMiddleware { func hitlStreamableToolCallMiddleware() compose.StreamableToolMiddleware { return func(next compose.StreamableToolEndpoint) compose.StreamableToolEndpoint { return func(ctx context.Context, input *compose.ToolInput) (*compose.StreamToolOutput, error) { + originalArgs := "" + editedArgs := "" if input != nil { if fn, ok := ctx.Value(hitlInterceptorKey{}).(HITLToolInterceptor); ok && fn != nil { + originalArgs = input.Arguments edited, err := fn(ctx, input.Name, input.Arguments) if err != nil { if IsHumanRejectError(err) { @@ -111,11 +161,26 @@ func hitlStreamableToolCallMiddleware() compose.StreamableToolMiddleware { return nil, err } if edited != "" { + editedArgs = edited input.Arguments = edited } } } - return next(ctx, input) + out, err := next(ctx, input) + if err != nil || out == nil { + return out, err + } + if hitlEditedArgumentsNotice(originalArgs, editedArgs) == "" { + return out, nil + } + result, collectErr := hitlCollectStringStream(out.Result) + if collectErr != nil { + return nil, collectErr + } + out.Result = schema.StreamReaderFromArray([]string{ + hitlPrependEditedArgumentsNotice(result, originalArgs, editedArgs), + }) + return out, nil } } }