Add concurrent tool calls

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 {

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
 }