feat: Add tool usage prompt tokens to usage calculation - #1827
feat: Add tool usage prompt tokens to usage calculation#1827RoodraNambisa wants to merge 2 commits into
Conversation
Added ToolUsePromptTokenCount to GeminiUsageMetadata and updated usage calculations in Gemini handlers to include tool use prompt tokens. Also improved logic for completion tokens and prompt token details to ensure accurate token accounting.
WalkthroughAdds ToolUsePromptTokenCount to GeminiUsageMetadata and updates Gemini handlers to include it in PromptTokens; adjusts CompletionTokens to include ThoughtsTokenCount and adds fallbacks to infer missing TextTokens or CompletionTokens. No public API signatures changed. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant Client
participant Handler as Gemini Handler
participant Gemini as Gemini API
Client->>Handler: Request (messages, tools)
Handler->>Gemini: Forward request
Gemini-->>Handler: Response + usage metadata
note right of Gemini #f0f4ff: includes PromptTokenCount,<br/>ToolUsePromptTokenCount, ThoughtsTokenCount, TotalTokens
rect rgba(230,245,255,0.6)
note right of Handler: Updated token accounting
Handler->>Handler: PromptTokens = PromptTokenCount + ToolUsePromptTokenCount
Handler->>Handler: CompletionTokens includes ThoughtsTokenCount
end
alt All prompt modalities zero
note right of Handler #fff7e6: Prompt detail fallback
Handler->>Handler: TextTokens = PromptTokens
end
alt TotalTokens > 0 && CompletionTokens == 0 && TotalTokens ≥ PromptTokens
note right of Handler #e8fff0: Completion fallback
Handler->>Handler: CompletionTokens = TotalTokens - PromptTokens
end
Handler-->>Client: Response + adjusted usage
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes Possibly related PRs
Poem
✨ Finishing touches
🧪 Generate unit tests
Tip 👮 Agentic pre-merge checks are now available in preview!Pro plan users can now enable pre-merge checks in their settings to enforce checklists before merging PRs.
Please see the documentation for more information. Example: reviews:
pre_merge_checks:
custom_checks:
- name: "Undocumented Breaking Changes"
mode: "warning"
instructions: |
Pass/fail criteria: All breaking changes to public APIs, CLI flags, environment variables, configuration keys, database schemas, or HTTP/GraphQL endpoints must be documented in the "Breaking Change" section of the PR description and in CHANGELOG.md. Exclude purely internal or private changes (e.g., code not exported from package entry points or explicitly marked as internal).Please share your feedback with us on this Discord post. 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 Pre-merge checks❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (4)
relay/channel/gemini/relay-gemini-native.go (1)
49-55: Missing IMAGE modality in prompt token detailsWe already handle AUDIO and TEXT; include IMAGE for completeness so fallbacks don’t misclassify image prompt tokens as text.
for _, detail := range geminiResponse.UsageMetadata.PromptTokensDetails { if detail.Modality == "AUDIO" { usage.PromptTokensDetails.AudioTokens = detail.TokenCount } else if detail.Modality == "TEXT" { usage.PromptTokensDetails.TextTokens = detail.TokenCount + } else if detail.Modality == "IMAGE" { + usage.PromptTokensDetails.ImageTokens = detail.TokenCount } }Apply the same change in both occurrences (initial and streaming blocks).
Also applies to: 132-138
relay/channel/gemini/relay-gemini.go (3)
941-947: Also map IMAGE modality to detailsParity with native handler and prevents misclassification in fallbacks.
for _, detail := range geminiResponse.UsageMetadata.PromptTokensDetails { if detail.Modality == "AUDIO" { usage.PromptTokensDetails.AudioTokens = detail.TokenCount } else if detail.Modality == "TEXT" { usage.PromptTokensDetails.TextTokens = detail.TokenCount + } else if detail.Modality == "IMAGE" { + usage.PromptTokensDetails.ImageTokens = detail.TokenCount } }Also applies to: 1061-1067
997-1004: Fix fallback: don’t count tool-use prompt tokens as text; guard against negativesWhen details are absent, set TextTokens to text-only count (PromptTokenCount), not aggregated PromptTokens (which include tool-use). Track the last seen tool-use count during streaming to compute accurately.
@@ - - if usage.PromptTokensDetails.TextTokens == 0 && usage.PromptTokensDetails.AudioTokens == 0 && usage.PromptTokensDetails.ImageTokens == 0 { - usage.PromptTokensDetails.TextTokens = usage.PromptTokens - } + // When all modalities missing, infer text-only by subtracting tool-use tokens seen in-stream. + if usage.PromptTokensDetails.TextTokens == 0 && usage.PromptTokensDetails.AudioTokens == 0 && usage.PromptTokensDetails.ImageTokens == 0 { + textOnly := usage.PromptTokens - toolUsePromptTokensLast + if textOnly < 0 { + textOnly = 0 + } + usage.PromptTokensDetails.TextTokens = textOnly + } if usage.TotalTokens > 0 && usage.CompletionTokens == 0 && usage.TotalTokens >= usage.PromptTokens { usage.CompletionTokens = usage.TotalTokens - usage.PromptTokens }And introduce/maintain the helper variable in this function:
@@ - var imageCount int + var imageCount int + // Track latest tool-use prompt tokens observed during the stream for accurate fallbacks. + toolUsePromptTokensLast := 0 @@ - if geminiResponse.UsageMetadata.TotalTokenCount != 0 { + if geminiResponse.UsageMetadata.TotalTokenCount != 0 { + toolUsePromptTokensLast = geminiResponse.UsageMetadata.ToolUsePromptTokenCount usage.PromptTokens = geminiResponse.UsageMetadata.PromptTokenCount + geminiResponse.UsageMetadata.ToolUsePromptTokenCount
1056-1058: Completion fallback is fine; consider symmetric prompt-details fallback (non-streaming)Optional: If all modalities are zero in non-streaming, set TextTokens to PromptTokenCount (excluding tool-use) for accuracy.
usage.CompletionTokenDetails.ReasoningTokens = geminiResponse.UsageMetadata.ThoughtsTokenCount + // Optional: fill missing prompt details with text-only count. + if usage.PromptTokensDetails.TextTokens == 0 && usage.PromptTokensDetails.AudioTokens == 0 && usage.PromptTokensDetails.ImageTokens == 0 { + textOnly := geminiResponse.UsageMetadata.PromptTokenCount + if textOnly < 0 { + textOnly = 0 + } + usage.PromptTokensDetails.TextTokens = textOnly + } if usage.TotalTokens > 0 && usage.CompletionTokens == 0 && usage.TotalTokens >= usage.PromptTokens { usage.CompletionTokens = usage.TotalTokens - usage.PromptTokens }
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
dto/gemini.go(1 hunks)relay/channel/gemini/relay-gemini-native.go(2 hunks)relay/channel/gemini/relay-gemini.go(3 hunks)
🧰 Additional context used
🧠 Learnings (1)
📚 Learning: 2025-06-21T03:37:41.726Z
Learnt from: 9Ninety
PR: QuantumNous/new-api#1273
File: relay/channel/gemini/relay-gemini.go:97-116
Timestamp: 2025-06-21T03:37:41.726Z
Learning: In relay/channel/gemini/relay-gemini.go, the thinking budget calculation logic (including the MaxOutputTokens multiplication) was introduced in PR #1247. PR #1273 focused specifically on decoupling the thoughts summary feature from thinking budget settings and did not modify the existing thinking budget behavior.
Applied to files:
relay/channel/gemini/relay-gemini-native.gorelay/channel/gemini/relay-gemini.go
🔇 Additional comments (4)
relay/channel/gemini/relay-gemini-native.go (2)
42-45: PromptTokens now include tool-use tokens — goodAggregation matches the new DTO field and intent.
42-45: Double-count risk: confirm provider token semantics.
Search returned no matches; verification inconclusive. Confirm whether geminiResponse.UsageMetadata.PromptTokenCount already includes ToolUsePromptTokenCount—if it does, remove the added ToolUsePromptTokenCount to avoid double-counting.
relay/channel/gemini/relay-gemini-native.go:42-45 (also applies to 128-133)relay/channel/gemini/relay-gemini.go (1)
937-939: Token aggregation updates — goodPromptTokens includes tool-use; CompletionTokens includes thoughts.
dto/gemini.go (1)
275-277: Approve — Gemini DTO addition verifiedDownstream handlers compute PromptTokens as PromptTokenCount + ToolUsePromptTokenCount; CompletionTokens fallback (Total - Prompt) remains where appropriate. No stale token-math found.
RoodraNambisa
changed the title
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (1)
relay/channel/gemini/relay-gemini.go (1)
1004-1012: Don’t overwrite provider/fallback usage; merge text estimate and keep prompt/tool tokens.Replacing the entire usage struct drops provider PromptTokens (including tool‑use) and any earlier fallbacks, and skips combining with image completion estimates. Merge only missing fields and always add text estimate when present.
Apply this diff:
- if usage.CompletionTokens == 0 { - str := responseText.String() - if len(str) > 0 { - usage = service.ResponseText2Usage(responseText.String(), info.UpstreamModelName, info.PromptTokens) - } else { - // 空补全，不需要使用量 - usage = &dto.Usage{} - } - } + str := responseText.String() + if len(str) > 0 { + // Estimate only text-side completion; preserve provider/fallback prompt/tool tokens. + est := service.ResponseText2Usage(str, info.UpstreamModelName, usage.PromptTokens) + if est.CompletionTokens > 0 { + if usage.CompletionTokens > 0 { + usage.CompletionTokens += est.CompletionTokens // add to any image fallback done earlier + } else { + usage.CompletionTokens = est.CompletionTokens + } + } + if usage.TotalTokens == 0 { + if usage.PromptTokens > 0 && usage.CompletionTokens > 0 { + usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens + } else if est.TotalTokens > 0 && usage.CompletionTokens == 0 { + usage.TotalTokens = est.TotalTokens + } + } + }
🧹 Nitpick comments (1)
relay/channel/gemini/relay-gemini.go (1)
991-995: Combine image and text fallbacks.This image fallback sets CompletionTokens and short‑circuits text estimation later. With the change above (Lines 1004‑1012), text tokens will be added on top; ensure both patches land together.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
relay/channel/gemini/relay-gemini.go(3 hunks)
🧰 Additional context used
🧠 Learnings (1)
📚 Learning: 2025-06-21T03:37:41.726Z
Learnt from: 9Ninety
PR: QuantumNous/new-api#1273
File: relay/channel/gemini/relay-gemini.go:97-116
Timestamp: 2025-06-21T03:37:41.726Z
Learning: In relay/channel/gemini/relay-gemini.go, the thinking budget calculation logic (including the MaxOutputTokens multiplication) was introduced in PR #1247. PR #1273 focused specifically on decoupling the thoughts summary feature from thinking budget settings and did not modify the existing thinking budget behavior.
Applied to files:
relay/channel/gemini/relay-gemini.go
🔇 Additional comments (3)
relay/channel/gemini/relay-gemini.go (3)
937-939: Correct: tool-use and thoughts now included in usage.Including ToolUsePromptTokenCount in PromptTokens and ThoughtsTokenCount in CompletionTokens aligns with the PR goal.
1046-1049: Correct: non‑streaming usage includes tool‑use and thoughts.Matches intended accounting; TotalTokens remains sourced from TotalTokenCount.
1052-1054: Safe completion fallback.The guard avoids negatives when computing TotalTokens - PromptTokens. Looks good.
| if usage.PromptTokensDetails.TextTokens == 0 && usage.PromptTokensDetails.AudioTokens == 0 && usage.PromptTokensDetails.ImageTokens == 0 { | ||
| usage.PromptTokensDetails.TextTokens = usage.PromptTokens | ||
| } | ||
| if usage.TotalTokens > 0 && usage.CompletionTokens == 0 && usage.TotalTokens >= usage.PromptTokens { | ||
| usage.CompletionTokens = usage.TotalTokens - usage.PromptTokens | ||
| } |
There was a problem hiding this comment.
Fix prompt-details fallback: IMAGE-only details get misattributed to TEXT.
Because ImageTokens is never populated, this fallback can assign all prompt tokens to TextTokens when the provider only returns IMAGE modality details. Populate ImageTokens from provider details to avoid this misattribution.
Apply this change where provider details are parsed:
for _, detail := range geminiResponse.UsageMetadata.PromptTokensDetails {
if detail.Modality == "AUDIO" {
usage.PromptTokensDetails.AudioTokens = detail.TokenCount
} else if detail.Modality == "TEXT" {
usage.PromptTokensDetails.TextTokens = detail.TokenCount
+ } else if detail.Modality == "IMAGE" {
+ usage.PromptTokensDetails.ImageTokens = detail.TokenCount
}
}Do the same in the non‑streaming path:
for _, detail := range geminiResponse.UsageMetadata.PromptTokensDetails {
if detail.Modality == "AUDIO" {
usage.PromptTokensDetails.AudioTokens = detail.TokenCount
} else if detail.Modality == "TEXT" {
usage.PromptTokensDetails.TextTokens = detail.TokenCount
+ } else if detail.Modality == "IMAGE" {
+ usage.PromptTokensDetails.ImageTokens = detail.TokenCount
}
}📝 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.
| if usage.PromptTokensDetails.TextTokens == 0 && usage.PromptTokensDetails.AudioTokens == 0 && usage.PromptTokensDetails.ImageTokens == 0 { | |
| usage.PromptTokensDetails.TextTokens = usage.PromptTokens | |
| } | |
| if usage.TotalTokens > 0 && usage.CompletionTokens == 0 && usage.TotalTokens >= usage.PromptTokens { | |
| usage.CompletionTokens = usage.TotalTokens - usage.PromptTokens | |
| } | |
| // streaming path: populate IMAGE modality too | |
| for _, detail := range geminiResponse.UsageMetadata.PromptTokensDetails { | |
| if detail.Modality == "AUDIO" { | |
| usage.PromptTokensDetails.AudioTokens = detail.TokenCount | |
| } else if detail.Modality == "TEXT" { | |
| usage.PromptTokensDetails.TextTokens = detail.TokenCount | |
| } else if detail.Modality == "IMAGE" { | |
| usage.PromptTokensDetails.ImageTokens = detail.TokenCount | |
| } | |
| } | |
| // non-streaming path: populate IMAGE modality too | |
| for _, detail := range geminiResponse.UsageMetadata.PromptTokensDetails { | |
| if detail.Modality == "AUDIO" { | |
| usage.PromptTokensDetails.AudioTokens = detail.TokenCount | |
| } else if detail.Modality == "TEXT" { | |
| usage.PromptTokensDetails.TextTokens = detail.TokenCount | |
| } else if detail.Modality == "IMAGE" { | |
| usage.PromptTokensDetails.ImageTokens = detail.TokenCount | |
| } | |
| } |
🤖 Prompt for AI Agents
In relay/channel/gemini/relay-gemini.go around lines 997 to 1002, the
prompt-details fallback currently assigns all prompt tokens to TextTokens when
ImageTokens is never populated; update the parsing earlier where provider
details are read (both streaming and non‑streaming paths) to populate
usage.PromptTokensDetails.ImageTokens from the provider's image-modality token
count when present, and then keep the fallback check but only trigger when
TextTokens, AudioTokens AND ImageTokens are all zero so IMAGE-only replies are
not misattributed to TEXT.
2 participants
Added ToolUsePromptTokenCount to GeminiUsageMetadata and updated usage calculations in Gemini handlers to include tool use prompt tokens. Also improved logic for completion tokens and prompt token details to ensure accurate token accounting.
Summary by CodeRabbit
New Features
Bug Fixes