Skip to content

phsl / new-api

Go to a project
Commit 89709dfa by CaIon
Options

feat: Set ContextKeyLocalCountTokens in NativeGeminiEmbeddingHandler for token tracking

parent ed3858ad
Showing with 3 additions and 0 deletions
relay/channel/gemini/relay-gemini-native.go
...@@ -5,6 +5,7 @@ import ( ...@@ -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 ...@@ -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)
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or sign in to comment