Skip to content

feat: Add tool usage prompt tokens to usage calculation #1827

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

Closed
RoodraNambisa wants to merge 2 commits into from
Closed

feat: Add tool usage prompt tokens to usage calculation #1827

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
3293c0e
Include tool use prompt tokens in usage calculation
RoodraNambisa Sep 17, 2025
9be615f
Update relay-gemini.go
RoodraNambisa Sep 17, 2025
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
1 change: 1 addition & 0 deletions dto/gemini.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -272,6 +272,7 @@ type GeminiUsageMetadata struct {
CandidatesTokenCount int `json:"candidatesTokenCount"`
TotalTokenCount int `json:"totalTokenCount"`
ThoughtsTokenCount int `json:"thoughtsTokenCount"`
ToolUsePromptTokenCount int `json:"toolUsePromptTokenCount"`
PromptTokensDetails []GeminiPromptTokensDetails `json:"promptTokensDetails"`
}

Expand Down
4 changes: 2 additions & 2 deletions relay/channel/gemini/relay-gemini-native.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,7 @@ func GeminiTextGenerationHandler(c *gin.Context, info *relaycommon.RelayInfo, re

// 计算使用量（基于 UsageMetadata）
usage := dto.Usage{
PromptTokens: geminiResponse.UsageMetadata.PromptTokenCount,
PromptTokens: geminiResponse.UsageMetadata.PromptTokenCount + geminiResponse.UsageMetadata.ToolUsePromptTokenCount,
CompletionTokens: geminiResponse.UsageMetadata.CandidatesTokenCount + geminiResponse.UsageMetadata.ThoughtsTokenCount,
TotalTokens: geminiResponse.UsageMetadata.TotalTokenCount,
}
Expand Down Expand Up @@ -125,7 +125,7 @@ func GeminiTextGenerationStreamHandler(c *gin.Context, info *relaycommon.RelayIn

// 更新使用量统计
if geminiResponse.UsageMetadata.TotalTokenCount != 0 {
usage.PromptTokens = geminiResponse.UsageMetadata.PromptTokenCount
usage.PromptTokens = geminiResponse.UsageMetadata.PromptTokenCount + geminiResponse.UsageMetadata.ToolUsePromptTokenCount
usage.CompletionTokens = geminiResponse.UsageMetadata.CandidatesTokenCount + geminiResponse.UsageMetadata.ThoughtsTokenCount
usage.TotalTokens = geminiResponse.UsageMetadata.TotalTokenCount
usage.CompletionTokenDetails.ReasoningTokens = geminiResponse.UsageMetadata.ThoughtsTokenCount
Expand Down
21 changes: 14 additions & 7 deletions relay/channel/gemini/relay-gemini.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -934,8 +934,8 @@ func GeminiChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *
response.Created = createAt
response.Model = info.UpstreamModelName
if geminiResponse.UsageMetadata.TotalTokenCount != 0 {
usage.PromptTokens = geminiResponse.UsageMetadata.PromptTokenCount
usage.CompletionTokens = geminiResponse.UsageMetadata.CandidatesTokenCount
usage.PromptTokens = geminiResponse.UsageMetadata.PromptTokenCount + geminiResponse.UsageMetadata.ToolUsePromptTokenCount
usage.CompletionTokens = geminiResponse.UsageMetadata.CandidatesTokenCount + geminiResponse.UsageMetadata.ThoughtsTokenCount
usage.CompletionTokenDetails.ReasoningTokens = geminiResponse.UsageMetadata.ThoughtsTokenCount
usage.TotalTokens = geminiResponse.UsageMetadata.TotalTokenCount
for _, detail := range geminiResponse.UsageMetadata.PromptTokensDetails {
Expand Down Expand Up @@ -994,8 +994,12 @@ func GeminiChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *
}
}

usage.PromptTokensDetails.TextTokens = usage.PromptTokens
usage.CompletionTokens = usage.TotalTokens - usage.PromptTokens
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
}
Comment on lines +997 to +1002

@coderabbitai coderabbitai Bot Sep 17, 2025

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.

⚠️ Potential issue

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.

Suggested change
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.


if usage.CompletionTokens == 0 {
str := responseText.String()
Expand Down Expand Up @@ -1039,13 +1043,16 @@ func GeminiChatHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.R
fullTextResponse := responseGeminiChat2OpenAI(c, &geminiResponse)
fullTextResponse.Model = info.UpstreamModelName
usage := dto.Usage{
PromptTokens: geminiResponse.UsageMetadata.PromptTokenCount,
CompletionTokens: geminiResponse.UsageMetadata.CandidatesTokenCount,
PromptTokens: geminiResponse.UsageMetadata.PromptTokenCount + geminiResponse.UsageMetadata.ToolUsePromptTokenCount,
CompletionTokens: geminiResponse.UsageMetadata.CandidatesTokenCount + geminiResponse.UsageMetadata.ThoughtsTokenCount,
TotalTokens: geminiResponse.UsageMetadata.TotalTokenCount,
}

usage.CompletionTokenDetails.ReasoningTokens = geminiResponse.UsageMetadata.ThoughtsTokenCount
usage.CompletionTokens = usage.TotalTokens - usage.PromptTokens
if usage.TotalTokens > 0 && usage.CompletionTokens == 0 && usage.TotalTokens >= usage.PromptTokens {
usage.CompletionTokens = usage.TotalTokens - usage.PromptTokens
}


for _, detail := range geminiResponse.UsageMetadata.PromptTokensDetails {
if detail.Modality == "AUDIO" {
Expand Down