Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit d0b35ed

Browse files
Calcium-IonCalcium-Ion
authored
Merge pull request QuantumNous#2959 from seefs001/fix/gemini-tool-use-token
fix: unify usage mapping and include toolUsePromptTokenCount
2 parents 4b058b4 + 4360393 commit d0b35ed

4 files changed

Lines changed: 386 additions & 71 deletions

File tree

‎dto/gemini.go‎

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -458,12 +458,14 @@ type GeminiChatResponse struct {
458458
}
459459

460460
type GeminiUsageMetadata struct {
461-
PromptTokenCount int `json:"promptTokenCount"`
462-
CandidatesTokenCount int `json:"candidatesTokenCount"`
463-
TotalTokenCount int `json:"totalTokenCount"`
464-
ThoughtsTokenCount int `json:"thoughtsTokenCount"`
465-
CachedContentTokenCount int `json:"cachedContentTokenCount"`
466-
PromptTokensDetails []GeminiPromptTokensDetails `json:"promptTokensDetails"`
461+
PromptTokenCount int `json:"promptTokenCount"`
462+
ToolUsePromptTokenCount int `json:"toolUsePromptTokenCount"`
463+
CandidatesTokenCount int `json:"candidatesTokenCount"`
464+
TotalTokenCount int `json:"totalTokenCount"`
465+
ThoughtsTokenCount int `json:"thoughtsTokenCount"`
466+
CachedContentTokenCount int `json:"cachedContentTokenCount"`
467+
PromptTokensDetails []GeminiPromptTokensDetails `json:"promptTokensDetails"`
468+
ToolUsePromptTokensDetails []GeminiPromptTokensDetails `json:"toolUsePromptTokensDetails"`
467469
}
468470

469471
type GeminiPromptTokensDetails struct {

‎relay/channel/gemini/relay-gemini-native.go‎

Lines changed: 1 addition & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -42,22 +42,7 @@ func GeminiTextGenerationHandler(c *gin.Context, info *relaycommon.RelayInfo, re
4242
}
4343

4444
// 计算使用量（基于 UsageMetadata）
45-
usage := dto.Usage{
46-
PromptTokens: geminiResponse.UsageMetadata.PromptTokenCount,
47-
CompletionTokens: geminiResponse.UsageMetadata.CandidatesTokenCount + geminiResponse.UsageMetadata.ThoughtsTokenCount,
48-
TotalTokens: geminiResponse.UsageMetadata.TotalTokenCount,
49-
}
50-
51-
usage.CompletionTokenDetails.ReasoningTokens = geminiResponse.UsageMetadata.ThoughtsTokenCount
52-
usage.PromptTokensDetails.CachedTokens = geminiResponse.UsageMetadata.CachedContentTokenCount
53-
54-
for _, detail := range geminiResponse.UsageMetadata.PromptTokensDetails {
55-
if detail.Modality == "AUDIO" {
56-
usage.PromptTokensDetails.AudioTokens = detail.TokenCount
57-
} else if detail.Modality == "TEXT" {
58-
usage.PromptTokensDetails.TextTokens = detail.TokenCount
59-
}
60-
}
45+
usage := buildUsageFromGeminiMetadata(geminiResponse.UsageMetadata, info.GetEstimatePromptTokens())
6146

6247
service.IOCopyBytesGracefully(c, resp, responseBody)
6348

‎relay/channel/gemini/relay-gemini.go‎

Lines changed: 44 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -1032,6 +1032,46 @@ func getResponseToolCall(item *dto.GeminiPart) *dto.ToolCallResponse {
10321032
}
10331033
}
10341034

