counting

Signed-off-by: yxia216 <yxia216@bloomberg.net>

Author: hustxiayang

internal/extproc/translator/openai_awsbedrock.go

@@ -42,6 +42,7 @@ type openAIToAWSBedrockTranslatorV1ChatCompletion struct {
 	// Translator is created for each request/response stream inside external processor, accordingly the role is not reused by multiple streams.
 	role         string
 	requestModel internalapi.RequestModel
+	toolIndex    *int64
 }
 
 // RequestBody implements [OpenAIChatCompletionTranslator.RequestBody].
@@ -598,6 +599,8 @@ func (o *openAIToAWSBedrockTranslatorV1ChatCompletion) ResponseBody(_ map[string
 		}
 		o.bufferedBody = append(o.bufferedBody, buf...)
 		o.extractAmazonEventStreamEvents()
+		toolIndex := int64(0)
+		o.toolIndex = &toolIndex
 
 		for i := range o.events {
 			event := &o.events[i]
@@ -771,6 +774,7 @@ func (o *openAIToAWSBedrockTranslatorV1ChatCompletion) convertEvent(event *awsbe
 				},
 			})
 		case event.Delta.ToolUse != nil:
+
 			chunk.Choices = append(chunk.Choices, openai.ChatCompletionResponseChunkChoice{
 				Index: 0,
 				Delta: &openai.ChatCompletionResponseChunkChoiceDelta{
@@ -781,7 +785,7 @@ func (o *openAIToAWSBedrockTranslatorV1ChatCompletion) convertEvent(event *awsbe
 								Arguments: event.Delta.ToolUse.Input,
 							},
 							Type:  openai.ChatCompletionMessageToolCallTypeFunction,
-							Index: 0,
+							Index: *o.toolIndex,
 						},
 					},
 				},
@@ -819,11 +823,12 @@ func (o *openAIToAWSBedrockTranslatorV1ChatCompletion) convertEvent(event *awsbe
 								Name: event.Start.ToolUse.Name,
 							},
 							Type:  openai.ChatCompletionMessageToolCallTypeFunction,
-							Index: 0,
+							Index: *o.toolIndex,
 						},
 					},
 				},
 			})
