diff --git a/model/price.go b/model/price.go index 4f509f5b..841e3515 100644 --- a/model/price.go +++ b/model/price.go @@ -207,9 +207,11 @@ func GetDefaultPrice() []*Price { "gemini-pro-vision": {[]float64{0.25, 0.75}, common.ChannelTypeGemini}, "gemini-1.0-pro": {[]float64{0.25, 0.75}, common.ChannelTypeGemini}, // $7 / 1 million tokens $21 / 1 million tokens - // 0.007$ / 1k tokens 0.021$ / 1k tokens - "gemini-1.5-pro": {[]float64{3.5, 10.5}, common.ChannelTypeGemini}, - "gemini-ultra": {[]float64{1, 1}, common.ChannelTypeGemini}, + "gemini-1.5-pro": {[]float64{1.75, 5.25}, common.ChannelTypeGemini}, + "gemini-1.5-pro-latest": {[]float64{1.75, 5.25}, common.ChannelTypeGemini}, + "gemini-1.5-flash": {[]float64{0.175, 0.265}, common.ChannelTypeGemini}, + "gemini-1.5-flash-latest": {[]float64{0.175, 0.265}, common.ChannelTypeGemini}, + "gemini-ultra": {[]float64{1, 1}, common.ChannelTypeGemini}, // ¥0.005 / 1k tokens "glm-3-turbo": {[]float64{0.3572, 0.3572}, common.ChannelTypeZhipu}, diff --git a/providers/gemini/chat.go b/providers/gemini/chat.go index 310160ba..d66299b7 100644 --- a/providers/gemini/chat.go +++ b/providers/gemini/chat.go @@ -173,10 +173,7 @@ func (p *GeminiProvider) convertToChatOpenai(response *GeminiChatResponse, reque openaiResponse.Choices = append(openaiResponse.Choices, choice) } - completionTokens := common.CountTokenText(response.GetResponseText(), response.Model) - - p.Usage.CompletionTokens = completionTokens - p.Usage.TotalTokens = p.Usage.PromptTokens + completionTokens + *p.Usage = convertOpenAIUsage(request.Model, response.UsageMetadata) openaiResponse.Usage = p.Usage return @@ -270,6 +267,60 @@ func (h *geminiStreamHandler) convertToOpenaiStream(geminiResponse *GeminiChatRe dataChan <- string(responseBody) } - h.Usage.CompletionTokens += common.CountTokenText(geminiResponse.GetResponseText(), h.Request.Model) - h.Usage.TotalTokens = h.Usage.PromptTokens + h.Usage.CompletionTokens + if geminiResponse.UsageMetadata != nil { + *h.Usage = convertOpenAIUsage(h.Request.Model, geminiResponse.UsageMetadata) + + } +} + +const tokenThreshold = 1000000 + +var modelAdjustRatios = map[string]int{ + "gemini-1.5-pro": 2, + "gemini-1.5-flash": 2, +} + +func adjustTokenCounts(modelName string, usage *GeminiUsageMetadata) { + if usage.PromptTokenCount <= tokenThreshold && usage.CandidatesTokenCount <= tokenThreshold { + return + } + + currentRatio := 1 + for model, r := range modelAdjustRatios { + if strings.HasPrefix(modelName, model) { + currentRatio = r + break + } + } + + if currentRatio == 1 { + return + } + + adjustTokenCount := func(count int) int { + if count > tokenThreshold { + return tokenThreshold + (count-tokenThreshold)*currentRatio + } + return count + } + + if usage.PromptTokenCount > tokenThreshold { + usage.PromptTokenCount = adjustTokenCount(usage.PromptTokenCount) + } + + if usage.CandidatesTokenCount > tokenThreshold { + usage.CandidatesTokenCount = adjustTokenCount(usage.CandidatesTokenCount) + } + + usage.TotalTokenCount = usage.PromptTokenCount + usage.CandidatesTokenCount +} + +func convertOpenAIUsage(modelName string, geminiUsage *GeminiUsageMetadata) types.Usage { + adjustTokenCounts(modelName, geminiUsage) + + return types.Usage{ + PromptTokens: geminiUsage.PromptTokenCount, + CompletionTokens: geminiUsage.CandidatesTokenCount, + TotalTokens: geminiUsage.TotalTokenCount, + } } diff --git a/providers/gemini/type.go b/providers/gemini/type.go index 46933cb8..d653dccf 100644 --- a/providers/gemini/type.go +++ b/providers/gemini/type.go @@ -94,11 +94,17 @@ type GeminiErrorResponse struct { type GeminiChatResponse struct { Candidates []GeminiChatCandidate `json:"candidates"` PromptFeedback GeminiChatPromptFeedback `json:"promptFeedback"` - Usage *types.Usage `json:"usage,omitempty"` + UsageMetadata *GeminiUsageMetadata `json:"usageMetadata,omitempty"` Model string `json:"model,omitempty"` GeminiErrorResponse } +type GeminiUsageMetadata struct { + PromptTokenCount int `json:"promptTokenCount"` + CandidatesTokenCount int `json:"candidatesTokenCount"` + TotalTokenCount int `json:"totalTokenCount"` +} + type GeminiChatCandidate struct { Content GeminiChatContent `json:"content"` FinishReason string `json:"finishReason"`