Skip to content

phsl / new-api

Go to a project
Commit a004db93 by RedwindA
Options

feat(Gemini): enhance budget clamping logic for Gemini models

parent 36a0d4d7
Showing with 50 additions and 59 deletions
relay/channel/gemini/relay-gemini.go
...@@ -42,8 +42,42 @@ const ( ...@@ -42,8 +42,42 @@ const (
pro25MinBudget = 128 pro25MinBudget = 128
pro25MaxBudget = 32768 pro25MaxBudget = 32768
flash25MaxBudget = 24576 flash25MaxBudget = 24576
flash25LiteMinBudget = 512
flash25LiteMaxBudget = 24576
) )
// clampThinkingBudget 根据模型名称将预算限制在允许的范围内
func clampThinkingBudget(modelName string, budget int) int {
isNew25Pro := strings.HasPrefix(modelName, "gemini-2.5-pro") &&
!strings.HasPrefix(modelName, "gemini-2.5-pro-preview-05-06") &&
!strings.HasPrefix(modelName, "gemini-2.5-pro-preview-03-25")
is25FlashLite := strings.HasPrefix(modelName, "gemini-2.5-flash-lite")
if is25FlashLite {
if budget < flash25LiteMinBudget {
return flash25LiteMinBudget
}
if budget > flash25LiteMaxBudget {
return flash25LiteMaxBudget
}
} else if isNew25Pro {
if budget < pro25MinBudget {
return pro25MinBudget
}
if budget > pro25MaxBudget {
return pro25MaxBudget
}
} else { // 其他模型
if budget < 0 {
return 0
}
if budget > flash25MaxBudget {
return flash25MaxBudget
}
}
return budget
}
// Setting safety to the lowest possible values since Gemini is already powerless enough // Setting safety to the lowest possible values since Gemini is already powerless enough
func CovertGemini2OpenAI(textRequest dto.GeneralOpenAIRequest, info *relaycommon.RelayInfo) (*GeminiChatRequest, error) { func CovertGemini2OpenAI(textRequest dto.GeneralOpenAIRequest, info *relaycommon.RelayInfo) (*GeminiChatRequest, error) {
...@@ -65,49 +99,31 @@ func CovertGemini2OpenAI(textRequest dto.GeneralOpenAIRequest, info *relaycommon ...@@ -65,49 +99,31 @@ func CovertGemini2OpenAI(textRequest dto.GeneralOpenAIRequest, info *relaycommon
} }
if model_setting.GetGeminiSettings().ThinkingAdapterEnabled { if model_setting.GetGeminiSettings().ThinkingAdapterEnabled {
// 新增逻辑：处理 -thinking-<budget> 格式 modelName := info.OriginModelName
if strings.Contains(info.OriginModelName, "-thinking-") { isNew25Pro := strings.HasPrefix(modelName, "gemini-2.5-pro") &&
parts := strings.SplitN(info.OriginModelName, "-thinking-", 2) !strings.HasPrefix(modelName, "gemini-2.5-pro-preview-05-06") &&
!strings.HasPrefix(modelName, "gemini-2.5-pro-preview-03-25")
is25FlashLite := strings.HasPrefix(modelName, "gemini-2.5-flash-lite")
if strings.Contains(modelName, "-thinking-") {
parts := strings.SplitN(modelName, "-thinking-", 2)
if len(parts) == 2 && parts[1] != "" { if len(parts) == 2 && parts[1] != "" {
if budgetTokens, err := strconv.Atoi(parts[1]); err == nil { if budgetTokens, err := strconv.Atoi(parts[1]); err == nil {
// 从模型名称成功解析预算 clampedBudget := clampThinkingBudget(modelName, budgetTokens)
isNew25Pro := strings.HasPrefix(info.OriginModelName, "gemini-2.5-pro") &&
!strings.HasPrefix(info.OriginModelName, "gemini-2.5-pro-preview-05-06") &&
!strings.HasPrefix(info.OriginModelName, "gemini-2.5-pro-preview-03-25")
if isNew25Pro {
// 新的2.5pro模型：ThinkingBudget范围为128-32768
if budgetTokens < pro25MinBudget {
budgetTokens = pro25MinBudget
} else if budgetTokens > pro25MaxBudget {
budgetTokens = pro25MaxBudget
}
} else {
// 其他模型：ThinkingBudget范围为0-24576
if budgetTokens < 0 {
budgetTokens = 0
} else if budgetTokens > flash25MaxBudget {
budgetTokens = flash25MaxBudget
}
}
geminiRequest.GenerationConfig.ThinkingConfig = &GeminiThinkingConfig{ geminiRequest.GenerationConfig.ThinkingConfig = &GeminiThinkingConfig{
ThinkingBudget: common.GetPointer(budgetTokens), ThinkingBudget: common.GetPointer(clampedBudget),
IncludeThoughts: true, IncludeThoughts: true,
} }
} }
// 如果解析失败，则不设置ThinkingConfig，静默处理
} }
} else if strings.HasSuffix(info.OriginModelName, "-thinking") { // 保留旧逻辑以兼容 } else if strings.HasSuffix(modelName, "-thinking") {
// 硬编码不支持 ThinkingBudget 的旧模型
unsupportedModels := []string{ unsupportedModels := []string{
"gemini-2.5-pro-preview-05-06", "gemini-2.5-pro-preview-05-06",
"gemini-2.5-pro-preview-03-25", "gemini-2.5-pro-preview-03-25",
} }
isUnsupported := false isUnsupported := false
for _, unsupportedModel := range unsupportedModels { for _, unsupportedModel := range unsupportedModels {
if strings.HasPrefix(info.OriginModelName, unsupportedModel) { if strings.HasPrefix(modelName, unsupportedModel) {
isUnsupported = true isUnsupported = true
break break
} }
...@@ -119,39 +135,14 @@ func CovertGemini2OpenAI(textRequest dto.GeneralOpenAIRequest, info *relaycommon ...@@ -119,39 +135,14 @@ func CovertGemini2OpenAI(textRequest dto.GeneralOpenAIRequest, info *relaycommon
} }
} else { } else {
budgetTokens := model_setting.GetGeminiSettings().ThinkingAdapterBudgetTokensPercentage * float64(geminiRequest.GenerationConfig.MaxOutputTokens) budgetTokens := model_setting.GetGeminiSettings().ThinkingAdapterBudgetTokensPercentage * float64(geminiRequest.GenerationConfig.MaxOutputTokens)
clampedBudget := clampThinkingBudget(modelName, int(budgetTokens))
// 检查是否为新的2.5pro模型（支持ThinkingBudget但有特殊范围）
isNew25Pro := strings.HasPrefix(info.OriginModelName, "gemini-2.5-pro") &&
!strings.HasPrefix(info.OriginModelName, "gemini-2.5-pro-preview-05-06") &&
!strings.HasPrefix(info.OriginModelName, "gemini-2.5-pro-preview-03-25")
if isNew25Pro {
// 新的2.5pro模型：ThinkingBudget范围为128-32768
if budgetTokens == 0 || budgetTokens < 128 {
budgetTokens = 128
} else if budgetTokens > 32768 {
budgetTokens = 32768
}
} else {
// 其他模型：ThinkingBudget范围为0-24576
if budgetTokens == 0 || budgetTokens > 24576 {
budgetTokens = 24576
}
}
geminiRequest.GenerationConfig.ThinkingConfig = &GeminiThinkingConfig{ geminiRequest.GenerationConfig.ThinkingConfig = &GeminiThinkingConfig{
ThinkingBudget: common.GetPointer(int(budgetTokens)), ThinkingBudget: common.GetPointer(clampedBudget),
IncludeThoughts: true, IncludeThoughts: true,
} }
} }
} else if strings.HasSuffix(info.OriginModelName, "-nothinking") { } else if strings.HasSuffix(modelName, "-nothinking") {
// 检查是否为新的2.5pro模型（不支持-nothinking，因为最低值只能为128） if !isNew25Pro && !is25FlashLite {
isNew25Pro := strings.HasPrefix(info.OriginModelName, "gemini-2.5-pro") &&
!strings.HasPrefix(info.OriginModelName, "gemini-2.5-pro-preview-05-06") &&
!strings.HasPrefix(info.OriginModelName, "gemini-2.5-pro-preview-03-25")
if !isNew25Pro {
// 只有非新2.5pro模型才支持-nothinking
geminiRequest.GenerationConfig.ThinkingConfig = &GeminiThinkingConfig{ geminiRequest.GenerationConfig.ThinkingConfig = &GeminiThinkingConfig{
ThinkingBudget: common.GetPointer(0), ThinkingBudget: common.GetPointer(0),
} }
......
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