Feat/wdz/tool arg handler (#257)

feat: add ToolArgumentsHandler to tool node

Author: meguminnnnnnnnn

compose/tool_node.go

@@ -64,9 +64,10 @@ func withExecutedTools(executedTools map[string]string) ToolsNodeOption {
 //	Invoke(ctx context.Context, input *schema.Message, opts ...ToolsNodeOption) ([]*schema.Message, error)
 //	Stream(ctx context.Context, input *schema.Message, opts ...ToolsNodeOption) (*schema.StreamReader[[]*schema.Message], error)
 type ToolsNode struct {
-	tuple               *toolsTuple
-	unknownToolHandler  func(ctx context.Context, name, input string) (string, error)
-	executeSequentially bool
+	tuple                *toolsTuple
+	unknownToolHandler   func(ctx context.Context, name, input string) (string, error)
+	executeSequentially  bool
+	toolArgumentsHandler func(ctx context.Context, name, input string) (string, error)
 }
 
 // ToolsNodeConfig is the config for ToolsNode.
@@ -91,6 +92,17 @@ type ToolsNodeConfig struct {
 	// When set to true, tool calls will be executed one after another in the order they appear in the input message.
 	// When set to false (default), tool calls will be executed in parallel.
 	ExecuteSequentially bool
+
+	// ToolArgumentsHandler allows handling of tool arguments before execution.
+	// When provided, this function will be called for each tool call to process the arguments.
+	// Parameters:
+	//   - ctx: The context for the tool call
+	//   - name: The name of the tool being called
+	//   - arguments: The original arguments string for the tool
+	// Returns:
+	//   - string: The processed arguments string to be used for tool execution
+	//   - error: Any error that occurred during preprocessing
+	ToolArgumentsHandler func(ctx context.Context, name, arguments string) (string, error)
 }
 
 // NewToolNode creates a new ToolsNode.
@@ -107,9 +119,10 @@ func NewToolNode(ctx context.Context, conf *ToolsNodeConfig) (*ToolsNode, error)
 	}
 
 	return &ToolsNode{
-		tuple:               tuple,
-		unknownToolHandler:  conf.UnknownToolsHandler,
-		executeSequentially: conf.ExecuteSequentially,
+		tuple:                tuple,
+		unknownToolHandler:   conf.UnknownToolsHandler,
+		executeSequentially:  conf.ExecuteSequentially,
+		toolArgumentsHandler: conf.ToolArgumentsHandler,
 	}, nil
 }
 
@@ -191,7 +204,7 @@ type toolCallTask struct {
 	err      error
 }
 
-func (tn *ToolsNode) genToolCallTasks(tuple *toolsTuple, input *schema.Message, executedTools map[string]string) ([]toolCallTask, error) {
+func (tn *ToolsNode) genToolCallTasks(ctx context.Context, tuple *toolsTuple, input *schema.Message, executedTools map[string]string, isStream bool) ([]toolCallTask, error) {
 	if input.Role != schema.Assistant {
 		return nil, fmt.Errorf("expected message role is Assistant, got %s", input.Role)
 	}
@@ -210,7 +223,11 @@ func (tn *ToolsNode) genToolCallTasks(tuple *toolsTuple, input *schema.Message,
 			toolCallTasks[i].arg = toolCall.Function.Arguments
 			toolCallTasks[i].callID = toolCall.ID
 			toolCallTasks[i].executed = true
-			toolCallTasks[i].output = result
+			if isStream {
+				toolCallTasks[i].sOutput = schema.StreamReaderFromArray([]string{result})
+			} else {
+				toolCallTasks[i].output = result
+			}
 			continue
 		}
 		index, ok := tuple.indexes[toolCall.Function.Name]
@@ -223,8 +240,16 @@ func (tn *ToolsNode) genToolCallTasks(tuple *toolsTuple, input *schema.Message,
 			toolCallTasks[i].r = tuple.rps[index]
 			toolCallTasks[i].meta = tuple.meta[index]
 			toolCallTasks[i].name = toolCall.Function.Name
-			toolCallTasks[i].arg = toolCall.Function.Arguments
 			toolCallTasks[i].callID = toolCall.ID
+			if tn.toolArgumentsHandler != nil {
+				arg, err := tn.toolArgumentsHandler(ctx, toolCall.Function.Name, toolCall.Function.Arguments)
+				if err != nil {
+					return nil, fmt.Errorf("failed to executed tool[name:%s arguments:%s] arguments handler: %w", toolCall.Function.Name, toolCall.Function.Arguments, err)
+				}
+				toolCallTasks[i].arg = arg
+			} else {
+				toolCallTasks[i].arg = toolCall.Function.Arguments
+			}
 		}
 	}
 
@@ -280,6 +305,9 @@ func runToolCallTaskByStream(ctx context.Context, task *toolCallTask, opts ...to
 
 func sequentialRunToolCall(ctx context.Context, run func(ctx2 context.Context, callTask *toolCallTask, opts ...tool.Option), tasks []toolCallTask, opts ...tool.Option) {
 	for i := 0; i < len(tasks); i++ {
+		if tasks[i].executed {
+			continue
+		}
 		run(ctx, &tasks[i], opts...)
 	}
 }
@@ -294,6 +322,9 @@ func parallelRunToolCall(ctx context.Context,
 
 	var wg sync.WaitGroup
 	for i := 1; i < len(tasks); i++ {
+		if tasks[i].executed {
+			continue
+		}
 		wg.Add(1)
 		go func(ctx_ context.Context, t *toolCallTask, opts ...tool.Option) {
 			defer wg.Done()
@@ -326,7 +357,7 @@ func (tn *ToolsNode) Invoke(ctx context.Context, input *schema.Message,
 		}
 	}
 
-	tasks, err := tn.genToolCallTasks(tuple, input, opt.executedTools)
+	tasks, err := tn.genToolCallTasks(ctx, tuple, input, opt.executedTools, false)
 	if err != nil {
 		return nil, err
 	}
@@ -385,7 +416,7 @@ func (tn *ToolsNode) Stream(ctx context.Context, input *schema.Message,
 		}
 	}
 
-	tasks, err := tn.genToolCallTasks(tuple, input, opt.executedTools)
+	tasks, err := tn.genToolCallTasks(ctx, tuple, input, opt.executedTools, true)
 	if err != nil {
 		return nil, err
 	}

compose/tool_node_test.go

@@ -734,7 +734,9 @@ func TestToolRerun(t *testing.T) {
 		},
 	}, info.RerunNodesExtra["tool node"])
 
-	result, err := r.Invoke(ctx, nil, WithCheckPointID("1"))
+	sr, err := r.Stream(ctx, nil, WithCheckPointID("1"))
+	assert.NoError(t, err)
+	result, err := concatStreamReader(sr)
 	assert.NoError(t, err)
 	assert.Equal(t, "tool1 input: inputtool2 input: inputtool3 input: inputtool4 input: input", result)
 }