diff --git a/relay/channel/gemini/relay-gemini-native.go b/relay/channel/gemini/relay-gemini-native.go index cf7920dc..3a017a11 100644 --- a/relay/channel/gemini/relay-gemini-native.go +++ b/relay/channel/gemini/relay-gemini-native.go @@ -75,6 +75,9 @@ func GeminiTextGenerationStreamHandler(c *gin.Context, resp *http.Response, info helper.SetEventStreamHeaders(c) + // 本地统计的completion tokens + localCompletionTokens := 0 + helper.StreamScannerHandler(c, resp, info, func(data string) bool { var geminiResponse GeminiChatResponse err := common.DecodeJsonStr(data, &geminiResponse) @@ -89,6 +92,12 @@ func GeminiTextGenerationStreamHandler(c *gin.Context, resp *http.Response, info if part.InlineData != nil && part.InlineData.MimeType != "" { imageCount++ } + // 本地统计completion tokens + textTokens, err := service.CountTextToken(part.Text, info.UpstreamModelName) + if err != nil { + common.LogError(c, "error counting text token: "+err.Error()) + } + localCompletionTokens += textTokens } } @@ -122,6 +131,12 @@ func GeminiTextGenerationStreamHandler(c *gin.Context, resp *http.Response, info } } + // 如果usage.CompletionTokens为0,则使用本地统计的completion tokens + if usage.CompletionTokens == 0 { + usage.CompletionTokens = localCompletionTokens + usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens + } + // 计算最终使用量 // usage.CompletionTokens = usage.TotalTokens - usage.PromptTokens