From 89709dfa74f33311bb23330948f09c2abda6ef5a Mon Sep 17 00:00:00 2001 From: CaIon Date: Sun, 23 Nov 2025 23:50:04 +0800 Subject: [PATCH] feat: Set ContextKeyLocalCountTokens in NativeGeminiEmbeddingHandler for token tracking --- relay/channel/gemini/relay-gemini-native.go | 3 +++ 1 file changed, 3 insertions(+) diff --git a/relay/channel/gemini/relay-gemini-native.go b/relay/channel/gemini/relay-gemini-native.go index 98d882ec..c3f7aa0a 100644 --- a/relay/channel/gemini/relay-gemini-native.go +++ b/relay/channel/gemini/relay-gemini-native.go @@ -5,6 +5,7 @@ import ( "net/http" "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/dto" "github.com/QuantumNous/new-api/logger" relaycommon "github.com/QuantumNous/new-api/relay/common" @@ -74,6 +75,8 @@ func NativeGeminiEmbeddingHandler(c *gin.Context, resp *http.Response, info *rel TotalTokens: info.PromptTokens, } + common.SetContextKey(c, constant.ContextKeyLocalCountTokens, true) + if info.IsGeminiBatchEmbedding { var geminiResponse dto.GeminiBatchEmbeddingResponse err = common.Unmarshal(responseBody, &geminiResponse)