Răsfoiți Sursa

feat: Set ContextKeyLocalCountTokens in NativeGeminiEmbeddingHandler for token tracking

CaIon 2 luni în urmă
părinte
comite
8f039b3a53
1 a modificat fișierele cu 3 adăugiri și 0 ștergeri
  1. 3 0
      relay/channel/gemini/relay-gemini-native.go

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