1035+
func buildUsageFromGeminiMetadata(metadata dto.GeminiUsageMetadata, fallbackPromptTokens int) dto.Usage {
1036+
promptTokens := metadata.PromptTokenCount + metadata.ToolUsePromptTokenCount
1037+
if promptTokens <= 0 && fallbackPromptTokens > 0 {
1038+
promptTokens = fallbackPromptTokens
1039+
}
1040+
1041+
usage := dto.Usage{
1042+
PromptTokens: promptTokens,
1043+
CompletionTokens: metadata.CandidatesTokenCount + metadata.ThoughtsTokenCount,
1044+
TotalTokens: metadata.TotalTokenCount,
1045+
}
1046+
usage.CompletionTokenDetails.ReasoningTokens = metadata.ThoughtsTokenCount
1047+
usage.PromptTokensDetails.CachedTokens = metadata.CachedContentTokenCount
1048+
1049+
for _, detail := range metadata.PromptTokensDetails {
1050+
if detail.Modality == "AUDIO" {
1051+
usage.PromptTokensDetails.AudioTokens += detail.TokenCount
1052+
} else if detail.Modality == "TEXT" {
1053+
usage.PromptTokensDetails.TextTokens += detail.TokenCount
1054+
}
1055+
}
1056+
for _, detail := range metadata.ToolUsePromptTokensDetails {
1057+
if detail.Modality == "AUDIO" {
1058+
usage.PromptTokensDetails.AudioTokens += detail.TokenCount
1059+
} else if detail.Modality == "TEXT" {
1060+
usage.PromptTokensDetails.TextTokens += detail.TokenCount
1061+
}
1062+
}
1063+
1064+
if usage.TotalTokens > 0 && usage.CompletionTokens <= 0 {
1065+
usage.CompletionTokens = usage.TotalTokens - usage.PromptTokens
1066+
}
1067+
1068+
if usage.PromptTokens > 0 && usage.PromptTokensDetails.TextTokens == 0 && usage.PromptTokensDetails.AudioTokens == 0 {
1069+
usage.PromptTokensDetails.TextTokens = usage.PromptTokens
1070+
}
1071+
1072+
return usage
1073+
}
1074+
10351075
func responseGeminiChat2OpenAI(c *gin.Context, response *dto.GeminiChatResponse) *dto.OpenAITextResponse {
10361076
fullTextResponse := dto.OpenAITextResponse{
10371077
Id: helper.GetResponseID(c),
@@ -1272,18 +1312,8 @@ func geminiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http
12721312

12731313
// 更新使用量统计
12741314
if geminiResponse.UsageMetadata.TotalTokenCount != 0 {
1275-
usage.PromptTokens = geminiResponse.UsageMetadata.PromptTokenCount
1276-
usage.CompletionTokens = geminiResponse.UsageMetadata.CandidatesTokenCount + geminiResponse.UsageMetadata.ThoughtsTokenCount
1277-
usage.CompletionTokenDetails.ReasoningTokens = geminiResponse.UsageMetadata.ThoughtsTokenCount
1278-
usage.TotalTokens = geminiResponse.UsageMetadata.TotalTokenCount
1279-
usage.PromptTokensDetails.CachedTokens = geminiResponse.UsageMetadata.CachedContentTokenCount
1280-
for _, detail := range geminiResponse.UsageMetadata.PromptTokensDetails {
1281-
if detail.Modality == "AUDIO" {
1282-
usage.PromptTokensDetails.AudioTokens = detail.TokenCount
1283-
} else if detail.Modality == "TEXT" {
1284-
usage.PromptTokensDetails.TextTokens = detail.TokenCount
1285-
}
1286-
}
1315+
mappedUsage := buildUsageFromGeminiMetadata(geminiResponse.UsageMetadata, info.GetEstimatePromptTokens())
1316+
*usage = mappedUsage
12871317
}
12881318

12891319
return callback(data, &geminiResponse)
@@ -1295,11 +1325,6 @@ func geminiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http
12951325
}
12961326
}
12971327

1298-
usage.PromptTokensDetails.TextTokens = usage.PromptTokens
1299-
if usage.TotalTokens > 0 {
1300-
usage.CompletionTokens = usage.TotalTokens - usage.PromptTokens
1301-
}
1302-
13031328
if usage.CompletionTokens <= 0 {
13041329
if info.ReceivedResponseCount > 0 {
13051330
usage = service.ResponseText2Usage(c, responseText.String(), info.UpstreamModelName, info.GetEstimatePromptTokens())
@@ -1416,21 +1441,7 @@ func GeminiChatHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.R
14161441
return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
14171442
}
14181443
if len(geminiResponse.Candidates) == 0 {
1419-
usage := dto.Usage{
1420-
PromptTokens: geminiResponse.UsageMetadata.PromptTokenCount,
1421-
}
1422-
usage.CompletionTokenDetails.ReasoningTokens = geminiResponse.UsageMetadata.ThoughtsTokenCount
1423-
usage.PromptTokensDetails.CachedTokens = geminiResponse.UsageMetadata.CachedContentTokenCount
1424-
for _, detail := range geminiResponse.UsageMetadata.PromptTokensDetails {
1425-
if detail.Modality == "AUDIO" {
1426-
usage.PromptTokensDetails.AudioTokens = detail.TokenCount
1427-
} else if detail.Modality == "TEXT" {
1428-
usage.PromptTokensDetails.TextTokens = detail.TokenCount
1429-
}
1430-
}
1431-
if usage.PromptTokens <= 0 {
1432-
usage.PromptTokens = info.GetEstimatePromptTokens()
1433-
}
1444+
usage := buildUsageFromGeminiMetadata(geminiResponse.UsageMetadata, info.GetEstimatePromptTokens())
14341445

14351446
var newAPIError *types.NewAPIError
14361447
if geminiResponse.PromptFeedback != nil && geminiResponse.PromptFeedback.BlockReason != nil {
@@ -1466,23 +1477,7 @@ func GeminiChatHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.R
14661477
}
14671478
fullTextResponse := responseGeminiChat2OpenAI(c, &geminiResponse)
14681479
fullTextResponse.Model = info.UpstreamModelName
1469-
usage := dto.Usage{
1470-
PromptTokens: geminiResponse.UsageMetadata.PromptTokenCount,
1471-
CompletionTokens: geminiResponse.UsageMetadata.CandidatesTokenCount,
1472-
TotalTokens: geminiResponse.UsageMetadata.TotalTokenCount,
1473-
}
1474-
1475-
usage.CompletionTokenDetails.ReasoningTokens = geminiResponse.UsageMetadata.ThoughtsTokenCount
1476-
usage.PromptTokensDetails.CachedTokens = geminiResponse.UsageMetadata.CachedContentTokenCount
1477-
usage.CompletionTokens = usage.TotalTokens - usage.PromptTokens
1478-
1479-
for _, detail := range geminiResponse.UsageMetadata.PromptTokensDetails {
1480-
if detail.Modality == "AUDIO" {
1481-
usage.PromptTokensDetails.AudioTokens = detail.TokenCount
1482-
} else if detail.Modality == "TEXT" {
1483-
usage.PromptTokensDetails.TextTokens = detail.TokenCount
1484-
}
1485-
}
1480+
usage := buildUsageFromGeminiMetadata(geminiResponse.UsageMetadata, info.GetEstimatePromptTokens())
14861481

14871482
fullTextResponse.Usage = usage
14881483

0 commit comments

Comments
 (0)