Merge pull request #2645 from loafoe/feat/bedrock-streaming

feat(bedrock): implement StreamingProvider for real-time token streaming

pkg/providers/bedrock/provider_bedrock.go

@@ -135,48 +135,23 @@ func NewProvider(ctx context.Context, opts ...Option) (*Provider, error) {
 	}, nil
 }
 
-// Chat sends messages to AWS Bedrock using the Converse API.
+// converseParams holds the shared request parameters for Converse and ConverseStream.
-func (p *Provider) Chat(
+type converseParams struct {
-	ctx context.Context,
+	messages       []types.Message
-	messages []Message,
+	system         []types.SystemContentBlock
-	tools []ToolDefinition,
+	inferenceConfig *types.InferenceConfiguration
-	model string,
+	toolConfig     *types.ToolConfiguration
-	options map[string]any,
+}
-) (*LLMResponse, error) {
-	// Apply request timeout if context doesn't already have a deadline.
-	// Use explicit timeout if set, otherwise fall back to common default.
-	effectiveTimeout := p.requestTimeout
-	if effectiveTimeout <= 0 {
-		effectiveTimeout = common.DefaultRequestTimeout
-	}
-	if _, hasDeadline := ctx.Deadline(); !hasDeadline {
-		var cancel context.CancelFunc
-		ctx, cancel = context.WithTimeout(ctx, effectiveTimeout)
-		defer cancel()
-	}
 
-	// Build the Converse API input
+func buildConverseParams(messages []Message, tools []ToolDefinition, options map[string]any) converseParams {
-	input := &bedrockruntime.ConverseInput{
-		ModelId: aws.String(model),
-	}
-
-	// Convert messages to Bedrock format
 	bedrockMessages, systemPrompts := convertMessages(messages)
-	input.Messages = bedrockMessages
 
-	// Set system prompts if any
-	if len(systemPrompts) > 0 {
-		input.System = systemPrompts
-	}
-
-	// Set inference configuration only when options are provided
 	var inferenceConfig *types.InferenceConfiguration
 
 	if maxTokens, ok := common.AsInt(options["max_tokens"]); ok && maxTokens > 0 {
 		if inferenceConfig == nil {
 			inferenceConfig = &types.InferenceConfiguration{}
 		}
-		// Clamp to int32 range to avoid overflow
 		if maxTokens > math.MaxInt32 {
 			maxTokens = math.MaxInt32
 		}
@@ -190,23 +165,53 @@ func (p *Provider) Chat(
 		inferenceConfig.Temperature = aws.Float32(float32(temp))
 	}
 
-	if inferenceConfig != nil {
+	var toolConfig *types.ToolConfiguration
-		input.InferenceConfig = inferenceConfig
-	}
-
-	// Convert tools to Bedrock format
-	// Only set ToolConfig if at least one valid tool was produced
 	if len(tools) > 0 {
-		toolConfig := convertTools(tools)
+		tc := convertTools(tools)
-		if len(toolConfig.Tools) > 0 {
+		if len(tc.Tools) > 0 {
-			input.ToolConfig = toolConfig
+			toolConfig = tc
 		}
 	}
 
-	// Call Bedrock Converse API
+	return converseParams{
+		messages:       bedrockMessages,
+		system:         systemPrompts,
+		inferenceConfig: inferenceConfig,
+		toolConfig:     toolConfig,
+	}
+}
+
+// Chat sends messages to AWS Bedrock using the Converse API.
+func (p *Provider) Chat(
+	ctx context.Context,
+	messages []Message,
+	tools []ToolDefinition,
+	model string,
+	options map[string]any,
+) (*LLMResponse, error) {
+	effectiveTimeout := p.requestTimeout
+	if effectiveTimeout <= 0 {
+		effectiveTimeout = common.DefaultRequestTimeout
+	}
+	if _, hasDeadline := ctx.Deadline(); !hasDeadline {
+		var cancel context.CancelFunc
+		ctx, cancel = context.WithTimeout(ctx, effectiveTimeout)
+		defer cancel()
+	}
+
+	params := buildConverseParams(messages, tools, options)
+	input := &bedrockruntime.ConverseInput{
+		ModelId:         aws.String(model),
+		Messages:        params.messages,
+		InferenceConfig: params.inferenceConfig,
+		ToolConfig:      params.toolConfig,
+	}
+	if len(params.system) > 0 {
+		input.System = params.system
+	}
+
 	output, err := p.client.Converse(ctx, input)
 	if err != nil {
-		// Check for SSO token expiration errors and provide actionable guidance
 		if isSSOTokenError(err) {
 			return nil, fmt.Errorf(
 				"bedrock converse: AWS credentials may have expired. If using AWS SSO, run 'aws sso login' to refresh: %w",
@@ -216,10 +221,199 @@ func (p *Provider) Chat(
 		return nil, fmt.Errorf("bedrock converse: %w", err)
 	}
 
-	// Parse the response
 	return parseResponse(output)
 }
 
+// ChatStream sends messages to AWS Bedrock using the ConverseStream API.
+// It streams the accumulated text so far via the onChunk callback and returns the complete response.
+func (p *Provider) ChatStream(
+	ctx context.Context,
+	messages []Message,
+	tools []ToolDefinition,
+	model string,
+	options map[string]any,
+	onChunk func(accumulated string),
+) (*LLMResponse, error) {
+	if p.requestTimeout > 0 {
+		if _, hasDeadline := ctx.Deadline(); !hasDeadline {
+			var cancel context.CancelFunc
+			ctx, cancel = context.WithTimeout(ctx, p.requestTimeout)
+			defer cancel()
+		}
+	}
+
+	params := buildConverseParams(messages, tools, options)
+	input := &bedrockruntime.ConverseStreamInput{
+		ModelId:         aws.String(model),
+		Messages:        params.messages,
+		InferenceConfig: params.inferenceConfig,
+		ToolConfig:      params.toolConfig,
+	}
+	if len(params.system) > 0 {
+		input.System = params.system
+	}
+
+	output, err := p.client.ConverseStream(ctx, input)
+	if err != nil {
+		if isSSOTokenError(err) {
+			return nil, fmt.Errorf(
+				"bedrock conversestream: AWS credentials may have expired. If using AWS SSO, run 'aws sso login' to refresh: %w",
+				err,
+			)
+		}
+		return nil, fmt.Errorf("bedrock conversestream: %w", err)
+	}
+
+	return parseStreamResponse(ctx, output.GetStream(), onChunk)
+}
+
+// converseStreamReader abstracts the Bedrock event stream so parseStreamResponse
+// can be unit-tested with a mock event source.
+type converseStreamReader interface {
+	Events() <-chan types.ConverseStreamOutput
+	Err() error
+	Close() error
+}
+
+// parseStreamResponse processes the ConverseStream event stream and accumulates the response.
+func parseStreamResponse(
+	ctx context.Context,
+	stream converseStreamReader,
+	onChunk func(accumulated string),
+) (resp *LLMResponse, err error) {
+	if stream == nil {
+		return nil, fmt.Errorf("bedrock conversestream: nil event stream")
+	}
+	defer func() {
+		if closeErr := stream.Close(); closeErr != nil {
+			if err == nil {
+				err = fmt.Errorf("bedrock conversestream: close event stream: %w", closeErr)
+			} else {
+				log.Printf("bedrock conversestream: close event stream: %v", closeErr)
+			}
+		}
+	}()
+
+	var textContent strings.Builder
+	finishReason := "stop"
+	var usage *UsageInfo
+	toolCalls := make([]ToolCall, 0)
+
+	// Track active tool use blocks by index
+	type toolAccum struct {
+		id       string
+		name     string
+		argsJSON strings.Builder
+	}
+	activeTools := map[int]*toolAccum{}
+
+	events := stream.Events()
+	for {
+		select {
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		case event, ok := <-events:
+			if !ok {
+				// Stream closed
+				goto done
+			}
+
+			switch e := event.(type) {
+			case *types.ConverseStreamOutputMemberContentBlockStart:
+				// New content block starting
+				if toolUse, ok := e.Value.Start.(*types.ContentBlockStartMemberToolUse); ok {
+					activeTools[int(aws.ToInt32(e.Value.ContentBlockIndex))] = &toolAccum{
+						id:   aws.ToString(toolUse.Value.ToolUseId),
+						name: aws.ToString(toolUse.Value.Name),
+					}
+				}
+
+			case *types.ConverseStreamOutputMemberContentBlockDelta:
+				// Content delta
+				switch delta := e.Value.Delta.(type) {
+				case *types.ContentBlockDeltaMemberText:
+					textContent.WriteString(delta.Value)
+					if onChunk != nil {
+						onChunk(textContent.String())
+					}
+				case *types.ContentBlockDeltaMemberToolUse:
+					idx := int(aws.ToInt32(e.Value.ContentBlockIndex))
+					if tool, exists := activeTools[idx]; exists {
+						tool.argsJSON.WriteString(aws.ToString(delta.Value.Input))
+					}
+				}
+
+			case *types.ConverseStreamOutputMemberContentBlockStop:
+				// Content block finished - finalize tool if it was a tool use
+				idx := int(aws.ToInt32(e.Value.ContentBlockIndex))
+				if tool, exists := activeTools[idx]; exists {
+					args := make(map[string]any)
+					argsStr := tool.argsJSON.String()
+					if argsStr != "" {
+						if err := json.Unmarshal([]byte(argsStr), &args); err != nil {
+							log.Printf("bedrock: stream: failed to parse tool arguments for %q: %v", tool.name, err)
+							args = map[string]any{"raw": argsStr}
+						}
+					}
+					funcArgs := argsStr
+					if argsJSON, marshalErr := json.Marshal(args); marshalErr == nil {
+						funcArgs = string(argsJSON)
+					}
+					toolCalls = append(toolCalls, ToolCall{
+						ID:        tool.id,
+						Name:      tool.name,
+						Arguments: args,
+						Function: &FunctionCall{
+							Name:      tool.name,
+							Arguments: funcArgs,
+						},
+					})
+					delete(activeTools, idx)
+				}
+
+			case *types.ConverseStreamOutputMemberMessageStop:
+				// Message complete
+				switch e.Value.StopReason {
+				case types.StopReasonToolUse:
+					finishReason = "tool_calls"
+				case types.StopReasonMaxTokens:
+					finishReason = "length"
+				case types.StopReasonEndTurn:
+					finishReason = "stop"
+				case types.StopReasonStopSequence:
+					finishReason = "stop"
+				case types.StopReasonContentFiltered:
+					finishReason = "content_filter"
+				default:
+					finishReason = "stop"
+				}
+
+			case *types.ConverseStreamOutputMemberMetadata:
+				// Usage metadata
+				if e.Value.Usage != nil {
+					usage = &UsageInfo{
+						PromptTokens:     int(aws.ToInt32(e.Value.Usage.InputTokens)),
+						CompletionTokens: int(aws.ToInt32(e.Value.Usage.OutputTokens)),
+						TotalTokens:      int(aws.ToInt32(e.Value.Usage.InputTokens)) + int(aws.ToInt32(e.Value.Usage.OutputTokens)),
+					}
+				}
+			}
+		}
+	}
+
+done:
+	if err := stream.Err(); err != nil {
+		return nil, fmt.Errorf("bedrock conversestream: %w", err)
+	}
+
+	return &LLMResponse{
+		Content:      textContent.String(),
+		ToolCalls:    toolCalls,
+		FinishReason: finishReason,
+		Usage:        usage,
+	}, nil
+}
+
 // GetDefaultModel returns an empty string as Bedrock models are user-configured.
 func (p *Provider) GetDefaultModel() string {
 	return ""

pkg/providers/bedrock/provider_bedrock_test.go

@@ -8,6 +8,7 @@
 package bedrock
 
 import (
+	"context"
 	"fmt"
 	"testing"
 
@@ -605,3 +606,272 @@ func TestIsSSOTokenError(t *testing.T) {
 		})
 	}
 }
+
+// mockStreamReader implements bedrockruntime.ConverseStreamOutputReader for testing.
+type mockStreamReader struct {
+	ch  chan types.ConverseStreamOutput
+	err error
+}
+
+func (r *mockStreamReader) Events() <-chan types.ConverseStreamOutput { return r.ch }
+func (r *mockStreamReader) Close() error                              { return nil }
+func (r *mockStreamReader) Err() error                                { return r.err }
+
+func newMockStream(events []types.ConverseStreamOutput) *bedrockruntime.ConverseStreamEventStream {
+	ch := make(chan types.ConverseStreamOutput, len(events))
+	for _, e := range events {
+		ch <- e
+	}
+	close(ch)
+
+	return bedrockruntime.NewConverseStreamEventStream(func(es *bedrockruntime.ConverseStreamEventStream) {
+		es.Reader = &mockStreamReader{ch: ch}
+	})
+}
+
+func TestParseStreamResponse_TextOnly(t *testing.T) {
+	events := []types.ConverseStreamOutput{
+		&types.ConverseStreamOutputMemberContentBlockDelta{
+			Value: types.ContentBlockDeltaEvent{
+				Delta:             &types.ContentBlockDeltaMemberText{Value: "Hello "},
+				ContentBlockIndex: aws.Int32(0),
+			},
+		},
+		&types.ConverseStreamOutputMemberContentBlockDelta{
+			Value: types.ContentBlockDeltaEvent{
+				Delta:             &types.ContentBlockDeltaMemberText{Value: "World"},
+				ContentBlockIndex: aws.Int32(0),
+			},
+		},
+		&types.ConverseStreamOutputMemberMessageStop{
+			Value: types.MessageStopEvent{StopReason: types.StopReasonEndTurn},
+		},
+		&types.ConverseStreamOutputMemberMetadata{
+			Value: types.ConverseStreamMetadataEvent{
+				Usage: &types.TokenUsage{
+					InputTokens:  aws.Int32(10),
+					OutputTokens: aws.Int32(5),
+				},
+			},
+		},
+	}
+
+	var chunks []string
+	stream := newMockStream(events)
+	resp, err := parseStreamResponse(context.Background(), stream, func(accumulated string) {
+		chunks = append(chunks, accumulated)
+	})
+
+	require.NoError(t, err)
+	assert.Equal(t, "Hello World", resp.Content)
+	assert.Equal(t, "stop", resp.FinishReason)
+	assert.Empty(t, resp.ToolCalls)
+	require.NotNil(t, resp.Usage)
+	assert.Equal(t, 10, resp.Usage.PromptTokens)
+	assert.Equal(t, 5, resp.Usage.CompletionTokens)
+	assert.Equal(t, 15, resp.Usage.TotalTokens)
+	assert.Equal(t, []string{"Hello ", "Hello World"}, chunks)
+}
+
+func TestParseStreamResponse_ToolCall(t *testing.T) {
+	events := []types.ConverseStreamOutput{
+		&types.ConverseStreamOutputMemberContentBlockStart{
+			Value: types.ContentBlockStartEvent{
+				ContentBlockIndex: aws.Int32(0),
+				Start: &types.ContentBlockStartMemberToolUse{
+					Value: types.ToolUseBlockStart{
+						ToolUseId: aws.String("call_1"),
+						Name:      aws.String("search"),
+					},
+				},
+			},
+		},
+		&types.ConverseStreamOutputMemberContentBlockDelta{
+			Value: types.ContentBlockDeltaEvent{
+				ContentBlockIndex: aws.Int32(0),
+				Delta: &types.ContentBlockDeltaMemberToolUse{
+					Value: types.ToolUseBlockDelta{Input: aws.String(`{"q":`)},
+				},
+			},
+		},
+		&types.ConverseStreamOutputMemberContentBlockDelta{
+			Value: types.ContentBlockDeltaEvent{
+				ContentBlockIndex: aws.Int32(0),
+				Delta: &types.ContentBlockDeltaMemberToolUse{
+					Value: types.ToolUseBlockDelta{Input: aws.String(`"test"}`)},
+				},
+			},
+		},
+		&types.ConverseStreamOutputMemberContentBlockStop{
+			Value: types.ContentBlockStopEvent{ContentBlockIndex: aws.Int32(0)},
+		},
+		&types.ConverseStreamOutputMemberMessageStop{
+			Value: types.MessageStopEvent{StopReason: types.StopReasonToolUse},
+		},
+	}
+
+	stream := newMockStream(events)
+	resp, err := parseStreamResponse(context.Background(), stream, nil)
+
+	require.NoError(t, err)
+	assert.Equal(t, "tool_calls", resp.FinishReason)
+	require.Len(t, resp.ToolCalls, 1)
+	assert.Equal(t, "call_1", resp.ToolCalls[0].ID)
+	assert.Equal(t, "search", resp.ToolCalls[0].Name)
+	assert.Equal(t, map[string]any{"q": "test"}, resp.ToolCalls[0].Arguments)
+	require.NotNil(t, resp.ToolCalls[0].Function)
+	assert.Equal(t, "search", resp.ToolCalls[0].Function.Name)
+	assert.Equal(t, `{"q":"test"}`, resp.ToolCalls[0].Function.Arguments)
+}
+
+func TestParseStreamResponse_TextAndToolCall(t *testing.T) {
+	events := []types.ConverseStreamOutput{
+		&types.ConverseStreamOutputMemberContentBlockDelta{
+			Value: types.ContentBlockDeltaEvent{
+				ContentBlockIndex: aws.Int32(0),
+				Delta:             &types.ContentBlockDeltaMemberText{Value: "Let me search that."},
+			},
+		},
+		&types.ConverseStreamOutputMemberContentBlockStart{
+			Value: types.ContentBlockStartEvent{
+				ContentBlockIndex: aws.Int32(1),
+				Start: &types.ContentBlockStartMemberToolUse{
+					Value: types.ToolUseBlockStart{
+						ToolUseId: aws.String("call_2"),
+						Name:      aws.String("web"),
+					},
+				},
+			},
+		},
+		&types.ConverseStreamOutputMemberContentBlockDelta{
+			Value: types.ContentBlockDeltaEvent{
+				ContentBlockIndex: aws.Int32(1),
+				Delta: &types.ContentBlockDeltaMemberToolUse{
+					Value: types.ToolUseBlockDelta{Input: aws.String(`{"url":"https://example.com"}`)},
+				},
+			},
+		},
+		&types.ConverseStreamOutputMemberContentBlockStop{
+			Value: types.ContentBlockStopEvent{ContentBlockIndex: aws.Int32(1)},
+		},
+		&types.ConverseStreamOutputMemberMessageStop{
+			Value: types.MessageStopEvent{StopReason: types.StopReasonToolUse},
+		},
+	}
+
+	var chunks []string
+	stream := newMockStream(events)
+	resp, err := parseStreamResponse(context.Background(), stream, func(accumulated string) {
+		chunks = append(chunks, accumulated)
+	})
+
+	require.NoError(t, err)
+	assert.Equal(t, "Let me search that.", resp.Content)
+	assert.Equal(t, "tool_calls", resp.FinishReason)
+	require.Len(t, resp.ToolCalls, 1)
+	assert.Equal(t, "web", resp.ToolCalls[0].Name)
+	assert.Equal(t, []string{"Let me search that."}, chunks)
+}
+
+func TestParseStreamResponse_ContextCancelled(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	// Use an unbuffered channel with no events so ctx.Done() is the only ready case.
+	ch := make(chan types.ConverseStreamOutput)
+
+	stream := bedrockruntime.NewConverseStreamEventStream(func(es *bedrockruntime.ConverseStreamEventStream) {
+		es.Reader = &mockStreamReader{ch: ch}
+	})
+
+	_, err := parseStreamResponse(ctx, stream, nil)
+	assert.ErrorIs(t, err, context.Canceled)
+}
+
+func TestParseStreamResponse_InvalidToolJSON(t *testing.T) {
+	events := []types.ConverseStreamOutput{
+		&types.ConverseStreamOutputMemberContentBlockStart{
+			Value: types.ContentBlockStartEvent{
+				ContentBlockIndex: aws.Int32(0),
+				Start: &types.ContentBlockStartMemberToolUse{
+					Value: types.ToolUseBlockStart{
+						ToolUseId: aws.String("call_bad"),
+						Name:      aws.String("broken"),
+					},
+				},
+			},
+		},
+		&types.ConverseStreamOutputMemberContentBlockDelta{
+			Value: types.ContentBlockDeltaEvent{
+				ContentBlockIndex: aws.Int32(0),
+				Delta: &types.ContentBlockDeltaMemberToolUse{
+					Value: types.ToolUseBlockDelta{Input: aws.String(`{not valid json`)},
+				},
+			},
+		},
+		&types.ConverseStreamOutputMemberContentBlockStop{
+			Value: types.ContentBlockStopEvent{ContentBlockIndex: aws.Int32(0)},
+		},
+		&types.ConverseStreamOutputMemberMessageStop{
+			Value: types.MessageStopEvent{StopReason: types.StopReasonToolUse},
+		},
+	}
+
+	stream := newMockStream(events)
+	resp, err := parseStreamResponse(context.Background(), stream, nil)
+
+	require.NoError(t, err)
+	require.Len(t, resp.ToolCalls, 1)
+	assert.Equal(t, map[string]any{"raw": `{not valid json`}, resp.ToolCalls[0].Arguments)
+	assert.JSONEq(t, `{"raw":"{not valid json"}`, resp.ToolCalls[0].Function.Arguments)
+}
+
+func TestParseStreamResponse_DefaultFinishReason(t *testing.T) {
+	events := []types.ConverseStreamOutput{
+		&types.ConverseStreamOutputMemberContentBlockDelta{
+			Value: types.ContentBlockDeltaEvent{
+				Delta:             &types.ContentBlockDeltaMemberText{Value: "partial"},
+				ContentBlockIndex: aws.Int32(0),
+			},
+		},
+	}
+
+	stream := newMockStream(events)
+	resp, err := parseStreamResponse(context.Background(), stream, nil)
+
+	require.NoError(t, err)
+	assert.Equal(t, "stop", resp.FinishReason)
+}
+
+func TestParseStreamResponse_NilStream(t *testing.T) {
+	_, err := parseStreamResponse(context.Background(), nil, nil)
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "nil event stream")
+}
+
+func TestParseStreamResponse_StopReasons(t *testing.T) {
+	tests := []struct {
+		reason   types.StopReason
+		expected string
+	}{
+		{types.StopReasonEndTurn, "stop"},
+		{types.StopReasonMaxTokens, "length"},
+		{types.StopReasonToolUse, "tool_calls"},
+		{types.StopReasonStopSequence, "stop"},
+		{types.StopReasonContentFiltered, "content_filter"},
+	}
+
+	for _, tt := range tests {
+		t.Run(string(tt.reason), func(t *testing.T) {
+			events := []types.ConverseStreamOutput{
+				&types.ConverseStreamOutputMemberMessageStop{
+					Value: types.MessageStopEvent{StopReason: tt.reason},
+				},
+			}
+			stream := newMockStream(events)
+			resp, err := parseStreamResponse(context.Background(), stream, nil)
+			require.NoError(t, err)
+			assert.Equal(t, tt.expected, resp.FinishReason)
+		})
+	}
+}