fix: gemini 原生格式流模式中断请求未计费

This commit is contained in:
creamlike1024
2025-06-20 23:01:10 +08:00
parent f5e80af0b3
commit a56d9ea98b

View File

@@ -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