Skip to content

feat: improve gemini thinking budget adaption #1247

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Calcium-Ion merged 1 commit into from
Jun 17, 2025
Merged

feat: improve gemini thinking budget adaption #1247

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
115 changes: 53 additions & 62 deletions relay/channel/gemini/relay-gemini.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,11 +39,45 @@ var geminiSupportedMimeTypes = map[string]bool{

// Gemini 允许的思考预算范围
const (
pro25MinBudget = 128
pro25MaxBudget = 32768
flash25MaxBudget = 24576
pro25MinBudget = 128
pro25MaxBudget = 32768
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
}
Comment on lines +49 to +79

@coderabbitai coderabbitai Bot Jun 17, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ Refactor suggestion

clampThinkingBudget mixes concerns & repeats model-detection logic

  1. The isNew25Pro / is25FlashLite checks will soon live in at least three places (here, in CovertGemini2OpenAI, and very likely in request validation). A single helper such as detectGeminiVariant(modelName) (variant enum) would avoid drift.

  2. For negative budgets the function silently upgrades to the model’s minimum (128 / 512). That might be desirable, but it is implicit. At minimum add a short doc-comment or, better, reject negatives and let the caller decide:

-	if budget < pro25MinBudget {
-		return pro25MinBudget
+	if budget <= 0 {
+		return pro25MinBudget // or: return 0, signalling “invalid”
  1. The max for “other models” is flash25MaxBudget; consider renaming that constant or adding a comment so the cap doesn’t look arbitrary.

Refactoring sketch (compressed):

type geminiVariant int
const (
	varUnknown geminiVariant = iota
	var25Pro
	var25FlashLite
	varOther
)

func detectVariant(m string) geminiVariant { … }

func clampBudget(v geminiVariant, b int) int { … }

This removes string-prefix checks from every caller.

🤖 Prompt for AI Agents 
In relay/channel/gemini/relay-gemini.go around lines 49 to 79, the
clampThinkingBudget function mixes model detection logic with budget clamping,
causing repeated string-prefix checks and implicit handling of negative budgets.
Refactor by extracting model detection into a separate helper function
detectGeminiVariant(modelName) returning a variant enum, then rewrite
clampThinkingBudget to accept this variant and budget, explicitly rejecting
negative budgets or documenting behavior clearly. Also, rename or comment the
flash25MaxBudget constant used as max for other models to clarify its purpose
and avoid confusion.


// Setting safety to the lowest possible values since Gemini is already powerless enough
func CovertGemini2OpenAI(textRequest dto.GeneralOpenAIRequest, info *relaycommon.RelayInfo) (*GeminiChatRequest, error) {

Expand All @@ -65,49 +99,31 @@ func CovertGemini2OpenAI(textRequest dto.GeneralOpenAIRequest, info *relaycommon
}

if model_setting.GetGeminiSettings().ThinkingAdapterEnabled {
// 新增逻辑：处理 -thinking-<budget> 格式
if strings.Contains(info.OriginModelName, "-thinking-") {
parts := strings.SplitN(info.OriginModelName, "-thinking-", 2)
modelName := info.OriginModelName
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 strings.Contains(modelName, "-thinking-") {
parts := strings.SplitN(modelName, "-thinking-", 2)
if len(parts) == 2 && parts[1] != "" {
if budgetTokens, err := strconv.Atoi(parts[1]); err == nil {
// 从模型名称成功解析预算
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
}
}

clampedBudget := clampThinkingBudget(modelName, budgetTokens)
geminiRequest.GenerationConfig.ThinkingConfig = &GeminiThinkingConfig{
ThinkingBudget: common.GetPointer(budgetTokens),
ThinkingBudget: common.GetPointer(clampedBudget),
IncludeThoughts: true,
}
}
// 如果解析失败，则不设置ThinkingConfig，静默处理
}
} else if strings.HasSuffix(info.OriginModelName, "-thinking") { // 保留旧逻辑以兼容
// 硬编码不支持 ThinkingBudget 的旧模型
} else if strings.HasSuffix(modelName, "-thinking") {
unsupportedModels := []string{
"gemini-2.5-pro-preview-05-06",
"gemini-2.5-pro-preview-03-25",
}

isUnsupported := false
for _, unsupportedModel := range unsupportedModels {
if strings.HasPrefix(info.OriginModelName, unsupportedModel) {
if strings.HasPrefix(modelName, unsupportedModel) {
isUnsupported = true
break
}
Expand All @@ -119,39 +135,14 @@ func CovertGemini2OpenAI(textRequest dto.GeneralOpenAIRequest, info *relaycommon
}
} else {
budgetTokens := model_setting.GetGeminiSettings().ThinkingAdapterBudgetTokensPercentage * float64(geminiRequest.GenerationConfig.MaxOutputTokens)

// 检查是否为新的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
}
}

clampedBudget := clampThinkingBudget(modelName, int(budgetTokens))
geminiRequest.GenerationConfig.ThinkingConfig = &GeminiThinkingConfig{
ThinkingBudget: common.GetPointer(int(budgetTokens)),
ThinkingBudget: common.GetPointer(clampedBudget),
IncludeThoughts: true,
}
}
} else if strings.HasSuffix(info.OriginModelName, "-nothinking") {
// 检查是否为新的2.5pro模型（不支持-nothinking，因为最低值只能为128）
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
} else if strings.HasSuffix(modelName, "-nothinking") {
if !isNew25Pro && !is25FlashLite {
geminiRequest.GenerationConfig.ThinkingConfig = &GeminiThinkingConfig{
ThinkingBudget: common.GetPointer(0),
}
Expand Down