Просмотр исходного кода

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

creamlike1024 8 месяцев назад
Родитель
Сommit
a56d9ea98b
1 измененных файлов с 15 добавлено и 0 удалено
  1. 15 0
      relay/channel/gemini/relay-gemini-native.go

+ 15 - 0
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