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)
+		toolResults := c.callTools(ctx, *message)
-		if err != nil {
-			return nil, fmt.Errorf("completions: %w", err)
-		}
 		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)
+		toolResults := c.callTools(ctx, message)
-		if err != nil {
-			return fmt.Errorf("completions stream: %w", err)
-		}
 		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,15 +4,26 @@ import (
 	"context"
 	"fmt"
 	"log/slog"
+	"runtime"
 	"strings"
+	"sync"
 )
 
-// callTools executes each tool call in the message and returns the
+// callTools executes each tool call in the message concurrently and returns
-// resulting tool result messages.
+// the resulting tool result messages in the same order as the input calls.
-func (c *Client) callTools(ctx context.Context, msg Message) (
+func (c *Client) callTools(ctx context.Context, msg Message) []Message {
-	results []Message, err error,
+	var (
-) {
+		results = make([]Message, len(msg.ToolCalls))
-	for _, tc := range 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 }()
+
 			c.log.InfoContext(
 				ctx,
 				"calling tool",
@@ -46,7 +57,7 @@ func (c *Client) callTools(ctx context.Context, msg Message) (
 				result = `{"ok": true, "result": null}`
 			}
 
-		toolResult := Message{
+			results[idx] = Message{
 				Role: RoleTool,
 				ContentParts: []ContentPart{{
 					Type: ContentTypeText, Text: result,
@@ -54,8 +65,9 @@ func (c *Client) callTools(ctx context.Context, msg Message) (
 				Name:       tc.Function.Name,
 				ToolCallID: tc.ID,
 			}
-		results = append(results, toolResult)
+		}(i, tc)
 	}
 
-	return results, nil
+	wg.Wait()
+	return results
 }