fix: unify usage mapping and include toolUsePromptTokenCount#2959
Open
seefs001 wants to merge 1 commit intoQuantumNous:mainfrom
Open
fix: unify usage mapping and include toolUsePromptTokenCount#2959seefs001 wants to merge 1 commit intoQuantumNous:mainfrom
seefs001 wants to merge 1 commit intoQuantumNous:mainfrom
Conversation
Contributor
WalkthroughThis PR adds separate tracking for tool-use prompt tokens in the Gemini DTO and centralizes token usage computation through a new helper function that properly excludes tool-use tokens from completion token calculations, ensuring accurate billing. Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~22 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing touches
🧪 Generate unit tests (beta)
Tip Issue Planner is now in beta. Read the docs and try it out! Share your feedback on Discord. 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: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@relay/channel/gemini/relay-gemini.go`:
- Around line 1064-1066: The fallback computation for CompletionTokens can
produce a negative value when fallbackPromptTokens > TotalTokens; in the block
that currently sets usage.CompletionTokens = usage.TotalTokens -
usage.PromptTokens, clamp the result to a non-negative value (e.g., set
usage.CompletionTokens to max(0, usage.TotalTokens - usage.PromptTokens)) so
CompletionTokens never becomes negative; ensure you reference the same variables
used in this file (usage.TotalTokens, usage.PromptTokens, usage.CompletionTokens
and the previously computed fallbackPromptTokens) and preserve the existing
branch logic.
Comment on lines
+1064
to
+1066
| if usage.TotalTokens > 0 && usage.CompletionTokens <= 0 { | ||
| usage.CompletionTokens = usage.TotalTokens - usage.PromptTokens | ||
| } |
Contributor
There was a problem hiding this comment.
Edge case: fallback CompletionTokens can become negative when fallbackPromptTokens exceeds TotalTokens.
When the API returns PromptTokenCount == 0 and CandidatesTokenCount == 0 but a nonzero TotalTokens, this path computes CompletionTokens = TotalTokens - PromptTokens. If fallbackPromptTokens (used on Line 1038) is larger than TotalTokens, CompletionTokens goes negative.
Consider clamping to zero:
Proposed fix
if usage.TotalTokens > 0 && usage.CompletionTokens <= 0 {
- usage.CompletionTokens = usage.TotalTokens - usage.PromptTokens
+ computed := usage.TotalTokens - usage.PromptTokens
+ if computed < 0 {
+ computed = 0
+ }
+ usage.CompletionTokens = computed
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| if usage.TotalTokens > 0 && usage.CompletionTokens <= 0 { | |
| usage.CompletionTokens = usage.TotalTokens - usage.PromptTokens | |
| } | |
| if usage.TotalTokens > 0 && usage.CompletionTokens <= 0 { | |
| computed := usage.TotalTokens - usage.PromptTokens | |
| if computed < 0 { | |
| computed = 0 | |
| } | |
| usage.CompletionTokens = computed | |
| } |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@relay/channel/gemini/relay-gemini.go` around lines 1064 - 1066, The fallback
computation for CompletionTokens can produce a negative value when
fallbackPromptTokens > TotalTokens; in the block that currently sets
usage.CompletionTokens = usage.TotalTokens - usage.PromptTokens, clamp the
result to a non-negative value (e.g., set usage.CompletionTokens to max(0,
usage.TotalTokens - usage.PromptTokens)) so CompletionTokens never becomes
negative; ensure you reference the same variables used in this file
(usage.TotalTokens, usage.PromptTokens, usage.CompletionTokens and the
previously computed fallbackPromptTokens) and preserve the existing branch
logic.
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
Summary by CodeRabbit
Improvements
Tests