fix-usage-chunk

Author: hustxiayang

internal/extproc/translator/openai_gcpvertexai.go

@@ -175,16 +175,6 @@ func (o *openAIToGCPVertexAITranslatorV1ChatCompletion) handleStreamingResponse(
 		// Convert GCP chunk to OpenAI chunk.
 		openAIChunk := o.convertGCPChunkToOpenAI(chunk)
 
-		// Extract token usage if present in this chunk (typically in the last chunk).
-		if chunk.UsageMetadata != nil {
-			tokenUsage = LLMTokenUsage{
-				InputTokens:       uint32(chunk.UsageMetadata.PromptTokenCount),        //nolint:gosec
-				OutputTokens:      uint32(chunk.UsageMetadata.CandidatesTokenCount),    //nolint:gosec
-				TotalTokens:       uint32(chunk.UsageMetadata.TotalTokenCount),         //nolint:gosec
-				CachedInputTokens: uint32(chunk.UsageMetadata.CachedContentTokenCount), //nolint:gosec
-			}
-		}
-
 		// Serialize to SSE format as expected by OpenAI API.
 		var chunkBytes []byte
 		chunkBytes, err = json.Marshal(openAIChunk)
@@ -198,6 +188,40 @@ func (o *openAIToGCPVertexAITranslatorV1ChatCompletion) handleStreamingResponse(
 		if span != nil {
 			span.RecordResponseChunk(openAIChunk)
 		}
+
+		// Extract token usage only in the last chunk.
+		if chunk.UsageMetadata != nil && chunk.UsageMetadata.PromptTokenCount > 0 {
+			// Convert usage to pointer if available.
+			usage := ptr.To(geminiUsageToOpenAIUsage(chunk.UsageMetadata))
+
+			usageChunk := openai.ChatCompletionResponseChunk{
+				Object:  "chat.completion.chunk",
+				Choices: []openai.ChatCompletionResponseChunkChoice{},
+				// usage is nil for all chunks other than the last chunk
+				Usage: usage,
+			}
+
+			// Serialize to SSE format as expected by OpenAI API.
+			var chunkBytes []byte
+			chunkBytes, err = json.Marshal(usageChunk)
+			if err != nil {
+				return nil, nil, LLMTokenUsage{}, "", fmt.Errorf("error marshaling OpenAI chunk: %w", err)
+			}
+			sseChunkBuf.WriteString("data: ")
+			sseChunkBuf.Write(chunkBytes)
+			sseChunkBuf.WriteString("\n\n")
+
+			if span != nil {
+				span.RecordResponseChunk(openAIChunk)
+			}
+
+			tokenUsage = LLMTokenUsage{
+				InputTokens:       uint32(chunk.UsageMetadata.PromptTokenCount),        //nolint:gosec
+				OutputTokens:      uint32(chunk.UsageMetadata.CandidatesTokenCount),    //nolint:gosec
+				TotalTokens:       uint32(chunk.UsageMetadata.TotalTokenCount),         //nolint:gosec
+				CachedInputTokens: uint32(chunk.UsageMetadata.CachedContentTokenCount), //nolint:gosec
+			}
+		}
 	}
 	mut := &extprocv3.BodyMutation_Body{
 		Body: sseChunkBuf.Bytes(),
@@ -251,16 +275,11 @@ func (o *openAIToGCPVertexAITranslatorV1ChatCompletion) convertGCPChunkToOpenAI(
 		choices = []openai.ChatCompletionResponseChunkChoice{}
 	}
 
-	// Convert usage to pointer if available.
-	var usage *openai.Usage
-	if chunk.UsageMetadata != nil {
-		usage = ptr.To(geminiUsageToOpenAIUsage(chunk.UsageMetadata))
-	}
-
 	return &openai.ChatCompletionResponseChunk{
 		Object:  "chat.completion.chunk",
 		Choices: choices,
-		Usage:   usage,
+		// usage is nil for all chunks other than the last chunk
+		Usage: nil,
 	}
 }
 

internal/extproc/translator/openai_gcpvertexai_test.go

@@ -1054,7 +1054,9 @@ func TestOpenAIToGCPVertexAITranslatorV1ChatCompletion_ResponseBody(t *testing.T
 			wantHeaderMut: nil,
 			wantBodyMut: &extprocv3.BodyMutation{
 				Mutation: &extprocv3.BodyMutation_Body{
-					Body: []byte(`data: {"choices":[{"index":0,"delta":{"content":"Hello","role":"assistant"}}],"object":"chat.completion.chunk","usage":{"prompt_tokens":5,"completion_tokens":3,"total_tokens":8,"completion_tokens_details":{},"prompt_tokens_details":{}}}
+					Body: []byte(`data: {"choices":[{"index":0,"delta":{"content":"Hello","role":"assistant"}}],"object":"chat.completion.chunk"}
+
+data: {"object":"chat.completion.chunk","usage":{"prompt_tokens":5,"completion_tokens":3,"total_tokens":8,"completion_tokens_details":{},"prompt_tokens_details":{}}}
 
 data: [DONE]
 `),