fix: count toolUsePrompt as input and normalize stream/chat usage - #2947
Closed
seefs001 wants to merge 1 commit into
Closed
fix: count toolUsePrompt as input and normalize stream/chat usage#2947seefs001 wants to merge 1 commit into
seefs001 wants to merge 1 commit into
Conversation
Contributor
WalkthroughAdds tool-use token tracking fields to GeminiUsageMetadata DTO, introduces helper functions to centralize usage construction from Gemini response metadata, and implements fallback logic for token count derivation across response handlers. Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~28 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 5 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Fix all issues with AI agents
In `@relay/channel/gemini/usage_test.go`:
- Around line 26-33: buildUsageFromGeminiMetadata currently sets promptTokens :=
meta.PromptTokenCount but never adds meta.ToolUsePromptTokenCount and never
iterates meta.ToolUsePromptDetails, so PromptTokens and
PromptTokensDetails.TextTokens/CachedTokens are missing tool-use contributions;
update buildUsageFromGeminiMetadata to add meta.ToolUsePromptTokenCount into the
promptTokens total (so usage.PromptTokens = meta.PromptTokenCount +
meta.ToolUsePromptTokenCount) and iterate meta.ToolUsePromptDetails (similar to
how meta.PromptTokensDetails are processed) to accumulate TEXT into
PromptTokensDetails.TextTokens and CACHED into PromptTokensDetails.CachedTokens
(and any other relevant token categories) before assigning to
usage.PromptTokensDetails and usage.PromptTokens.
In `@relay/channel/gemini/usage.go`:
- Around line 16-57: In buildUsageFromGeminiMetadata, promptTokens is set from
meta.PromptTokenCount but never includes meta.ToolUsePromptTokenCount and
meta.ToolUsePromptDetails are never applied; update the promptTokens calculation
to add meta.ToolUsePromptTokenCount (when present) and iterate
meta.ToolUsePromptDetails to accumulate their TokenCount into
usage.PromptTokensDetails (e.g., a ToolUseTokens field or add into modality
buckets) and into promptTokens so the input = promptTokenCount +
toolUsePromptTokenCount invariant holds; adjust the existing loop over
meta.PromptTokensDetails to also incorporate tool-use modality breakdowns (map
AUDIO/TEXT to usage.PromptTokensDetails.AudioTokens/TextTokens or to a
ToolUseTokens field) and ensure downstream derived values (InputTokens,
PromptTokensDetails.TextTokens fallback) reflect the added tool-use tokens.
🧹 Nitpick comments (1)
relay/channel/gemini/relay-gemini.go (1)
1287-1289: Minor: TextTokens fallback is duplicated frombuildUsageFromGeminiMetadata.Lines 55–57 of
usage.goalready apply the sameTextTokens = PromptTokensfallback. This code is harmless but redundant in the streaming path since*usage = latestUsagealready includes the fallback result.
Comment on lines
+26
to
+33
| usage := buildUsageFromGeminiMetadata(meta, 0) | ||
|
|
||
| require.Equal(t, 18480, usage.PromptTokens) | ||
| require.Equal(t, 2209, usage.CompletionTokens) | ||
| require.Equal(t, 20689, usage.TotalTokens) | ||
| require.Equal(t, 1120, usage.CompletionTokenDetails.ReasoningTokens) | ||
| require.Equal(t, 17, usage.PromptTokensDetails.CachedTokens) | ||
| require.Equal(t, 18480, usage.PromptTokensDetails.TextTokens) |
Contributor
There was a problem hiding this comment.
Tests will fail: buildUsageFromGeminiMetadata does not add ToolUsePromptTokenCount to PromptTokens.
Tests 1–3 all expect PromptTokens = PromptTokenCount + ToolUsePromptTokenCount, but usage.go line 17 only sets promptTokens := meta.PromptTokenCount without adding meta.ToolUsePromptTokenCount. For example, test 1 expects PromptTokens=18480 (151+18329) but the code will produce 151.
Similarly, Line 33 expects TextTokens=18480, but the code only iterates meta.PromptTokensDetails (yielding TEXT=151) and the fallback on line 55–57 of usage.go won't trigger since TextTokens != 0. The ToolUsePromptDetails slice is never iterated.
The root cause is in usage.go — see the comment there for the proposed fix.
🤖 Prompt for AI Agents
In `@relay/channel/gemini/usage_test.go` around lines 26 - 33,
buildUsageFromGeminiMetadata currently sets promptTokens :=
meta.PromptTokenCount but never adds meta.ToolUsePromptTokenCount and never
iterates meta.ToolUsePromptDetails, so PromptTokens and
PromptTokensDetails.TextTokens/CachedTokens are missing tool-use contributions;
update buildUsageFromGeminiMetadata to add meta.ToolUsePromptTokenCount into the
promptTokens total (so usage.PromptTokens = meta.PromptTokenCount +
meta.ToolUsePromptTokenCount) and iterate meta.ToolUsePromptDetails (similar to
how meta.PromptTokensDetails are processed) to accumulate TEXT into
PromptTokensDetails.TextTokens and CACHED into PromptTokensDetails.CachedTokens
(and any other relevant token categories) before assigning to
usage.PromptTokensDetails and usage.PromptTokens.
Comment on lines
+16
to
+57
| func buildUsageFromGeminiMetadata(meta dto.GeminiUsageMetadata, estimatedPromptTokens int) dto.Usage { | ||
| promptTokens := meta.PromptTokenCount | ||
| if promptTokens <= 0 { | ||
| promptTokens = estimatedPromptTokens | ||
| } | ||
|
|
||
| completionTokens := meta.CandidatesTokenCount + meta.ThoughtsTokenCount | ||
| if completionTokens <= 0 && meta.TotalTokenCount > 0 { | ||
| completionTokens = meta.TotalTokenCount - meta.PromptTokenCount - meta.ToolUsePromptTokenCount | ||
| } | ||
| if completionTokens < 0 { | ||
| completionTokens = 0 | ||
| } | ||
|
|
||
| totalTokens := meta.TotalTokenCount | ||
| if totalTokens <= 0 { | ||
| totalTokens = promptTokens + completionTokens | ||
| } | ||
|
|
||
| usage := dto.Usage{ | ||
| PromptTokens: promptTokens, | ||
| CompletionTokens: completionTokens, | ||
| TotalTokens: totalTokens, | ||
| InputTokens: promptTokens, | ||
| OutputTokens: completionTokens, | ||
| } | ||
|
|
||
| usage.CompletionTokenDetails.ReasoningTokens = meta.ThoughtsTokenCount | ||
| usage.PromptTokensDetails.CachedTokens = meta.CachedContentTokenCount | ||
|
|
||
| for _, detail := range meta.PromptTokensDetails { | ||
| switch detail.Modality { | ||
| case "AUDIO": | ||
| usage.PromptTokensDetails.AudioTokens = detail.TokenCount | ||
| case "TEXT": | ||
| usage.PromptTokensDetails.TextTokens = detail.TokenCount | ||
| } | ||
| } | ||
|
|
||
| if usage.PromptTokensDetails.TextTokens == 0 && usage.PromptTokens > 0 { | ||
| usage.PromptTokensDetails.TextTokens = usage.PromptTokens | ||
| } |
Contributor
There was a problem hiding this comment.
Critical: ToolUsePromptTokenCount is never added to promptTokens, defeating the purpose of this PR.
The PR's stated objective is input = promptTokenCount + toolUsePromptTokenCount, but line 17 only assigns meta.PromptTokenCount without adding meta.ToolUsePromptTokenCount. Additionally, meta.ToolUsePromptDetails is never iterated, so tool-use modality breakdowns are lost.
All three tests that include non-zero ToolUsePromptTokenCount will fail.
🐛 Proposed fix
func buildUsageFromGeminiMetadata(meta dto.GeminiUsageMetadata, estimatedPromptTokens int) dto.Usage {
promptTokens := meta.PromptTokenCount
if promptTokens <= 0 {
promptTokens = estimatedPromptTokens
}
+ promptTokens += meta.ToolUsePromptTokenCount
completionTokens := meta.CandidatesTokenCount + meta.ThoughtsTokenCount
if completionTokens <= 0 && meta.TotalTokenCount > 0 {
completionTokens = meta.TotalTokenCount - meta.PromptTokenCount - meta.ToolUsePromptTokenCount
}
if completionTokens < 0 {
completionTokens = 0
}
totalTokens := meta.TotalTokenCount
if totalTokens <= 0 {
totalTokens = promptTokens + completionTokens
}
usage := dto.Usage{
PromptTokens: promptTokens,
CompletionTokens: completionTokens,
TotalTokens: totalTokens,
InputTokens: promptTokens,
OutputTokens: completionTokens,
}
usage.CompletionTokenDetails.ReasoningTokens = meta.ThoughtsTokenCount
usage.PromptTokensDetails.CachedTokens = meta.CachedContentTokenCount
for _, detail := range meta.PromptTokensDetails {
switch detail.Modality {
case "AUDIO":
usage.PromptTokensDetails.AudioTokens = detail.TokenCount
case "TEXT":
usage.PromptTokensDetails.TextTokens = detail.TokenCount
}
}
+ for _, detail := range meta.ToolUsePromptDetails {
+ switch detail.Modality {
+ case "AUDIO":
+ usage.PromptTokensDetails.AudioTokens += detail.TokenCount
+ case "TEXT":
+ usage.PromptTokensDetails.TextTokens += detail.TokenCount
+ }
+ }
if usage.PromptTokensDetails.TextTokens == 0 && usage.PromptTokens > 0 {
usage.PromptTokensDetails.TextTokens = usage.PromptTokens
}
return usage
}🤖 Prompt for AI Agents
In `@relay/channel/gemini/usage.go` around lines 16 - 57, In
buildUsageFromGeminiMetadata, promptTokens is set from meta.PromptTokenCount but
never includes meta.ToolUsePromptTokenCount and meta.ToolUsePromptDetails are
never applied; update the promptTokens calculation to add
meta.ToolUsePromptTokenCount (when present) and iterate
meta.ToolUsePromptDetails to accumulate their TokenCount into
usage.PromptTokensDetails (e.g., a ToolUseTokens field or add into modality
buckets) and into promptTokens so the input = promptTokenCount +
toolUsePromptTokenCount invariant holds; adjust the existing loop over
meta.PromptTokensDetails to also incorporate tool-use modality breakdowns (map
AUDIO/TEXT to usage.PromptTokensDetails.AudioTokens/TextTokens or to a
ToolUseTokens field) and ensure downstream derived values (InputTokens,
PromptTokensDetails.TextTokens fallback) reflect the added tool-use tokens.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
fix #2935
input = promptTokenCount + toolUsePromptTokenCount
output = candidatesTokenCount + thoughtsTokenCount / totalTokenCount - input
针对流和非流两种情况
Summary by CodeRabbit
Bug Fixes
Tests