luispater commited on
Commit
712a84b
·
unverified ·
1 Parent(s): 261f82c

Fixed: #291

Browse files

**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.

internal/runtime/executor/aistudio_executor.go CHANGED
@@ -264,6 +264,7 @@ func (e *AIStudioExecutor) translateRequest(req cliproxyexecutor.Request, opts c
264
  }
265
  payload = util.ApplyGeminiThinkingConfig(payload, budgetOverride, includeOverride)
266
  }
 
267
  payload = util.StripThinkingConfigIfUnsupported(req.Model, payload)
268
  payload = fixGeminiImageAspectRatio(req.Model, payload)
269
  payload = applyPayloadConfig(e.cfg, req.Model, payload)
 
264
  }
265
  payload = util.ApplyGeminiThinkingConfig(payload, budgetOverride, includeOverride)
266
  }
267
+ payload = util.ConvertThinkingLevelToBudget(payload)
268
  payload = util.StripThinkingConfigIfUnsupported(req.Model, payload)
269
  payload = fixGeminiImageAspectRatio(req.Model, payload)
270
  payload = applyPayloadConfig(e.cfg, req.Model, payload)
internal/util/gemini_thinking.go CHANGED
@@ -5,6 +5,7 @@ import (
5
  "strconv"
6
  "strings"
7
 
 
8
  "github.com/tidwall/sjson"
9
  )
10
 
@@ -212,3 +213,49 @@ func StripThinkingConfigIfUnsupported(model string, body []byte) []byte {
212
  updated, _ = sjson.DeleteBytes(updated, "generationConfig.thinkingConfig")
213
  return updated
214
  }
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5
  "strconv"
6
  "strings"
7
 
8
+ "github.com/tidwall/gjson"
9
  "github.com/tidwall/sjson"
10
  )
11
 
 
213
  updated, _ = sjson.DeleteBytes(updated, "generationConfig.thinkingConfig")
214
  return updated
215
  }
216
+
217
+ // ConvertThinkingLevelToBudget checks for "generationConfig.thinkingConfig.thinkingLevel"
218
+ // and converts it to "thinkingBudget".
219
+ // "high" -> 32768
220
+ // "low" -> 128
221
+ // It removes "thinkingLevel" after conversion.
222
+ func ConvertThinkingLevelToBudget(body []byte) []byte {
223
+ levelPath := "generationConfig.thinkingConfig.thinkingLevel"
224
+ res := gjson.GetBytes(body, levelPath)
225
+ if !res.Exists() {
226
+ return body
227
+ }
228
+
229
+ level := strings.ToLower(res.String())
230
+ var budget int
231
+ switch level {
232
+ case "high":
233
+ budget = 32768
234
+ case "low":
235
+ budget = 128
236
+ default:
237
+ // If unknown level, we might just leave it or default.
238
+ // User only specified high and low. We'll assume we shouldn't touch it if it's something else,
239
+ // or maybe we should just remove the invalid level?
240
+ // For safety adhering to strict instructions: "If high... if low...".
241
+ // If it's something else, the upstream might fail anyway if we leave it,
242
+ // but let's just delete the level if we processed it.
243
+ // Actually, let's check if we need to do anything for other values.
244
+ // For now, only handle high/low.
245
+ return body
246
+ }
247
+
248
+ // Set budget
249
+ budgetPath := "generationConfig.thinkingConfig.thinkingBudget"
250
+ updated, err := sjson.SetBytes(body, budgetPath, budget)
251
+ if err != nil {
252
+ return body
253
+ }
254
+
255
+ // Remove level
256
+ updated, err = sjson.DeleteBytes(updated, levelPath)
257
+ if err != nil {
258
+ return body
259
+ }
260
+ return updated
261
+ }