+			*o.toolIndex++
 		}
 	case event.StopReason != nil:
 		chunk.Choices = append(chunk.Choices, openai.ChatCompletionResponseChunkChoice{

internal/extproc/translator/openai_gcpanthropic_stream.go

@@ -40,7 +40,8 @@ type streamingToolCall struct {
 type anthropicStreamParser struct {
 	buffer          bytes.Buffer
 	activeMessageID string
-	activeToolCalls map[int]*streamingToolCall
+	activeToolCalls map[int64]*streamingToolCall
+	toolIndex       *int64
 	tokenUsage      LLMTokenUsage
 	stopReason      anthropic.StopReason
 	requestModel    internalapi.RequestModel
@@ -51,7 +52,7 @@ type anthropicStreamParser struct {
 func newAnthropicStreamParser(requestModel string) *anthropicStreamParser {
 	return &anthropicStreamParser{
 		requestModel:    requestModel,
-		activeToolCalls: make(map[int]*streamingToolCall),
+		activeToolCalls: make(map[int64]*streamingToolCall),
 	}
 }
 
@@ -124,15 +125,15 @@ func (p *anthropicStreamParser) Process(body io.Reader, endOfStream bool, span t
 
 		// Add active tool calls to the final chunk.
 		var toolCalls []openai.ChatCompletionChunkChoiceDeltaToolCall
-		for _, tool := range p.activeToolCalls {
+		for toolIndex, tool := range p.activeToolCalls {
 			toolCalls = append(toolCalls, openai.ChatCompletionChunkChoiceDeltaToolCall{
 				ID:   &tool.id,
 				Type: openai.ChatCompletionMessageToolCallTypeFunction,
 				Function: openai.ChatCompletionMessageToolCallFunctionParam{
 					Name:      tool.name,
 					Arguments: tool.inputJSON,
 				},
-				Index: 0,
+				Index: toolIndex,
 			})
 		}
 
@@ -197,6 +198,7 @@ func (p *anthropicStreamParser) handleAnthropicStreamEvent(eventType []byte, dat
 		p.activeMessageID = event.Message.ID
 		p.tokenUsage.InputTokens = uint32(event.Message.Usage.InputTokens)                 //nolint:gosec
 		p.tokenUsage.CachedInputTokens += uint32(event.Message.Usage.CacheReadInputTokens) //nolint:gosec
+		*p.toolIndex = 0
 		return nil, nil
 
 	case string(constant.ValueOf[constant.ContentBlockStart]()):
@@ -205,7 +207,6 @@ func (p *anthropicStreamParser) handleAnthropicStreamEvent(eventType []byte, dat
 			return nil, fmt.Errorf("failed to unmarshal content_block_start: %w", err)
 		}
 		if event.ContentBlock.Type == string(constant.ValueOf[constant.ToolUse]()) || event.ContentBlock.Type == string(constant.ValueOf[constant.ServerToolUse]()) {
-			toolIdx := int(event.Index)
 			var argsJSON string
 			// Check if the input field is provided directly in the start event.
 			if event.ContentBlock.Input != nil {
@@ -228,15 +229,15 @@ func (p *anthropicStreamParser) handleAnthropicStreamEvent(eventType []byte, dat
 			}
 
 			// Store the complete input JSON in our state.
-			p.activeToolCalls[toolIdx] = &streamingToolCall{
+			p.activeToolCalls[*p.toolIndex] = &streamingToolCall{
 				id:        event.ContentBlock.ID,
 				name:      event.ContentBlock.Name,
 				inputJSON: argsJSON,
 			}
 			delta := openai.ChatCompletionResponseChunkChoiceDelta{
 				ToolCalls: []openai.ChatCompletionChunkChoiceDeltaToolCall{
 					{
-						Index: 0,
+						Index: *p.toolIndex,
 						ID:    &event.ContentBlock.ID,
 						Type:  openai.ChatCompletionMessageToolCallTypeFunction,
 						Function: openai.ChatCompletionMessageToolCallFunctionParam{
@@ -247,6 +248,7 @@ func (p *anthropicStreamParser) handleAnthropicStreamEvent(eventType []byte, dat
 					},
 				},
 			}
+			*p.toolIndex++
 			return p.constructOpenAIChatCompletionChunk(delta, ""), nil
 		}
 		if event.ContentBlock.Type == string(constant.ValueOf[constant.Thinking]()) {
@@ -284,14 +286,14 @@ func (p *anthropicStreamParser) handleAnthropicStreamEvent(eventType []byte, dat
 			delta := openai.ChatCompletionResponseChunkChoiceDelta{Content: &event.Delta.Text}
 			return p.constructOpenAIChatCompletionChunk(delta, ""), nil
 		case string(constant.ValueOf[constant.InputJSONDelta]()):
-			tool, ok := p.activeToolCalls[int(event.Index)]
+			tool, ok := p.activeToolCalls[*p.toolIndex]
 			if !ok {
 				return nil, fmt.Errorf("received input_json_delta for unknown tool at index %d", event.Index)
 			}
 			delta := openai.ChatCompletionResponseChunkChoiceDelta{
 				ToolCalls: []openai.ChatCompletionChunkChoiceDeltaToolCall{
 					{
-						Index: 0,
+						Index: *p.toolIndex,
 						Function: openai.ChatCompletionMessageToolCallFunctionParam{
 							Arguments: event.Delta.PartialJSON,
 						},
@@ -308,7 +310,7 @@ func (p *anthropicStreamParser) handleAnthropicStreamEvent(eventType []byte, dat
 		if err := json.Unmarshal(data, &event); err != nil {
 			return nil, fmt.Errorf("unmarshal content_block_stop: %w", err)
 		}
-		delete(p.activeToolCalls, int(event.Index))
+		delete(p.activeToolCalls, *p.toolIndex)
 		return nil, nil
 
 	case string(constant.ValueOf[constant.MessageStop]()):