Only advance streamed tool-call index for tool-use blocks

Incrementing toolIndex on every ContentBlockStop caused it to drift
when non-tool blocks (e.g. text) preceded tool-use blocks, making
Bedrock tool calls disappear from the accumulated response.

Track whether the current content block is a tool-use block and only
bump the index when a tool-use block ends.

Signed-off-by: Bryan Frimin <bryan@getprobo.com>
This commit is contained in:
Bryan Frimin
2026-03-13 18:26:38 +01:00
parent 7c77c217c7
commit 3b9a13b5d9

View File

@@ -347,6 +347,7 @@ type bedrockStream struct {
current llm.ChatCompletionStreamEvent
err error
toolIndex int
inToolUse bool
}
func newBedrockStream(eventStream *bedrockruntime.ConverseStreamEventStream) *bedrockStream {
@@ -387,6 +388,7 @@ func (s *bedrockStream) mapEvent(event types.ConverseStreamOutput) (llm.ChatComp
switch e := event.(type) {
case *types.ConverseStreamOutputMemberContentBlockStart:
if start, ok := e.Value.Start.(*types.ContentBlockStartMemberToolUse); ok {
s.inToolUse = true
return llm.ChatCompletionStreamEvent{
Delta: llm.MessageDelta{
ToolCalls: []llm.ToolCallDelta{{
@@ -397,6 +399,7 @@ func (s *bedrockStream) mapEvent(event types.ConverseStreamOutput) (llm.ChatComp
},
}, true
}
s.inToolUse = false
return llm.ChatCompletionStreamEvent{}, false
case *types.ConverseStreamOutputMemberContentBlockDelta:
@@ -418,7 +421,10 @@ func (s *bedrockStream) mapEvent(event types.ConverseStreamOutput) (llm.ChatComp
return llm.ChatCompletionStreamEvent{}, false
case *types.ConverseStreamOutputMemberContentBlockStop:
s.toolIndex++
if s.inToolUse {
s.toolIndex++
s.inToolUse = false
}
return llm.ChatCompletionStreamEvent{}, false
case *types.ConverseStreamOutputMemberMessageStop: