🔖 chore: update gemini model and token calculation.

This commit is contained in:
MartialBE 2024-05-15 03:49:22 +08:00
parent e9d3b4654f
commit 46aea4731b
No known key found for this signature in database
GPG Key ID: 27C0267EC84B0A5C
3 changed files with 69 additions and 10 deletions

View File

@ -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},

View File

@ -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,
}
}

View File

@ -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"`