fix: add reasoning tokens into token metrics in stream mode for gemini models (#2259)

**Description**
The previous implementation on tokens tracking is like 
```
if chunk.UsageMetadata.CandidatesTokenCount >= 0 {
				tokenUsage.SetOutputTokens(uint32(chunk.UsageMetadata.CandidatesTokenCount)) //nolint:gosec
			}
```
The output tokens should include both the reasoning tokens and the
candidates tokens. Thus fix this by reusing the the variable `usage`,
just like non-streaming mode. Also updated the tests to include
reasoning tokens

---------

Signed-off-by: yxia216 <yxia216@bloomberg.net>
Co-authored-by: Ignasi Barrera <ignasi@tetrate.io>

Author: hustxiayang

internal/translator/openai_gcpvertexai.go

@@ -249,21 +249,11 @@ func (o *openAIToGCPVertexAITranslatorV1ChatCompletion) handleStreamingResponse(
 				span.RecordResponseChunk(usageChunk)
 			}
 
-			if chunk.UsageMetadata.PromptTokenCount >= 0 {
-				tokenUsage.SetInputTokens(uint32(chunk.UsageMetadata.PromptTokenCount)) //nolint:gosec
-			}
-			if chunk.UsageMetadata.CandidatesTokenCount >= 0 {
-				tokenUsage.SetOutputTokens(uint32(chunk.UsageMetadata.CandidatesTokenCount)) //nolint:gosec
-			}
-			if chunk.UsageMetadata.TotalTokenCount >= 0 {
-				tokenUsage.SetTotalTokens(uint32(chunk.UsageMetadata.TotalTokenCount)) //nolint:gosec
-			}
-			if chunk.UsageMetadata.CachedContentTokenCount >= 0 {
-				tokenUsage.SetCachedInputTokens(uint32(chunk.UsageMetadata.CachedContentTokenCount)) //nolint:gosec
-			}
-			if chunk.UsageMetadata.ThoughtsTokenCount >= 0 {
-				tokenUsage.SetReasoningTokens(uint32(chunk.UsageMetadata.ThoughtsTokenCount)) //nolint:gosec
-			}
+			tokenUsage.SetInputTokens(uint32(usage.PromptTokens))                                //nolint:gosec
+			tokenUsage.SetOutputTokens(uint32(usage.CompletionTokens))                           //nolint:gosec
+			tokenUsage.SetTotalTokens(uint32(usage.TotalTokens))                                 //nolint:gosec
+			tokenUsage.SetCachedInputTokens(uint32(usage.PromptTokensDetails.CachedTokens))      //nolint:gosec
+			tokenUsage.SetReasoningTokens(uint32(usage.CompletionTokensDetails.ReasoningTokens)) //nolint:gosec
 		}
 	}
 

internal/translator/openai_gcpvertexai_test.go

@@ -1230,7 +1230,7 @@ data: [DONE]
 			},
 			body: `data: {"candidates":[{"content":{"parts":[{"text":"let me think step by step and reply you.", "thought": true}]}}]}
 
-data: {"candidates":[{"content":{"parts":[{"text":"Hello"}]}}],"usageMetadata":{"promptTokenCount":5,"candidatesTokenCount":3,"totalTokenCount":8}}`,
+data: {"candidates":[{"content":{"parts":[{"text":"Hello"}]}}],"usageMetadata":{"promptTokenCount":5,"candidatesTokenCount":3,"totalTokenCount":18,"thoughtsTokenCount":10}}`,
 			stream:        true,
 			endOfStream:   true,
 			wantError:     false,
@@ -1239,11 +1239,11 @@ data: {"candidates":[{"content":{"parts":[{"text":"Hello"}]}}],"usageMetadata":{
 
 data: {"choices":[{"index":0,"delta":{"content":"Hello","role":"assistant"}}],"object":"chat.completion.chunk"}
 
-data: {"choices":[],"object":"chat.completion.chunk","usage":{"prompt_tokens":5,"completion_tokens":3,"total_tokens":8,"completion_tokens_details":{},"prompt_tokens_details":{}}}
+data: {"choices":[],"object":"chat.completion.chunk","usage":{"prompt_tokens":5,"completion_tokens":13,"total_tokens":18,"completion_tokens_details":{"reasoning_tokens":10},"prompt_tokens_details":{}}}
 
 data: [DONE]
 `),
-			wantTokenUsage: tokenUsageFrom(5, 0, -1, 3, 8, 0), // Does not support Cache Creation.
+			wantTokenUsage: tokenUsageFrom(5, 0, -1, 13, 18, 10), // Does not support Cache Creation.
 		},
 		{
 			name: "stream chunks with thought signature on text part",
@@ -1252,7 +1252,7 @@ data: [DONE]
 			},
 			body: `data: {"candidates":[{"content":{"parts":[{"text":"let me think about this.", "thought": true}]}}]}
 
-data: {"candidates":[{"content":{"parts":[{"text":"The answer is 42.", "thoughtSignature": "dGVzdHNpZ25hdHVyZQ=="}]}}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":8,"totalTokenCount":18}}`,
+data: {"candidates":[{"content":{"parts":[{"text":"The answer is 42.", "thoughtSignature": "dGVzdHNpZ25hdHVyZQ=="}]}}],"usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":8,"totalTokenCount":33,"thoughtsTokenCount":15}}`,
 			stream:        true,
 			endOfStream:   true,
 			wantError:     false,
@@ -1261,11 +1261,11 @@ data: {"candidates":[{"content":{"parts":[{"text":"The answer is 42.", "thoughtS
 
 data: {"choices":[{"index":0,"delta":{"content":"The answer is 42.","role":"assistant","reasoning_content":{"signature":"dGVzdHNpZ25hdHVyZQ=="}}}],"object":"chat.completion.chunk"}
 
-data: {"choices":[],"object":"chat.completion.chunk","usage":{"prompt_tokens":10,"completion_tokens":8,"total_tokens":18,"completion_tokens_details":{},"prompt_tokens_details":{}}}
+data: {"choices":[],"object":"chat.completion.chunk","usage":{"prompt_tokens":10,"completion_tokens":23,"total_tokens":33,"completion_tokens_details":{"reasoning_tokens":15},"prompt_tokens_details":{}}}
 
 data: [DONE]
 `),
-			wantTokenUsage: tokenUsageFrom(10, 0, -1, 8, 18, 0),
+			wantTokenUsage: tokenUsageFrom(10, 0, -1, 23, 33, 15),
 		},
 	}