Browse Source

fix: gemini thinking tokens count #1014

CaIon 8 months ago
parent
commit
b64480b750
3 changed files with 10 additions and 3 deletions
  1. 1 0
      dto/openai_request.go
  2. 6 1
      relay/channel/ali/constants.go
  3. 3 2
      relay/channel/gemini/relay-gemini.go

+ 1 - 0
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"`
 }
 

+ 6 - 1
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",
 }
 

+ 3 - 2
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)