fix: add reasoning tokens into token metrics in stream mode for gemini models#2259
Conversation
Signed-off-by: yxia216 <yxia216@bloomberg.net>
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request refactors how token usage is extracted and recorded during streaming responses in the OpenAI to GCP Vertex AI translator, switching from chunk.UsageMetadata to direct fields on the usage object. However, accessing nested fields like usage.PromptTokensDetails.CachedTokens and usage.CompletionTokensDetails.ReasoningTokens without nil checks introduces a risk of nil pointer dereference panics if those details are not populated in the stream. It is recommended to add defensive nil checks for these fields.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| tokenUsage.SetCachedInputTokens(uint32(usage.PromptTokensDetails.CachedTokens)) //nolint:gosec | ||
| tokenUsage.SetReasoningTokens(uint32(usage.CompletionTokensDetails.ReasoningTokens)) //nolint:gosec |
There was a problem hiding this comment.
Directly accessing usage.PromptTokensDetails and usage.CompletionTokensDetails without nil checks can lead to a nil pointer dereference panic if these fields are not populated in the streaming chunk. Adding defensive nil checks (similar to the non-streaming implementation) prevents potential runtime panics.
if usage.PromptTokensDetails != nil {
tokenUsage.SetCachedInputTokens(uint32(usage.PromptTokensDetails.CachedTokens)) //nolint:gosec
}
if usage.CompletionTokensDetails != nil {
tokenUsage.SetReasoningTokens(uint32(usage.CompletionTokensDetails.ReasoningTokens)) //nolint:gosec
}There was a problem hiding this comment.
in function geminiUsageToOpenAIUsage, these values are always set
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #2259 +/- ##
==========================================
- Coverage 84.90% 84.88% -0.02%
==========================================
Files 144 144
Lines 21259 21254 -5
==========================================
- Hits 18049 18042 -7
- Misses 2131 2132 +1
- Partials 1079 1080 +1 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
/retest |
Description
The previous implementation on tokens tracking is like
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