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
98 98
99 translated = applyThinkingMetadataCLI(translated, req.Metadata, req.Model)99 translated = applyThinkingMetadataCLI(translated, req.Metadata, req.Model)
100 translated = util.ApplyGemini3ThinkingLevelFromMetadataCLI(req.Model, req.Metadata, translated)100 translated = util.ApplyGemini3ThinkingLevelFromMetadataCLI(req.Model, req.Metadata, translated)
101 translated = util.ApplyDefaultThinkingIfNeededCLI(req.Model, translated)101 translated = util.ApplyDefaultThinkingIfNeededCLI(req.Model, req.Metadata, translated)
102 translated = normalizeAntigravityThinking(req.Model, translated, isClaude)102 translated = normalizeAntigravityThinking(req.Model, translated, isClaude)
103 translated = applyPayloadConfigWithRoot(e.cfg, req.Model, "antigravity", "request", translated)103 translated = applyPayloadConfigWithRoot(e.cfg, req.Model, "antigravity", "request", translated)
104 104
@@ -193,7 +193,7 @@ func (e *AntigravityExecutor) executeClaudeNonStream(ctx context.Context, auth * · 88 unmodified lines
193 193
194 translated = applyThinkingMetadataCLI(translated, req.Metadata, req.Model)194 translated = applyThinkingMetadataCLI(translated, req.Metadata, req.Model)
195 translated = util.ApplyGemini3ThinkingLevelFromMetadataCLI(req.Model, req.Metadata, translated)195 translated = util.ApplyGemini3ThinkingLevelFromMetadataCLI(req.Model, req.Metadata, translated)
196 translated = util.ApplyDefaultThinkingIfNeededCLI(req.Model, translated)196 translated = util.ApplyDefaultThinkingIfNeededCLI(req.Model, req.Metadata, translated)
197 translated = normalizeAntigravityThinking(req.Model, translated, true)197 translated = normalizeAntigravityThinking(req.Model, translated, true)
198 translated = applyPayloadConfigWithRoot(e.cfg, req.Model, "antigravity", "request", translated)198 translated = applyPayloadConfigWithRoot(e.cfg, req.Model, "antigravity", "request", translated)
199 199
@@ -529,7 +529,7 @@ func (e *AntigravityExecutor) ExecuteStream(ctx context.Context, auth *cliproxya · 329 unmodified lines
529 529
530 translated = applyThinkingMetadataCLI(translated, req.Metadata, req.Model)530 translated = applyThinkingMetadataCLI(translated, req.Metadata, req.Model)
531 translated = util.ApplyGemini3ThinkingLevelFromMetadataCLI(req.Model, req.Metadata, translated)531 translated = util.ApplyGemini3ThinkingLevelFromMetadataCLI(req.Model, req.Metadata, translated)
532 translated = util.ApplyDefaultThinkingIfNeededCLI(req.Model, translated)532 translated = util.ApplyDefaultThinkingIfNeededCLI(req.Model, req.Metadata, translated)
533 translated = normalizeAntigravityThinking(req.Model, translated, isClaude)533 translated = normalizeAntigravityThinking(req.Model, translated, isClaude)
534 translated = applyPayloadConfigWithRoot(e.cfg, req.Model, "antigravity", "request", translated)534 translated = applyPayloadConfigWithRoot(e.cfg, req.Model, "antigravity", "request", translated)
535 535
@@ -698,7 +698,7 @@ func (e *AntigravityExecutor) CountTokens(ctx context.Context, auth *cliproxyaut · 162 unmodified lines
698 for idx, baseURL := range baseURLs {698 for idx, baseURL := range baseURLs {
699 payload := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), false)699 payload := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), false)
700 payload = applyThinkingMetadataCLI(payload, req.Metadata, req.Model)700 payload = applyThinkingMetadataCLI(payload, req.Metadata, req.Model)
701 payload = util.ApplyDefaultThinkingIfNeededCLI(req.Model, payload)701 payload = util.ApplyDefaultThinkingIfNeededCLI(req.Model, req.Metadata, payload)
702 payload = normalizeAntigravityThinking(req.Model, payload, isClaude)702 payload = normalizeAntigravityThinking(req.Model, payload, isClaude)
703 payload = deleteJSONField(payload, "project")703 payload = deleteJSONField(payload, "project")
704 payload = deleteJSONField(payload, "model")704 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
80 basePayload := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), false)80 basePayload := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), false)
81 basePayload = applyThinkingMetadataCLI(basePayload, req.Metadata, req.Model)81 basePayload = applyThinkingMetadataCLI(basePayload, req.Metadata, req.Model)
82 basePayload = util.ApplyGemini3ThinkingLevelFromMetadataCLI(req.Model, req.Metadata, basePayload)82 basePayload = util.ApplyGemini3ThinkingLevelFromMetadataCLI(req.Model, req.Metadata, basePayload)
83 basePayload = util.ApplyDefaultThinkingIfNeededCLI(req.Model, basePayload)83 basePayload = util.ApplyDefaultThinkingIfNeededCLI(req.Model, req.Metadata, basePayload)
84 basePayload = util.NormalizeGeminiCLIThinkingBudget(req.Model, basePayload)84 basePayload = util.NormalizeGeminiCLIThinkingBudget(req.Model, basePayload)
85 basePayload = util.StripThinkingConfigIfUnsupported(req.Model, basePayload)85 basePayload = util.StripThinkingConfigIfUnsupported(req.Model, basePayload)
86 basePayload = fixGeminiCLIImageAspectRatio(req.Model, basePayload)86 basePayload = fixGeminiCLIImageAspectRatio(req.Model, basePayload)
@@ -219,7 +219,7 @@ func (e *GeminiCLIExecutor) ExecuteStream(ctx context.Context, auth *cliproxyaut · 132 unmodified lines
219 basePayload := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), true)219 basePayload := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), true)
220 basePayload = applyThinkingMetadataCLI(basePayload, req.Metadata, req.Model)220 basePayload = applyThinkingMetadataCLI(basePayload, req.Metadata, req.Model)
221 basePayload = util.ApplyGemini3ThinkingLevelFromMetadataCLI(req.Model, req.Metadata, basePayload)221 basePayload = util.ApplyGemini3ThinkingLevelFromMetadataCLI(req.Model, req.Metadata, basePayload)
222 basePayload = util.ApplyDefaultThinkingIfNeededCLI(req.Model, basePayload)222 basePayload = util.ApplyDefaultThinkingIfNeededCLI(req.Model, req.Metadata, basePayload)
223 basePayload = util.NormalizeGeminiCLIThinkingBudget(req.Model, basePayload)223 basePayload = util.NormalizeGeminiCLIThinkingBudget(req.Model, basePayload)
224 basePayload = util.StripThinkingConfigIfUnsupported(req.Model, basePayload)224 basePayload = util.StripThinkingConfigIfUnsupported(req.Model, basePayload)
225 basePayload = fixGeminiCLIImageAspectRatio(req.Model, basePayload)225 basePayload = fixGeminiCLIImageAspectRatio(req.Model, basePayload)
internal/runtime/executor/payload_helpers.go+12−6
@@ -14,15 +14,18 @@ import ( · 13 unmodified lines
14// ApplyThinkingMetadata applies thinking config from model suffix metadata (e.g., (high), (8192))14// ApplyThinkingMetadata applies thinking config from model suffix metadata (e.g., (high), (8192))
15// for standard Gemini format payloads. It normalizes the budget when the model supports thinking.15// for standard Gemini format payloads. It normalizes the budget when the model supports thinking.
16func ApplyThinkingMetadata(payload []byte, metadata map[string]any, model string) []byte {16func 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)
18 if !ok || (budgetOverride == nil && includeOverride == nil) {21 if !ok || (budgetOverride == nil && includeOverride == nil) {
19 return payload22 return payload
20 }23 }
21 if !util.ModelSupportsThinking(model) {24 if !util.ModelSupportsThinking(lookupModel) {
22 return payload25 return payload
23 }26 }
24 if budgetOverride != nil {27 if budgetOverride != nil {
25 norm := util.NormalizeThinkingBudget(model, *budgetOverride)28 norm := util.NormalizeThinkingBudget(lookupModel, *budgetOverride)
26 budgetOverride = &norm29 budgetOverride = &norm
27 }30 }
28 return util.ApplyGeminiThinkingConfig(payload, budgetOverride, includeOverride)31 return util.ApplyGeminiThinkingConfig(payload, budgetOverride, includeOverride)
@@ -31,15 +34,18 @@ func ApplyThinkingMetadata(payload []byte, metadata map[string]any, model string · 2 unmodified lines
31// applyThinkingMetadataCLI applies thinking config from model suffix metadata (e.g., (high), (8192))34// applyThinkingMetadataCLI applies thinking config from model suffix metadata (e.g., (high), (8192))
32// for Gemini CLI format payloads (nested under "request"). It normalizes the budget when the model supports thinking.35// for Gemini CLI format payloads (nested under "request"). It normalizes the budget when the model supports thinking.
33func applyThinkingMetadataCLI(payload []byte, metadata map[string]any, model string) []byte {36func 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)
35 if !ok || (budgetOverride == nil && includeOverride == nil) {41 if !ok || (budgetOverride == nil && includeOverride == nil) {
36 return payload42 return payload
37 }43 }
38 if !util.ModelSupportsThinking(model) {44 if !util.ModelSupportsThinking(lookupModel) {
39 return payload45 return payload
40 }46 }
41 if budgetOverride != nil {47 if budgetOverride != nil {
42 norm := util.NormalizeThinkingBudget(model, *budgetOverride)48 norm := util.NormalizeThinkingBudget(lookupModel, *budgetOverride)
43 budgetOverride = &norm49 budgetOverride = &norm
44 }50 }
45 return util.ApplyGeminiCLIThinkingConfig(payload, budgetOverride, includeOverride)51 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
290// For standard Gemini API format (generationConfig.thinkingConfig path).290// For standard Gemini API format (generationConfig.thinkingConfig path).
291// This handles the case where reasoning_effort is specified via model name suffix (e.g., model(minimal)).291// This handles the case where reasoning_effort is specified via model name suffix (e.g., model(minimal)).
292func ApplyGemini3ThinkingLevelFromMetadata(model string, metadata map[string]any, body []byte) []byte {292func 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) {
294 return body296 return body
295 }297 }
296 effort, ok := ReasoningEffortFromMetadata(metadata)298 effort, ok := ReasoningEffortFromMetadata(metadata)
297 if !ok || effort == "" {299 if !ok || effort == "" {
298 return body300 return body
299 }301 }
300 // Validate and apply the thinkingLevel302 // Validate and apply the thinkingLevel using the model that matches Gemini 3 pattern
301 if level, valid := ValidateGemini3ThinkingLevel(model, effort); valid {303 checkModel := model
304 if IsGemini3Model(lookupModel) {
305 checkModel = lookupModel
306 }
307 if level, valid := ValidateGemini3ThinkingLevel(checkModel, effort); valid {
302 return ApplyGeminiThinkingLevel(body, level, nil)308 return ApplyGeminiThinkingLevel(body, level, nil)
303 }309 }
304 return body310 return body
@@ -308,15 +314,21 @@ func ApplyGemini3ThinkingLevelFromMetadata(model string, metadata map[string]any · 3 unmodified lines
308// For Gemini CLI API format (request.generationConfig.thinkingConfig path).314// For Gemini CLI API format (request.generationConfig.thinkingConfig path).
309// This handles the case where reasoning_effort is specified via model name suffix (e.g., model(minimal)).315// This handles the case where reasoning_effort is specified via model name suffix (e.g., model(minimal)).
310func ApplyGemini3ThinkingLevelFromMetadataCLI(model string, metadata map[string]any, body []byte) []byte {316func 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) {
312 return body320 return body
313 }321 }
314 effort, ok := ReasoningEffortFromMetadata(metadata)322 effort, ok := ReasoningEffortFromMetadata(metadata)
315 if !ok || effort == "" {323 if !ok || effort == "" {
316 return body324 return body
317 }325 }
318 // Validate and apply the thinkingLevel326 // Validate and apply the thinkingLevel using the model that matches Gemini 3 pattern
319 if level, valid := ValidateGemini3ThinkingLevel(model, effort); valid {327 checkModel := model
328 if IsGemini3Model(lookupModel) {
329 checkModel = lookupModel
330 }
331 if level, valid := ValidateGemini3ThinkingLevel(checkModel, effort); valid {
320 return ApplyGeminiCLIThinkingLevel(body, level, nil)332 return ApplyGeminiCLIThinkingLevel(body, level, nil)
321 }333 }
322 return body334 return body
@@ -326,15 +338,17 @@ func ApplyGemini3ThinkingLevelFromMetadataCLI(model string, metadata map[string] · 3 unmodified lines
326// For Gemini CLI API format (request.generationConfig.thinkingConfig path).338// For Gemini CLI API format (request.generationConfig.thinkingConfig path).
327// Returns the modified body if thinkingConfig was added, otherwise returns the original.339// Returns the modified body if thinkingConfig was added, otherwise returns the original.
328// For Gemini 3 models, uses thinkingLevel instead of thinkingBudget per Google's documentation.340// For Gemini 3 models, uses thinkingLevel instead of thinkingBudget per Google's documentation.
329func ApplyDefaultThinkingIfNeededCLI(model string, body []byte) []byte {341func ApplyDefaultThinkingIfNeededCLI(model string, metadata map[string]any, body []byte) []byte {
330 if !ModelHasDefaultThinking(model) {342 // Use the alias from metadata if available for model property lookup
343 lookupModel := ResolveOriginalModel(model, metadata)
344 if !ModelHasDefaultThinking(lookupModel) && !ModelHasDefaultThinking(model) {
331 return body345 return body
332 }346 }
333 if gjson.GetBytes(body, "request.generationConfig.thinkingConfig").Exists() {347 if gjson.GetBytes(body, "request.generationConfig.thinkingConfig").Exists() {
334 return body348 return body
335 }349 }
336 // Gemini 3 models use thinkingLevel instead of thinkingBudget350 // Gemini 3 models use thinkingLevel instead of thinkingBudget
337 if IsGemini3Model(model) {351 if IsGemini3Model(lookupModel) || IsGemini3Model(model) {
338 // Don't set a default - let the API use its dynamic default ("high")352 // Don't set a default - let the API use its dynamic default ("high")
339 // Only set includeThoughts353 // Only set includeThoughts
340 updated, _ := sjson.SetBytes(body, "request.generationConfig.thinkingConfig.includeThoughts", true)354 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
81 out[k] = v81 out[k] = v
82 }82 }
83 }83 }
84 out[util.ModelMappingOriginalModelMetadataKey] = upstreamModel84 // 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
85 return upstreamModel, out87 return upstreamModel, out
86}88}
87 89