diff --git a/dto/gemini.go b/dto/gemini.go index 0fd74c639d2c..92ed9070396a 100644 --- a/dto/gemini.go +++ b/dto/gemini.go @@ -457,8 +457,10 @@ type GeminiUsageMetadata struct { CandidatesTokenCount int `json:"candidatesTokenCount"` TotalTokenCount int `json:"totalTokenCount"` ThoughtsTokenCount int `json:"thoughtsTokenCount"` + ToolUsePromptTokenCount int `json:"toolUsePromptTokenCount"` CachedContentTokenCount int `json:"cachedContentTokenCount"` PromptTokensDetails []GeminiPromptTokensDetails `json:"promptTokensDetails"` + ToolUsePromptDetails []GeminiPromptTokensDetails `json:"toolUsePromptTokensDetails"` } type GeminiPromptTokensDetails struct { diff --git a/relay/channel/gemini/relay-gemini-native.go b/relay/channel/gemini/relay-gemini-native.go index 39485b16f143..44f8fbec6dcb 100644 --- a/relay/channel/gemini/relay-gemini-native.go +++ b/relay/channel/gemini/relay-gemini-native.go @@ -41,23 +41,7 @@ func GeminiTextGenerationHandler(c *gin.Context, info *relaycommon.RelayInfo, re common.SetContextKey(c, constant.ContextKeyAdminRejectReason, fmt.Sprintf("gemini_block_reason=%s", *geminiResponse.PromptFeedback.BlockReason)) } - // 计算使用量（基于 UsageMetadata） - usage := dto.Usage{ - PromptTokens: geminiResponse.UsageMetadata.PromptTokenCount, - CompletionTokens: geminiResponse.UsageMetadata.CandidatesTokenCount + geminiResponse.UsageMetadata.ThoughtsTokenCount, - TotalTokens: geminiResponse.UsageMetadata.TotalTokenCount, - } - - usage.CompletionTokenDetails.ReasoningTokens = geminiResponse.UsageMetadata.ThoughtsTokenCount - usage.PromptTokensDetails.CachedTokens = geminiResponse.UsageMetadata.CachedContentTokenCount - - 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 - } - } + usage := buildUsageFromGeminiMetadata(geminiResponse.UsageMetadata, info.GetEstimatePromptTokens()) service.IOCopyBytesGracefully(c, resp, responseBody) diff --git a/relay/channel/gemini/relay-gemini.go b/relay/channel/gemini/relay-gemini.go index b10ec06c7b16..b7d19dc77cfe 100644 --- a/relay/channel/gemini/relay-gemini.go +++ b/relay/channel/gemini/relay-gemini.go @@ -1270,20 +1270,9 @@ func geminiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http } } - // 更新使用量统计 - if geminiResponse.UsageMetadata.TotalTokenCount != 0 { - usage.PromptTokens = geminiResponse.UsageMetadata.PromptTokenCount - usage.CompletionTokens = geminiResponse.UsageMetadata.CandidatesTokenCount + geminiResponse.UsageMetadata.ThoughtsTokenCount - usage.CompletionTokenDetails.ReasoningTokens = geminiResponse.UsageMetadata.ThoughtsTokenCount - usage.TotalTokens = geminiResponse.UsageMetadata.TotalTokenCount - usage.PromptTokensDetails.CachedTokens = geminiResponse.UsageMetadata.CachedContentTokenCount - 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 - } - } + if hasGeminiUsageMetadata(geminiResponse.UsageMetadata) { + latestUsage := buildUsageFromGeminiMetadata(geminiResponse.UsageMetadata, info.GetEstimatePromptTokens()) + *usage = latestUsage } return callback(data, &geminiResponse) @@ -1295,9 +1284,11 @@ func geminiStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http } } - usage.PromptTokensDetails.TextTokens = usage.PromptTokens - if usage.TotalTokens > 0 { - usage.CompletionTokens = usage.TotalTokens - usage.PromptTokens + if usage.PromptTokensDetails.TextTokens == 0 && usage.PromptTokens > 0 { + usage.PromptTokensDetails.TextTokens = usage.PromptTokens + } + if usage.TotalTokens <= 0 && (usage.PromptTokens > 0 || usage.CompletionTokens > 0) { + usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens } if usage.CompletionTokens <= 0 { @@ -1416,21 +1407,7 @@ func GeminiChatHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.R return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError) } if len(geminiResponse.Candidates) == 0 { - usage := dto.Usage{ - PromptTokens: geminiResponse.UsageMetadata.PromptTokenCount, - } - usage.CompletionTokenDetails.ReasoningTokens = geminiResponse.UsageMetadata.ThoughtsTokenCount - usage.PromptTokensDetails.CachedTokens = geminiResponse.UsageMetadata.CachedContentTokenCount - 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 - } - } - if usage.PromptTokens <= 0 { - usage.PromptTokens = info.GetEstimatePromptTokens() - } + usage := buildUsageFromGeminiMetadata(geminiResponse.UsageMetadata, info.GetEstimatePromptTokens()) var newAPIError *types.NewAPIError if geminiResponse.PromptFeedback != nil && geminiResponse.PromptFeedback.BlockReason != nil { @@ -1466,23 +1443,7 @@ 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, - TotalTokens: geminiResponse.UsageMetadata.TotalTokenCount, - } - - usage.CompletionTokenDetails.ReasoningTokens = geminiResponse.UsageMetadata.ThoughtsTokenCount - usage.PromptTokensDetails.CachedTokens = geminiResponse.UsageMetadata.CachedContentTokenCount - usage.CompletionTokens = usage.TotalTokens - usage.PromptTokens - - 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 - } - } + usage := buildUsageFromGeminiMetadata(geminiResponse.UsageMetadata, info.GetEstimatePromptTokens()) fullTextResponse.Usage = usage diff --git a/relay/channel/gemini/usage.go b/relay/channel/gemini/usage.go new file mode 100644 index 000000000000..66ad3fc106e0 --- /dev/null +++ b/relay/channel/gemini/usage.go @@ -0,0 +1,60 @@ +package gemini + +import "github.com/QuantumNous/new-api/dto" + +func hasGeminiUsageMetadata(meta dto.GeminiUsageMetadata) bool { + return meta.PromptTokenCount > 0 || + meta.CandidatesTokenCount > 0 || + meta.ThoughtsTokenCount > 0 || + meta.TotalTokenCount > 0 || + meta.ToolUsePromptTokenCount > 0 || + meta.CachedContentTokenCount > 0 || + len(meta.PromptTokensDetails) > 0 || + len(meta.ToolUsePromptDetails) > 0 +} + +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 + } + + return usage +} diff --git a/relay/channel/gemini/usage_test.go b/relay/channel/gemini/usage_test.go new file mode 100644 index 000000000000..eb1422f1a910 --- /dev/null +++ b/relay/channel/gemini/usage_test.go @@ -0,0 +1,74 @@ +package gemini + +import ( + "testing" + + "github.com/QuantumNous/new-api/dto" + "github.com/stretchr/testify/require" +) + +func TestBuildUsageFromGeminiMetadata_ExcludesToolUseFromCompletion(t *testing.T) { + meta := dto.GeminiUsageMetadata{ + PromptTokenCount: 151, + CandidatesTokenCount: 1089, + ThoughtsTokenCount: 1120, + TotalTokenCount: 20689, + ToolUsePromptTokenCount: 18329, + CachedContentTokenCount: 17, + PromptTokensDetails: []dto.GeminiPromptTokensDetails{ + {Modality: "TEXT", TokenCount: 151}, + }, + ToolUsePromptDetails: []dto.GeminiPromptTokensDetails{ + {Modality: "TEXT", TokenCount: 18329}, + }, + } + + 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) +} + +func TestBuildUsageFromGeminiMetadata_FallsBackToTotalPromptAndToolUse(t *testing.T) { + meta := dto.GeminiUsageMetadata{ + PromptTokenCount: 100, + TotalTokenCount: 1000, + ToolUsePromptTokenCount: 700, + } + + usage := buildUsageFromGeminiMetadata(meta, 0) + + require.Equal(t, 800, usage.PromptTokens) + require.Equal(t, 200, usage.CompletionTokens) + require.Equal(t, 1000, usage.TotalTokens) +} + +func TestBuildUsageFromGeminiMetadata_NegativeCompletionClampedToZero(t *testing.T) { + meta := dto.GeminiUsageMetadata{ + PromptTokenCount: 300, + TotalTokenCount: 200, + ToolUsePromptTokenCount: 50, + } + + usage := buildUsageFromGeminiMetadata(meta, 0) + + require.Equal(t, 350, usage.PromptTokens) + require.Equal(t, 0, usage.CompletionTokens) + require.Equal(t, 200, usage.TotalTokens) +} + +func TestBuildUsageFromGeminiMetadata_UsesEstimatedPromptWhenMissing(t *testing.T) { + meta := dto.GeminiUsageMetadata{ + CandidatesTokenCount: 20, + } + + usage := buildUsageFromGeminiMetadata(meta, 123) + + require.Equal(t, 123, usage.PromptTokens) + require.Equal(t, 20, usage.CompletionTokens) + require.Equal(t, 143, usage.TotalTokens) +}