**feat(executor): add thinking level to budget conversion utility**

- Introduced `ConvertThinkingLevelToBudget` to map thinking level ("high"/"low") to corresponding budget values.
- Applied the utility in `aistudio_executor.go` before stripping unsupported configs.
- Updated dependencies to include `tidwall/gjson` for JSON parsing.
This commit is contained in:
Luis Pater
2025-11-21 00:48:12 +08:00
parent db81331ae8
commit cbcfeb92cc
2 changed files with 48 additions and 0 deletions

View File

@@ -264,6 +264,7 @@ func (e *AIStudioExecutor) translateRequest(req cliproxyexecutor.Request, opts c
} }
payload = util.ApplyGeminiThinkingConfig(payload, budgetOverride, includeOverride) payload = util.ApplyGeminiThinkingConfig(payload, budgetOverride, includeOverride)
} }
payload = util.ConvertThinkingLevelToBudget(payload)
payload = util.StripThinkingConfigIfUnsupported(req.Model, payload) payload = util.StripThinkingConfigIfUnsupported(req.Model, payload)
payload = fixGeminiImageAspectRatio(req.Model, payload) payload = fixGeminiImageAspectRatio(req.Model, payload)
payload = applyPayloadConfig(e.cfg, req.Model, payload) payload = applyPayloadConfig(e.cfg, req.Model, payload)

View File

@@ -5,6 +5,7 @@ import (
"strconv" "strconv"
"strings" "strings"
"github.com/tidwall/gjson"
"github.com/tidwall/sjson" "github.com/tidwall/sjson"
) )
@@ -212,3 +213,49 @@ func StripThinkingConfigIfUnsupported(model string, body []byte) []byte {
updated, _ = sjson.DeleteBytes(updated, "generationConfig.thinkingConfig") updated, _ = sjson.DeleteBytes(updated, "generationConfig.thinkingConfig")
return updated return updated
} }
// ConvertThinkingLevelToBudget checks for "generationConfig.thinkingConfig.thinkingLevel"
// and converts it to "thinkingBudget".
// "high" -> 32768
// "low" -> 128
// It removes "thinkingLevel" after conversion.
func ConvertThinkingLevelToBudget(body []byte) []byte {
levelPath := "generationConfig.thinkingConfig.thinkingLevel"
res := gjson.GetBytes(body, levelPath)
if !res.Exists() {
return body
}
level := strings.ToLower(res.String())
var budget int
switch level {
case "high":
budget = 32768
case "low":
budget = 128
default:
// If unknown level, we might just leave it or default.
// User only specified high and low. We'll assume we shouldn't touch it if it's something else,
// or maybe we should just remove the invalid level?
// For safety adhering to strict instructions: "If high... if low...".
// If it's something else, the upstream might fail anyway if we leave it,
// but let's just delete the level if we processed it.
// Actually, let's check if we need to do anything for other values.
// For now, only handle high/low.
return body
}
// Set budget
budgetPath := "generationConfig.thinkingConfig.thinkingBudget"
updated, err := sjson.SetBytes(body, budgetPath, budget)
if err != nil {
return body
}
// Remove level
updated, err = sjson.DeleteBytes(updated, levelPath)
if err != nil {
return body
}
return updated
}