github/new-api
1
0
Fork 0
mirror of https://github.com/QuantumNous/new-api.git synced 2026-04-28 00:18:37 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: Set ContextKeyLocalCountTokens in NativeGeminiEmbeddingHandler for token tracking

Browse Source
This commit is contained in:
CaIon
2025-11-23 23:50:04 +08:00
parent c939686509
commit 8f039b3a53
1 changed files with 3 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
relay/channel/gemini/relay-gemini-native.go
View File

@@ -5,6 +5,7 @@ import (
"net/http" "net/http"
"github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/dto" "github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/logger" "github.com/QuantumNous/new-api/logger"
relaycommon "github.com/QuantumNous/new-api/relay/common" 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, TotalTokens: info.PromptTokens,
} }
common.SetContextKey(c, constant.ContextKeyLocalCountTokens, true)
if info.IsGeminiBatchEmbedding { if info.IsGeminiBatchEmbedding {
var geminiResponse dto.GeminiBatchEmbeddingResponse var geminiResponse dto.GeminiBatchEmbeddingResponse
err = common.Unmarshal(responseBody, &geminiResponse) err = common.Unmarshal(responseBody, &geminiResponse)