diff --git a/internal/llm/client_completions.go b/internal/llm/client_completions.go index cc29ec0..9763ed3 100644 --- a/internal/llm/client_completions.go +++ b/internal/llm/client_completions.go @@ -97,10 +97,7 @@ func (c *Client) Completions( Timings: finalTimings, }, nil } - toolResults, err := c.callTools(ctx, *message) - if err != nil { - return nil, fmt.Errorf("completions: %w", err) - } + toolResults := c.callTools(ctx, *message) messages = append(messages, toolResults...) } @@ -126,6 +123,7 @@ func (c *Client) completions( Tools: c.tools, ThinkingBudgetTokens: thinkingBudgetTokens, ChatTemplateKWARGS: kwargs, + ParallelToolCalls: new(true), }) if err != nil { return nil, fmt.Errorf( @@ -286,10 +284,7 @@ func (c *Client) CompletionsStream( if len(message.ToolCalls) == 0 { return nil } - toolResults, err := c.callTools(ctx, message) - if err != nil { - return fmt.Errorf("completions stream: %w", err) - } + toolResults := c.callTools(ctx, message) for _, tr := range toolResults { events <- StreamEvent{ Type: streamEventMessage, @@ -328,6 +323,7 @@ func (c *Client) completionsStream(ctx context.Context, messages []Message, mode Tools: c.tools, ThinkingBudgetTokens: thinkingBudgetTokens, ChatTemplateKWARGS: kwargs, + ParallelToolCalls: new(true), StreamOptions: &StreamOptions{IncludeUsage: true}, }) if err != nil { diff --git a/internal/llm/client_tools.go b/internal/llm/client_tools.go index f246cd3..2f5ac15 100644 --- a/internal/llm/client_tools.go +++ b/internal/llm/client_tools.go @@ -4,58 +4,70 @@ import ( "context" "fmt" "log/slog" + "runtime" "strings" + "sync" ) -// callTools executes each tool call in the message and returns the -// resulting tool result messages. -func (c *Client) callTools(ctx context.Context, msg Message) ( - results []Message, err error, -) { - for _, tc := range msg.ToolCalls { - c.log.InfoContext( - ctx, - "calling tool", - slog.String("name", tc.Function.Name), - slog.String("args", tc.Function.Arguments), - ) +// callTools executes each tool call in the message concurrently and returns +// the resulting tool result messages in the same order as the input calls. +func (c *Client) callTools(ctx context.Context, msg Message) []Message { + var ( + results = make([]Message, len(msg.ToolCalls)) + sem = make(chan struct{}, runtime.NumCPU()*2) + wg sync.WaitGroup + ) + for i, tc := range msg.ToolCalls { + wg.Add(1) + go func(idx int, tc ToolCall) { + defer wg.Done() + sem <- struct{}{} + defer func() { <-sem }() - result, err := c.registry.Execute( - ctx, tc.Function.Name, tc.Function.Arguments, - ) - if err != nil { - c.log.ErrorContext( - ctx, - "failed to call tool", - slog.Any("error", err), - slog.String("name", tc.Function.Name), - ) - result = fmt.Sprintf( - `{"ok": false, "error": %q}`, err, - ) - } else { c.log.InfoContext( ctx, - "called tool", + "calling tool", slog.String("name", tc.Function.Name), + slog.String("args", tc.Function.Arguments), ) - } - // Content cannot be empty. - if strings.TrimSpace(result) == "" { - result = `{"ok": true, "result": null}` - } + result, err := c.registry.Execute( + ctx, tc.Function.Name, tc.Function.Arguments, + ) + if err != nil { + c.log.ErrorContext( + ctx, + "failed to call tool", + slog.Any("error", err), + slog.String("name", tc.Function.Name), + ) + result = fmt.Sprintf( + `{"ok": false, "error": %q}`, err, + ) + } else { + c.log.InfoContext( + ctx, + "called tool", + slog.String("name", tc.Function.Name), + ) + } - toolResult := Message{ - Role: RoleTool, - ContentParts: []ContentPart{{ - Type: ContentTypeText, Text: result, - }}, - Name: tc.Function.Name, - ToolCallID: tc.ID, - } - results = append(results, toolResult) + // Content cannot be empty. + if strings.TrimSpace(result) == "" { + result = `{"ok": true, "result": null}` + } + + results[idx] = Message{ + Role: RoleTool, + ContentParts: []ContentPart{{ + Type: ContentTypeText, Text: result, + }}, + Name: tc.Function.Name, + ToolCallID: tc.ID, + } + }(i, tc) } - return results, nil + wg.Wait() + return results }