Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
aabchoo merged 6 commits into from
Jul 2, 2026
Merged

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

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
43e84ea
fix-reasong-tokens-token-metrics
hustxiayang Jun 18, 2026
f32d197
Merge branch 'main' into fix-gemini-streaming-reasoning-tokens
hustxiayang Jun 18, 2026
0d31f10
tests
hustxiayang Jun 18, 2026
362ebe6
Merge branch 'main' into fix-gemini-streaming-reasoning-tokens
nacx Jun 22, 2026
8f1dc52
Merge branch 'main' into fix-gemini-streaming-reasoning-tokens
hustxiayang Jul 2, 2026
36e01f1
Merge branch 'main' into fix-gemini-streaming-reasoning-tokens
hustxiayang Jul 2, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 5 additions & 15 deletions internal/translator/openai_gcpvertexai.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Comment on lines +255 to +256

@gemini-code-assist gemini-code-assist Bot Jun 18, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

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
			}

@hustxiayang hustxiayang Jun 18, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

in function geminiUsageToOpenAIUsage, these values are always set

}
}

Expand Down
12 changes: 6 additions & 6 deletions internal/translator/openai_gcpvertexai_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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",
Expand All @@ -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,
Expand All @@ -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),
},
}

Expand Down