History

fix(thinking): use model alias for thinking config resolution in mapped models

hkfires committed 10 months ago

Commit 89db4e91 parent e3324195 files · +44 −22Browse files

Unified · Split

5 files
  1. internal/runtime/executor/antigravity_executor.go +4 −4
  2. internal/runtime/executor/gemini_cli_executor.go +2 −2
  3. internal/runtime/executor/payload_helpers.go +12 −6
  4. internal/util/gemini_thinking.go +23 −9
  5. sdk/cliproxy/auth/model_name_mappings.go +3 −1
internal/runtime/executor/antigravity_executor.go+4−4
@@ -98,7 +98,7 @@ func (e *AntigravityExecutor) Execute(ctx context.Context, auth *cliproxyauth.Au · 97 unmodified lines
9898
9999 translated = applyThinkingMetadataCLI(translated, req.Metadata, req.Model)
100100 translated = util.ApplyGemini3ThinkingLevelFromMetadataCLI(req.Model, req.Metadata, translated)
101− translated = util.ApplyDefaultThinkingIfNeededCLI(req.Model, translated)
101+ translated = util.ApplyDefaultThinkingIfNeededCLI(req.Model, req.Metadata, translated)
102102 translated = normalizeAntigravityThinking(req.Model, translated, isClaude)
103103 translated = applyPayloadConfigWithRoot(e.cfg, req.Model, "antigravity", "request", translated)
104104
@@ -193,7 +193,7 @@ func (e *AntigravityExecutor) executeClaudeNonStream(ctx context.Context, auth * · 88 unmodified lines
193193
194194 translated = applyThinkingMetadataCLI(translated, req.Metadata, req.Model)
195195 translated = util.ApplyGemini3ThinkingLevelFromMetadataCLI(req.Model, req.Metadata, translated)
196− translated = util.ApplyDefaultThinkingIfNeededCLI(req.Model, translated)
196+ translated = util.ApplyDefaultThinkingIfNeededCLI(req.Model, req.Metadata, translated)
197197 translated = normalizeAntigravityThinking(req.Model, translated, true)
198198 translated = applyPayloadConfigWithRoot(e.cfg, req.Model, "antigravity", "request", translated)
199199
@@ -529,7 +529,7 @@ func (e *AntigravityExecutor) ExecuteStream(ctx context.Context, auth *cliproxya · 329 unmodified lines
529529
530530 translated = applyThinkingMetadataCLI(translated, req.Metadata, req.Model)
531531 translated = util.ApplyGemini3ThinkingLevelFromMetadataCLI(req.Model, req.Metadata, translated)
532− translated = util.ApplyDefaultThinkingIfNeededCLI(req.Model, translated)
532+ translated = util.ApplyDefaultThinkingIfNeededCLI(req.Model, req.Metadata, translated)
533533 translated = normalizeAntigravityThinking(req.Model, translated, isClaude)
534534 translated = applyPayloadConfigWithRoot(e.cfg, req.Model, "antigravity", "request", translated)
535535
@@ -698,7 +698,7 @@ func (e *AntigravityExecutor) CountTokens(ctx context.Context, auth *cliproxyaut · 162 unmodified lines
698698 for idx, baseURL := range baseURLs {
699699 payload := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), false)
700700 payload = applyThinkingMetadataCLI(payload, req.Metadata, req.Model)
701− payload = util.ApplyDefaultThinkingIfNeededCLI(req.Model, payload)
701+ payload = util.ApplyDefaultThinkingIfNeededCLI(req.Model, req.Metadata, payload)
702702 payload = normalizeAntigravityThinking(req.Model, payload, isClaude)
703703 payload = deleteJSONField(payload, "project")
704704 payload = deleteJSONField(payload, "model")
internal/runtime/executor/gemini_cli_executor.go+2−2
@@ -80,7 +80,7 @@ func (e *GeminiCLIExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth · 79 unmodified lines
8080 basePayload := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), false)
8181 basePayload = applyThinkingMetadataCLI(basePayload, req.Metadata, req.Model)
8282 basePayload = util.ApplyGemini3ThinkingLevelFromMetadataCLI(req.Model, req.Metadata, basePayload)
83− basePayload = util.ApplyDefaultThinkingIfNeededCLI(req.Model, basePayload)
83+ basePayload = util.ApplyDefaultThinkingIfNeededCLI(req.Model, req.Metadata, basePayload)
8484 basePayload = util.NormalizeGeminiCLIThinkingBudget(req.Model, basePayload)
8585 basePayload = util.StripThinkingConfigIfUnsupported(req.Model, basePayload)
8686 basePayload = fixGeminiCLIImageAspectRatio(req.Model, basePayload)
@@ -219,7 +219,7 @@ func (e *GeminiCLIExecutor) ExecuteStream(ctx context.Context, auth *cliproxyaut · 132 unmodified lines
219219 basePayload := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), true)
220220 basePayload = applyThinkingMetadataCLI(basePayload, req.Metadata, req.Model)
221221 basePayload = util.ApplyGemini3ThinkingLevelFromMetadataCLI(req.Model, req.Metadata, basePayload)
222− basePayload = util.ApplyDefaultThinkingIfNeededCLI(req.Model, basePayload)
222+ basePayload = util.ApplyDefaultThinkingIfNeededCLI(req.Model, req.Metadata, basePayload)
223223 basePayload = util.NormalizeGeminiCLIThinkingBudget(req.Model, basePayload)
224224 basePayload = util.StripThinkingConfigIfUnsupported(req.Model, basePayload)
225225 basePayload = fixGeminiCLIImageAspectRatio(req.Model, basePayload)
internal/runtime/executor/payload_helpers.go+12−6
@@ -14,15 +14,18 @@ import ( · 13 unmodified lines
1414 // ApplyThinkingMetadata applies thinking config from model suffix metadata (e.g., (high), (8192))
1515 // for standard Gemini format payloads. It normalizes the budget when the model supports thinking.
1616 func ApplyThinkingMetadata(payload []byte, metadata map[string]any, model string) []byte {
17− budgetOverride, includeOverride, ok := util.ResolveThinkingConfigFromMetadata(model, metadata)
17+ // Use the alias from metadata if available, as it's registered in the global registry
18+ // with thinking metadata; the upstream model name may not be registered.
19+ lookupModel := util.ResolveOriginalModel(model, metadata)
20+ budgetOverride, includeOverride, ok := util.ResolveThinkingConfigFromMetadata(lookupModel, metadata)
1821 if !ok || (budgetOverride == nil && includeOverride == nil) {
1922 return payload
2023 }
21− if !util.ModelSupportsThinking(model) {
24+ if !util.ModelSupportsThinking(lookupModel) {
2225 return payload
2326 }
2427 if budgetOverride != nil {
25− norm := util.NormalizeThinkingBudget(model, *budgetOverride)
28+ norm := util.NormalizeThinkingBudget(lookupModel, *budgetOverride)
2629 budgetOverride = &norm
2730 }
2831 return util.ApplyGeminiThinkingConfig(payload, budgetOverride, includeOverride)
@@ -31,15 +34,18 @@ func ApplyThinkingMetadata(payload []byte, metadata map[string]any, model string · 2 unmodified lines
3134 // applyThinkingMetadataCLI applies thinking config from model suffix metadata (e.g., (high), (8192))
3235 // for Gemini CLI format payloads (nested under "request"). It normalizes the budget when the model supports thinking.
3336 func applyThinkingMetadataCLI(payload []byte, metadata map[string]any, model string) []byte {
34− budgetOverride, includeOverride, ok := util.ResolveThinkingConfigFromMetadata(model, metadata)
37+ // Use the alias from metadata if available, as it's registered in the global registry
38+ // with thinking metadata; the upstream model name may not be registered.
39+ lookupModel := util.ResolveOriginalModel(model, metadata)
40+ budgetOverride, includeOverride, ok := util.ResolveThinkingConfigFromMetadata(lookupModel, metadata)
3541 if !ok || (budgetOverride == nil && includeOverride == nil) {
3642 return payload
3743 }
38− if !util.ModelSupportsThinking(model) {
44+ if !util.ModelSupportsThinking(lookupModel) {
3945 return payload
4046 }
4147 if budgetOverride != nil {
42− norm := util.NormalizeThinkingBudget(model, *budgetOverride)
48+ norm := util.NormalizeThinkingBudget(lookupModel, *budgetOverride)
4349 budgetOverride = &norm
4450 }
4551 return util.ApplyGeminiCLIThinkingConfig(payload, budgetOverride, includeOverride)
internal/util/gemini_thinking.go+23−9
@@ -290,15 +290,21 @@ func ApplyDefaultThinkingIfNeeded(model string, body []byte) []byte { · 289 unmodified lines
290290 // For standard Gemini API format (generationConfig.thinkingConfig path).
291291 // This handles the case where reasoning_effort is specified via model name suffix (e.g., model(minimal)).
292292 func ApplyGemini3ThinkingLevelFromMetadata(model string, metadata map[string]any, body []byte) []byte {
293− if !IsGemini3Model(model) {
293+ // Use the alias from metadata if available for model type detection
294+ lookupModel := ResolveOriginalModel(model, metadata)
295+ if !IsGemini3Model(lookupModel) && !IsGemini3Model(model) {
294296 return body
295297 }
296298 effort, ok := ReasoningEffortFromMetadata(metadata)
297299 if !ok || effort == "" {
298300 return body
299301 }
300− // Validate and apply the thinkingLevel
301− if level, valid := ValidateGemini3ThinkingLevel(model, effort); valid {
302+ // Validate and apply the thinkingLevel using the model that matches Gemini 3 pattern
303+ checkModel := model
304+ if IsGemini3Model(lookupModel) {
305+ checkModel = lookupModel
306+ }
307+ if level, valid := ValidateGemini3ThinkingLevel(checkModel, effort); valid {
302308 return ApplyGeminiThinkingLevel(body, level, nil)
303309 }
304310 return body
@@ -308,15 +314,21 @@ func ApplyGemini3ThinkingLevelFromMetadata(model string, metadata map[string]any · 3 unmodified lines
308314 // For Gemini CLI API format (request.generationConfig.thinkingConfig path).
309315 // This handles the case where reasoning_effort is specified via model name suffix (e.g., model(minimal)).
310316 func ApplyGemini3ThinkingLevelFromMetadataCLI(model string, metadata map[string]any, body []byte) []byte {
311− if !IsGemini3Model(model) {
317+ // Use the alias from metadata if available for model type detection
318+ lookupModel := ResolveOriginalModel(model, metadata)
319+ if !IsGemini3Model(lookupModel) && !IsGemini3Model(model) {
312320 return body
313321 }
314322 effort, ok := ReasoningEffortFromMetadata(metadata)
315323 if !ok || effort == "" {
316324 return body
317325 }
318− // Validate and apply the thinkingLevel
319− if level, valid := ValidateGemini3ThinkingLevel(model, effort); valid {
326+ // Validate and apply the thinkingLevel using the model that matches Gemini 3 pattern
327+ checkModel := model
328+ if IsGemini3Model(lookupModel) {
329+ checkModel = lookupModel
330+ }
331+ if level, valid := ValidateGemini3ThinkingLevel(checkModel, effort); valid {
320332 return ApplyGeminiCLIThinkingLevel(body, level, nil)
321333 }
322334 return body
@@ -326,15 +338,17 @@ func ApplyGemini3ThinkingLevelFromMetadataCLI(model string, metadata map[string] · 3 unmodified lines
326338 // For Gemini CLI API format (request.generationConfig.thinkingConfig path).
327339 // Returns the modified body if thinkingConfig was added, otherwise returns the original.
328340 // For Gemini 3 models, uses thinkingLevel instead of thinkingBudget per Google's documentation.
329−func ApplyDefaultThinkingIfNeededCLI(model string, body []byte) []byte {
330− if !ModelHasDefaultThinking(model) {
341+func ApplyDefaultThinkingIfNeededCLI(model string, metadata map[string]any, body []byte) []byte {
342+ // Use the alias from metadata if available for model property lookup
343+ lookupModel := ResolveOriginalModel(model, metadata)
344+ if !ModelHasDefaultThinking(lookupModel) && !ModelHasDefaultThinking(model) {
331345 return body
332346 }
333347 if gjson.GetBytes(body, "request.generationConfig.thinkingConfig").Exists() {
334348 return body
335349 }
336350 // Gemini 3 models use thinkingLevel instead of thinkingBudget
337− if IsGemini3Model(model) {
351+ if IsGemini3Model(lookupModel) || IsGemini3Model(model) {
338352 // Don't set a default - let the API use its dynamic default ("high")
339353 // Only set includeThoughts
340354 updated, _ := sjson.SetBytes(body, "request.generationConfig.thinkingConfig.includeThoughts", true)
sdk/cliproxy/auth/model_name_mappings.go+3−1
@@ -81,7 +81,9 @@ func (m *Manager) applyOAuthModelMapping(auth *Auth, requestedModel string, meta · 80 unmodified lines
8181 out[k] = v
8282 }
8383 }
84− out[util.ModelMappingOriginalModelMetadataKey] = upstreamModel
84+ // Store the requested alias (e.g., "gp") so downstream can use it to look up
85+ // model metadata from the global registry where it was registered under this alias.
86+ out[util.ModelMappingOriginalModelMetadataKey] = requestedModel
8587 return upstreamModel, out
8688 }
8789