From b64480b7500f2a6df6c75aa64eda5b142b01f70a Mon Sep 17 00:00:00 2001
From: CaIon <1808837298@qq.com>
Date: Tue, 29 Apr 2025 16:21:20 +0800
Subject: [PATCH] fix: gemini thinking tokens count #1014

---
 dto/openai_request.go                | 1 +
 relay/channel/ali/constants.go       | 7 ++++++-
 relay/channel/gemini/relay-gemini.go | 5 +++--
 3 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/dto/openai_request.go b/dto/openai_request.go
index 28aac20f..652d8cce 100644
--- a/dto/openai_request.go
+++ b/dto/openai_request.go
@@ -51,6 +51,7 @@ type GeneralOpenAIRequest struct {
 	Dimensions       int               `json:"dimensions,omitempty"`
 	Modalities       any               `json:"modalities,omitempty"`
 	Audio            any               `json:"audio,omitempty"`
+	EnableThinking   any               `json:"enable_thinking,omitempty"` // ali
 	ExtraBody        any               `json:"extra_body,omitempty"`
 }
 
diff --git a/relay/channel/ali/constants.go b/relay/channel/ali/constants.go
index 6f6658c5..46de5e40 100644
--- a/relay/channel/ali/constants.go
+++ b/relay/channel/ali/constants.go
@@ -1,7 +1,12 @@
 package ali
 
 var ModelList = []string{
-	"qwen-turbo", "qwen-plus", "qwen-max", "qwen-max-longcontext",
+	"qwen-turbo",
+	"qwen-plus",
+	"qwen-max",
+	"qwen-max-longcontext",
+	"qwq-32b",
+	"qwen3-235b-a22b",
 	"text-embedding-v1",
 }
 
diff --git a/relay/channel/gemini/relay-gemini.go b/relay/channel/gemini/relay-gemini.go
index c4f4af31..dbe65528 100644
--- a/relay/channel/gemini/relay-gemini.go
+++ b/relay/channel/gemini/relay-gemini.go
@@ -670,6 +670,7 @@ func GeminiChatStreamHandler(c *gin.Context, resp *http.Response, info *relaycom
 			usage.PromptTokens = geminiResponse.UsageMetadata.PromptTokenCount
 			usage.CompletionTokens = geminiResponse.UsageMetadata.CandidatesTokenCount
 			usage.CompletionTokenDetails.ReasoningTokens = geminiResponse.UsageMetadata.ThoughtsTokenCount
+			usage.TotalTokens = geminiResponse.UsageMetadata.TotalTokenCount
 		}
 		err = helper.ObjectData(c, response)
 		if err != nil {
@@ -690,9 +691,8 @@ func GeminiChatStreamHandler(c *gin.Context, resp *http.Response, info *relaycom
 		}
 	}
 
-	usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens
 	usage.PromptTokensDetails.TextTokens = usage.PromptTokens
-	//usage.CompletionTokenDetails.TextTokens = usage.CompletionTokens
+	usage.CompletionTokens = usage.TotalTokens - usage.PromptTokens
 
 	if info.ShouldIncludeUsage {
 		response = helper.GenerateFinalUsageResponse(id, createAt, info.UpstreamModelName, *usage)
@@ -740,6 +740,7 @@ func GeminiChatHandler(c *gin.Context, resp *http.Response, info *relaycommon.Re
 	}
 
 	usage.CompletionTokenDetails.ReasoningTokens = geminiResponse.UsageMetadata.ThoughtsTokenCount
+	usage.CompletionTokens = usage.TotalTokens - usage.PromptTokens
 
 	fullTextResponse.Usage = usage
 	jsonResponse, err := json.Marshal(fullTextResponse)