refactor(thinking): remove legacy utilities and simplify model mapping

This commit is contained in:
hkfires
2026-01-14 19:11:04 +08:00
parent 33d66959e9
commit 2262479365
9 changed files with 43 additions and 865 deletions

View File

@@ -1,49 +0,0 @@
package util
import (
"github.com/tidwall/gjson"
"github.com/tidwall/sjson"
)
// ApplyClaudeThinkingConfig applies thinking configuration to a Claude API request payload.
// It sets the thinking.type to "enabled" and thinking.budget_tokens to the specified budget.
// If budget is nil or the payload already has thinking config, it returns the payload unchanged.
func ApplyClaudeThinkingConfig(body []byte, budget *int) []byte {
if budget == nil {
return body
}
if gjson.GetBytes(body, "thinking").Exists() {
return body
}
if *budget <= 0 {
return body
}
updated := body
updated, _ = sjson.SetBytes(updated, "thinking.type", "enabled")
updated, _ = sjson.SetBytes(updated, "thinking.budget_tokens", *budget)
return updated
}
// ResolveClaudeThinkingConfig resolves thinking configuration from metadata for Claude models.
// It uses the unified ResolveThinkingConfigFromMetadata and normalizes the budget.
// Returns the normalized budget (nil if thinking should not be enabled) and whether it matched.
func ResolveClaudeThinkingConfig(modelName string, metadata map[string]any) (*int, bool) {
if !ModelSupportsThinking(modelName) {
return nil, false
}
budget, include, matched := ResolveThinkingConfigFromMetadata(modelName, metadata)
if !matched {
return nil, false
}
if include != nil && !*include {
return nil, true
}
if budget == nil {
return nil, true
}
normalized := NormalizeThinkingBudget(modelName, *budget)
if normalized <= 0 {
return nil, true
}
return &normalized, true
}

View File

@@ -8,12 +8,6 @@ import (
"github.com/tidwall/sjson"
)
const (
GeminiThinkingBudgetMetadataKey = "gemini_thinking_budget"
GeminiIncludeThoughtsMetadataKey = "gemini_include_thoughts"
GeminiOriginalModelMetadataKey = "gemini_original_model"
)
// Gemini model family detection patterns
var (
gemini3Pattern = regexp.MustCompile(`(?i)^gemini[_-]?3[_-]`)
@@ -297,104 +291,6 @@ func ApplyDefaultThinkingIfNeeded(model string, body []byte) []byte {
return updated
}
// ApplyGemini3ThinkingLevelFromMetadata applies thinkingLevel from metadata for Gemini 3 models.
// For standard Gemini API format (generationConfig.thinkingConfig path).
// This handles the case where reasoning_effort is specified via model name suffix (e.g., model(minimal))
// or numeric budget suffix (e.g., model(1000)) which gets converted to a thinkingLevel.
func ApplyGemini3ThinkingLevelFromMetadata(model string, metadata map[string]any, body []byte) []byte {
// Use the alias from metadata if available for model type detection
lookupModel := ResolveOriginalModel(model, metadata)
if !IsGemini3Model(lookupModel) && !IsGemini3Model(model) {
return body
}
// Determine which model to use for validation
checkModel := model
if IsGemini3Model(lookupModel) {
checkModel = lookupModel
}
// First try to get effort string from metadata
effort, ok := ReasoningEffortFromMetadata(metadata)
if ok && effort != "" {
if level, valid := ValidateGemini3ThinkingLevel(checkModel, effort); valid {
return ApplyGeminiThinkingLevel(body, level, nil)
}
}
// Fallback: check for numeric budget and convert to thinkingLevel
budget, _, _, matched := ThinkingFromMetadata(metadata)
if matched && budget != nil {
if level, valid := ThinkingBudgetToGemini3Level(checkModel, *budget); valid {
return ApplyGeminiThinkingLevel(body, level, nil)
}
}
return body
}
// ApplyGemini3ThinkingLevelFromMetadataCLI applies thinkingLevel from metadata for Gemini 3 models.
// For Gemini CLI API format (request.generationConfig.thinkingConfig path).
// This handles the case where reasoning_effort is specified via model name suffix (e.g., model(minimal))
// or numeric budget suffix (e.g., model(1000)) which gets converted to a thinkingLevel.
func ApplyGemini3ThinkingLevelFromMetadataCLI(model string, metadata map[string]any, body []byte) []byte {
// Use the alias from metadata if available for model type detection
lookupModel := ResolveOriginalModel(model, metadata)
if !IsGemini3Model(lookupModel) && !IsGemini3Model(model) {
return body
}
// Determine which model to use for validation
checkModel := model
if IsGemini3Model(lookupModel) {
checkModel = lookupModel
}
// First try to get effort string from metadata
effort, ok := ReasoningEffortFromMetadata(metadata)
if ok && effort != "" {
if level, valid := ValidateGemini3ThinkingLevel(checkModel, effort); valid {
return ApplyGeminiCLIThinkingLevel(body, level, nil)
}
}
// Fallback: check for numeric budget and convert to thinkingLevel
budget, _, _, matched := ThinkingFromMetadata(metadata)
if matched && budget != nil {
if level, valid := ThinkingBudgetToGemini3Level(checkModel, *budget); valid {
return ApplyGeminiCLIThinkingLevel(body, level, nil)
}
}
return body
}
// ApplyDefaultThinkingIfNeededCLI injects default thinkingConfig for models that require it.
// For Gemini CLI API format (request.generationConfig.thinkingConfig path).
// Returns the modified body if thinkingConfig was added, otherwise returns the original.
// For Gemini 3 models, uses thinkingLevel instead of thinkingBudget per Google's documentation.
func ApplyDefaultThinkingIfNeededCLI(model string, metadata map[string]any, body []byte) []byte {
// Use the alias from metadata if available for model property lookup
lookupModel := ResolveOriginalModel(model, metadata)
if !ModelHasDefaultThinking(lookupModel) && !ModelHasDefaultThinking(model) {
return body
}
if gjson.GetBytes(body, "request.generationConfig.thinkingConfig").Exists() {
return body
}
// Gemini 3 models use thinkingLevel instead of thinkingBudget
if IsGemini3Model(lookupModel) || IsGemini3Model(model) {
// Don't set a default - let the API use its dynamic default ("high")
// Only set includeThoughts
updated, _ := sjson.SetBytes(body, "request.generationConfig.thinkingConfig.includeThoughts", true)
return updated
}
// Gemini 2.5 and other models use thinkingBudget
updated, _ := sjson.SetBytes(body, "request.generationConfig.thinkingConfig.thinkingBudget", -1)
updated, _ = sjson.SetBytes(updated, "request.generationConfig.thinkingConfig.include_thoughts", true)
return updated
}
// StripThinkingConfigIfUnsupported removes thinkingConfig from the request body
// when the target model does not advertise Thinking capability. It cleans both
// standard Gemini and Gemini CLI JSON envelopes. This acts as a final safety net

View File

@@ -91,106 +91,6 @@ func thinkingRangeFromRegistry(model string) (found bool, min int, max int, zero
return false, 0, 0, false, false
}
// GetModelThinkingLevels returns the discrete reasoning effort levels for the model.
// Returns nil if the model has no thinking support or no levels defined.
//
// Deprecated: Access modelInfo.Thinking.Levels directly.
func GetModelThinkingLevels(model string) []string {
if model == "" {
return nil
}
info := registry.GetGlobalRegistry().GetModelInfo(model)
if info == nil || info.Thinking == nil {
return nil
}
return info.Thinking.Levels
}
// ModelUsesThinkingLevels reports whether the model uses discrete reasoning
// effort levels instead of numeric budgets.
//
// Deprecated: Check len(modelInfo.Thinking.Levels) > 0.
func ModelUsesThinkingLevels(model string) bool {
levels := GetModelThinkingLevels(model)
return len(levels) > 0
}
// NormalizeReasoningEffortLevel validates and normalizes a reasoning effort
// level for the given model. Returns false when the level is not supported.
//
// Deprecated: Use thinking.ValidateConfig for level validation.
func NormalizeReasoningEffortLevel(model, effort string) (string, bool) {
levels := GetModelThinkingLevels(model)
if len(levels) == 0 {
return "", false
}
loweredEffort := strings.ToLower(strings.TrimSpace(effort))
for _, lvl := range levels {
if strings.ToLower(lvl) == loweredEffort {
return lvl, true
}
}
return "", false
}
// IsOpenAICompatibilityModel reports whether the model is registered as an OpenAI-compatibility model.
// These models may not advertise Thinking metadata in the registry.
//
// Deprecated: Check modelInfo.Type == "openai-compatibility".
func IsOpenAICompatibilityModel(model string) bool {
if model == "" {
return false
}
info := registry.GetGlobalRegistry().GetModelInfo(model)
if info == nil {
return false
}
return strings.EqualFold(strings.TrimSpace(info.Type), "openai-compatibility")
}
// ThinkingEffortToBudget maps a reasoning effort level to a numeric thinking budget (tokens),
// clamping the result to the model's supported range.
//
// Mappings (values are normalized to model's supported range):
// - "none" -> 0
// - "auto" -> -1
// - "minimal" -> 512
// - "low" -> 1024
// - "medium" -> 8192
// - "high" -> 24576
// - "xhigh" -> 32768
//
// Returns false when the effort level is empty or unsupported.
//
// Deprecated: Use thinking.ConvertLevelToBudget instead.
func ThinkingEffortToBudget(model, effort string) (int, bool) {
if effort == "" {
return 0, false
}
normalized, ok := NormalizeReasoningEffortLevel(model, effort)
if !ok {
normalized = strings.ToLower(strings.TrimSpace(effort))
}
switch normalized {
case "none":
return 0, true
case "auto":
return NormalizeThinkingBudget(model, -1), true
case "minimal":
return NormalizeThinkingBudget(model, 512), true
case "low":
return NormalizeThinkingBudget(model, 1024), true
case "medium":
return NormalizeThinkingBudget(model, 8192), true
case "high":
return NormalizeThinkingBudget(model, 24576), true
case "xhigh":
return NormalizeThinkingBudget(model, 32768), true
default:
return 0, false
}
}
// ThinkingLevelToBudget maps a Gemini thinkingLevel to a numeric thinking budget (tokens).
//
// Mappings:
@@ -220,44 +120,3 @@ func ThinkingLevelToBudget(level string) (int, bool) {
return 0, false
}
}
// ThinkingBudgetToEffort maps a numeric thinking budget (tokens)
// to a reasoning effort level for level-based models.
//
// Mappings:
// - 0 -> "none" (or lowest supported level if model doesn't support "none")
// - -1 -> "auto"
// - 1..1024 -> "low"
// - 1025..8192 -> "medium"
// - 8193..24576 -> "high"
// - 24577.. -> highest supported level for the model (defaults to "xhigh")
//
// Returns false when the budget is unsupported (negative values other than -1).
//
// Deprecated: Use thinking.ConvertBudgetToLevel instead.
func ThinkingBudgetToEffort(model string, budget int) (string, bool) {
switch {
case budget == -1:
return "auto", true
case budget < -1:
return "", false
case budget == 0:
if levels := GetModelThinkingLevels(model); len(levels) > 0 {
return levels[0], true
}
return "none", true
case budget > 0 && budget <= 1024:
return "low", true
case budget <= 8192:
return "medium", true
case budget <= 24576:
return "high", true
case budget > 24576:
if levels := GetModelThinkingLevels(model); len(levels) > 0 {
return levels[len(levels)-1], true
}
return "xhigh", true
default:
return "", false
}
}

View File

@@ -1,130 +0,0 @@
package util
import (
"go/ast"
"go/parser"
"go/token"
"os"
"path/filepath"
"runtime"
"strings"
"testing"
)
func TestThinkingUtilDeprecationComments(t *testing.T) {
dir, err := thinkingSourceDir()
if err != nil {
t.Fatalf("resolve thinking source dir: %v", err)
}
// Test thinking.go deprecation comments
t.Run("thinking.go", func(t *testing.T) {
docs := parseFuncDocs(t, filepath.Join(dir, "thinking.go"))
tests := []struct {
funcName string
want string
}{
{"ModelSupportsThinking", "Deprecated: Use thinking.ApplyThinking with modelInfo.Thinking check."},
{"NormalizeThinkingBudget", "Deprecated: Use thinking.ValidateConfig for budget normalization."},
{"ThinkingEffortToBudget", "Deprecated: Use thinking.ConvertLevelToBudget instead."},
{"ThinkingBudgetToEffort", "Deprecated: Use thinking.ConvertBudgetToLevel instead."},
{"GetModelThinkingLevels", "Deprecated: Access modelInfo.Thinking.Levels directly."},
{"ModelUsesThinkingLevels", "Deprecated: Check len(modelInfo.Thinking.Levels) > 0."},
{"NormalizeReasoningEffortLevel", "Deprecated: Use thinking.ValidateConfig for level validation."},
{"IsOpenAICompatibilityModel", "Deprecated: Check modelInfo.Type == \"openai-compatibility\"."},
{"ThinkingLevelToBudget", "Deprecated: Use thinking.ConvertLevelToBudget instead."},
}
for _, tt := range tests {
t.Run(tt.funcName, func(t *testing.T) {
doc, ok := docs[tt.funcName]
if !ok {
t.Fatalf("missing function %q in thinking.go", tt.funcName)
}
if !strings.Contains(doc, tt.want) {
t.Fatalf("missing deprecation note for %s: want %q, got %q", tt.funcName, tt.want, doc)
}
})
}
})
// Test thinking_suffix.go deprecation comments
t.Run("thinking_suffix.go", func(t *testing.T) {
docs := parseFuncDocs(t, filepath.Join(dir, "thinking_suffix.go"))
tests := []struct {
funcName string
want string
}{
{"NormalizeThinkingModel", "Deprecated: Use thinking.ParseSuffix instead."},
{"ThinkingFromMetadata", "Deprecated: Access ThinkingConfig fields directly."},
{"ResolveThinkingConfigFromMetadata", "Deprecated: Use thinking.ApplyThinking instead."},
{"ReasoningEffortFromMetadata", "Deprecated: Use thinking.ConvertBudgetToLevel instead."},
{"ResolveOriginalModel", "Deprecated: Parse model suffix with thinking.ParseSuffix."},
}
for _, tt := range tests {
t.Run(tt.funcName, func(t *testing.T) {
doc, ok := docs[tt.funcName]
if !ok {
t.Fatalf("missing function %q in thinking_suffix.go", tt.funcName)
}
if !strings.Contains(doc, tt.want) {
t.Fatalf("missing deprecation note for %s: want %q, got %q", tt.funcName, tt.want, doc)
}
})
}
})
// Test thinking_text.go deprecation comments
t.Run("thinking_text.go", func(t *testing.T) {
docs := parseFuncDocs(t, filepath.Join(dir, "thinking_text.go"))
tests := []struct {
funcName string
want string
}{
{"GetThinkingText", "Deprecated: Use thinking package for thinking text extraction."},
{"GetThinkingTextFromJSON", "Deprecated: Use thinking package for thinking text extraction."},
{"SanitizeThinkingPart", "Deprecated: Use thinking package for thinking part sanitization."},
{"StripCacheControl", "Deprecated: Use thinking package for cache control stripping."},
}
for _, tt := range tests {
t.Run(tt.funcName, func(t *testing.T) {
doc, ok := docs[tt.funcName]
if !ok {
t.Fatalf("missing function %q in thinking_text.go", tt.funcName)
}
if !strings.Contains(doc, tt.want) {
t.Fatalf("missing deprecation note for %s: want %q, got %q", tt.funcName, tt.want, doc)
}
})
}
})
}
func parseFuncDocs(t *testing.T, path string) map[string]string {
t.Helper()
fset := token.NewFileSet()
file, err := parser.ParseFile(fset, path, nil, parser.ParseComments)
if err != nil {
t.Fatalf("parse %s: %v", path, err)
}
docs := map[string]string{}
for _, decl := range file.Decls {
fn, ok := decl.(*ast.FuncDecl)
if !ok || fn.Recv != nil {
continue
}
if fn.Doc == nil {
docs[fn.Name.Name] = ""
continue
}
docs[fn.Name.Name] = fn.Doc.Text()
}
return docs
}
func thinkingSourceDir() (string, error) {
_, thisFile, _, ok := runtime.Caller(0)
if !ok {
return "", os.ErrNotExist
}
return filepath.Dir(thisFile), nil
}

View File

@@ -1,319 +0,0 @@
package util
import (
"encoding/json"
"strconv"
"strings"
)
const (
// Deprecated: No longer used. Thinking configuration is now passed via
// model name suffix and processed by thinking.ApplyThinking().
ThinkingBudgetMetadataKey = "thinking_budget"
// Deprecated: No longer used. See ThinkingBudgetMetadataKey.
ThinkingIncludeThoughtsMetadataKey = "thinking_include_thoughts"
// Deprecated: No longer used. See ThinkingBudgetMetadataKey.
ReasoningEffortMetadataKey = "reasoning_effort"
// Deprecated: No longer used. The original model name (with suffix) is now
// preserved directly in the model field. Use thinking.ParseSuffix() to
// extract the base model name if needed.
ThinkingOriginalModelMetadataKey = "thinking_original_model"
// ModelMappingOriginalModelMetadataKey stores the client-requested model alias
// for OAuth model name mappings. This is NOT deprecated.
ModelMappingOriginalModelMetadataKey = "model_mapping_original_model"
)
// NormalizeThinkingModel parses dynamic thinking suffixes on model names and returns
// the normalized base model with extracted metadata. Supported pattern:
//
// Deprecated: Use thinking.ParseSuffix instead.
// - "(<value>)" where value can be:
// - A numeric budget (e.g., "(8192)", "(16384)")
// - A reasoning effort level (e.g., "(high)", "(medium)", "(low)")
//
// Examples:
// - "claude-sonnet-4-5-20250929(16384)" → budget=16384
// - "gpt-5.1(high)" → reasoning_effort="high"
// - "gemini-2.5-pro(32768)" → budget=32768
//
// Note: Empty parentheses "()" are not supported and will be ignored.
func NormalizeThinkingModel(modelName string) (string, map[string]any) {
if modelName == "" {
return modelName, nil
}
baseModel := modelName
var (
budgetOverride *int
reasoningEffort *string
matched bool
)
// Match "(<value>)" pattern at the end of the model name
if idx := strings.LastIndex(modelName, "("); idx != -1 {
if !strings.HasSuffix(modelName, ")") {
// Incomplete parenthesis, ignore
return baseModel, nil
}
value := modelName[idx+1 : len(modelName)-1] // Extract content between ( and )
if value == "" {
// Empty parentheses not supported
return baseModel, nil
}
candidateBase := modelName[:idx]
// Auto-detect: pure numeric → budget, string → reasoning effort level
if parsed, ok := parseIntPrefix(value); ok {
// Numeric value: treat as thinking budget
baseModel = candidateBase
budgetOverride = &parsed
matched = true
} else {
// String value: treat as reasoning effort level
baseModel = candidateBase
raw := strings.ToLower(strings.TrimSpace(value))
if raw != "" {
reasoningEffort = &raw
matched = true
}
}
}
if !matched {
return baseModel, nil
}
metadata := map[string]any{
ThinkingOriginalModelMetadataKey: modelName,
}
if budgetOverride != nil {
metadata[ThinkingBudgetMetadataKey] = *budgetOverride
}
if reasoningEffort != nil {
metadata[ReasoningEffortMetadataKey] = *reasoningEffort
}
return baseModel, metadata
}
// ThinkingFromMetadata extracts thinking overrides from metadata produced by NormalizeThinkingModel.
// It accepts both the new generic keys and legacy Gemini-specific keys.
//
// Deprecated: Access ThinkingConfig fields directly.
func ThinkingFromMetadata(metadata map[string]any) (*int, *bool, *string, bool) {
if len(metadata) == 0 {
return nil, nil, nil, false
}
var (
budgetPtr *int
includePtr *bool
effortPtr *string
matched bool
)
readBudget := func(key string) {
if budgetPtr != nil {
return
}
if raw, ok := metadata[key]; ok {
if v, okNumber := parseNumberToInt(raw); okNumber {
budget := v
budgetPtr = &budget
matched = true
}
}
}
readInclude := func(key string) {
if includePtr != nil {
return
}
if raw, ok := metadata[key]; ok {
switch v := raw.(type) {
case bool:
val := v
includePtr = &val
matched = true
case *bool:
if v != nil {
val := *v
includePtr = &val
matched = true
}
}
}
}
readEffort := func(key string) {
if effortPtr != nil {
return
}
if raw, ok := metadata[key]; ok {
if val, okStr := raw.(string); okStr && strings.TrimSpace(val) != "" {
normalized := strings.ToLower(strings.TrimSpace(val))
effortPtr = &normalized
matched = true
}
}
}
readBudget(ThinkingBudgetMetadataKey)
readBudget(GeminiThinkingBudgetMetadataKey)
readInclude(ThinkingIncludeThoughtsMetadataKey)
readInclude(GeminiIncludeThoughtsMetadataKey)
readEffort(ReasoningEffortMetadataKey)
readEffort("reasoning.effort")
return budgetPtr, includePtr, effortPtr, matched
}
// ResolveThinkingConfigFromMetadata derives thinking budget/include overrides,
// converting reasoning effort strings into budgets when possible.
//
// Deprecated: Use thinking.ApplyThinking instead.
func ResolveThinkingConfigFromMetadata(model string, metadata map[string]any) (*int, *bool, bool) {
budget, include, effort, matched := ThinkingFromMetadata(metadata)
if !matched {
return nil, nil, false
}
// Level-based models (OpenAI-style) do not accept numeric thinking budgets in
// Claude/Gemini-style protocols, so we don't derive budgets for them here.
if ModelUsesThinkingLevels(model) {
return nil, nil, false
}
if budget == nil && effort != nil {
if derived, ok := ThinkingEffortToBudget(model, *effort); ok {
budget = &derived
}
}
return budget, include, budget != nil || include != nil || effort != nil
}
// ReasoningEffortFromMetadata resolves a reasoning effort string from metadata,
// inferring "auto" and "none" when budgets request dynamic or disabled thinking.
//
// Deprecated: Use thinking.ConvertBudgetToLevel instead.
func ReasoningEffortFromMetadata(metadata map[string]any) (string, bool) {
budget, include, effort, matched := ThinkingFromMetadata(metadata)
if !matched {
return "", false
}
if effort != nil && *effort != "" {
return strings.ToLower(strings.TrimSpace(*effort)), true
}
if budget != nil {
switch *budget {
case -1:
return "auto", true
case 0:
return "none", true
}
}
if include != nil && !*include {
return "none", true
}
return "", true
}
// ResolveOriginalModel returns the original model name stored in metadata (if present),
// otherwise falls back to the provided model.
//
// Deprecated: Parse model suffix with thinking.ParseSuffix.
func ResolveOriginalModel(model string, metadata map[string]any) string {
normalize := func(name string) string {
if name == "" {
return ""
}
if base, _ := NormalizeThinkingModel(name); base != "" {
return base
}
return strings.TrimSpace(name)
}
if metadata != nil {
if v, ok := metadata[ModelMappingOriginalModelMetadataKey]; ok {
if s, okStr := v.(string); okStr && strings.TrimSpace(s) != "" {
if base := normalize(s); base != "" {
return base
}
}
}
if v, ok := metadata[ThinkingOriginalModelMetadataKey]; ok {
if s, okStr := v.(string); okStr && strings.TrimSpace(s) != "" {
if base := normalize(s); base != "" {
return base
}
}
}
if v, ok := metadata[GeminiOriginalModelMetadataKey]; ok {
if s, okStr := v.(string); okStr && strings.TrimSpace(s) != "" {
if base := normalize(s); base != "" {
return base
}
}
}
}
// Fallback: try to re-normalize the model name when metadata was dropped.
if base := normalize(model); base != "" {
return base
}
return model
}
func parseIntPrefix(value string) (int, bool) {
if value == "" {
return 0, false
}
digits := strings.TrimLeft(value, "-")
if digits == "" {
return 0, false
}
end := len(digits)
for i := 0; i < len(digits); i++ {
if digits[i] < '0' || digits[i] > '9' {
end = i
break
}
}
if end == 0 {
return 0, false
}
val, err := strconv.Atoi(digits[:end])
if err != nil {
return 0, false
}
return val, true
}
func parseNumberToInt(raw any) (int, bool) {
switch v := raw.(type) {
case int:
return v, true
case int32:
return int(v), true
case int64:
return int(v), true
case float64:
return int(v), true
case json.Number:
if val, err := v.Int64(); err == nil {
return int(val), true
}
case string:
if strings.TrimSpace(v) == "" {
return 0, false
}
if parsed, err := strconv.Atoi(strings.TrimSpace(v)); err == nil {
return parsed, true
}
}
return 0, false
}

View File

@@ -149,53 +149,32 @@ func TestApplyAPIKeyModelMapping(t *testing.T) {
_, _ = mgr.Register(ctx, apiKeyAuth)
tests := []struct {
name string
auth *Auth
inputModel string
wantModel string
wantOriginal string
expectMapping bool
name string
auth *Auth
inputModel string
wantModel string
}{
{
name: "api_key auth with alias",
auth: apiKeyAuth,
inputModel: "g25p(8192)",
wantModel: "gemini-2.5-pro-exp-03-25(8192)",
wantOriginal: "g25p(8192)",
expectMapping: true,
name: "api_key auth with alias",
auth: apiKeyAuth,
inputModel: "g25p(8192)",
wantModel: "gemini-2.5-pro-exp-03-25(8192)",
},
{
name: "oauth auth passthrough",
auth: oauthAuth,
inputModel: "some-model",
wantModel: "some-model",
expectMapping: false,
name: "oauth auth passthrough",
auth: oauthAuth,
inputModel: "some-model",
wantModel: "some-model",
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
metadata := map[string]any{"existing": "value"}
resolvedModel, resultMeta := mgr.applyAPIKeyModelMapping(tt.auth, tt.inputModel, metadata)
resolvedModel := mgr.applyAPIKeyModelMapping(tt.auth, tt.inputModel)
if resolvedModel != tt.wantModel {
t.Errorf("model = %q, want %q", resolvedModel, tt.wantModel)
}
if resultMeta["existing"] != "value" {
t.Error("existing metadata not preserved")
}
original, hasOriginal := resultMeta["model_mapping_original_model"].(string)
if tt.expectMapping {
if !hasOriginal || original != tt.wantOriginal {
t.Errorf("original model = %q, want %q", original, tt.wantOriginal)
}
} else {
if hasOriginal {
t.Error("should not set model_mapping_original_model for non-api_key auth")
}
}
})
}
}

View File

@@ -752,9 +752,9 @@ func (m *Manager) executeWithProvider(ctx context.Context, provider string, req
execCtx = context.WithValue(execCtx, "cliproxy.roundtripper", rt)
}
execReq := req
execReq.Model, execReq.Metadata = rewriteModelForAuth(routeModel, req.Metadata, auth)
execReq.Model, execReq.Metadata = m.applyOAuthModelMapping(auth, execReq.Model, execReq.Metadata)
execReq.Model, execReq.Metadata = m.applyAPIKeyModelMapping(auth, execReq.Model, execReq.Metadata)
execReq.Model = rewriteModelForAuth(routeModel, auth)
execReq.Model = m.applyOAuthModelMapping(auth, execReq.Model)
execReq.Model = m.applyAPIKeyModelMapping(auth, execReq.Model)
resp, errExec := executor.Execute(execCtx, auth, execReq, opts)
result := Result{AuthID: auth.ID, Provider: provider, Model: routeModel, Success: errExec == nil}
if errExec != nil {
@@ -801,9 +801,9 @@ func (m *Manager) executeCountWithProvider(ctx context.Context, provider string,
execCtx = context.WithValue(execCtx, "cliproxy.roundtripper", rt)
}
execReq := req
execReq.Model, execReq.Metadata = rewriteModelForAuth(routeModel, req.Metadata, auth)
execReq.Model, execReq.Metadata = m.applyOAuthModelMapping(auth, execReq.Model, execReq.Metadata)
execReq.Model, execReq.Metadata = m.applyAPIKeyModelMapping(auth, execReq.Model, execReq.Metadata)
execReq.Model = rewriteModelForAuth(routeModel, auth)
execReq.Model = m.applyOAuthModelMapping(auth, execReq.Model)
execReq.Model = m.applyAPIKeyModelMapping(auth, execReq.Model)
resp, errExec := executor.CountTokens(execCtx, auth, execReq, opts)
result := Result{AuthID: auth.ID, Provider: provider, Model: routeModel, Success: errExec == nil}
if errExec != nil {
@@ -850,9 +850,9 @@ func (m *Manager) executeStreamWithProvider(ctx context.Context, provider string
execCtx = context.WithValue(execCtx, "cliproxy.roundtripper", rt)
}
execReq := req
execReq.Model, execReq.Metadata = rewriteModelForAuth(routeModel, req.Metadata, auth)
execReq.Model, execReq.Metadata = m.applyOAuthModelMapping(auth, execReq.Model, execReq.Metadata)
execReq.Model, execReq.Metadata = m.applyAPIKeyModelMapping(auth, execReq.Model, execReq.Metadata)
execReq.Model = rewriteModelForAuth(routeModel, auth)
execReq.Model = m.applyOAuthModelMapping(auth, execReq.Model)
execReq.Model = m.applyAPIKeyModelMapping(auth, execReq.Model)
chunks, errStream := executor.ExecuteStream(execCtx, auth, execReq, opts)
if errStream != nil {
rerr := &Error{Message: errStream.Error()}
@@ -890,72 +890,39 @@ func (m *Manager) executeStreamWithProvider(ctx context.Context, provider string
}
}
func rewriteModelForAuth(model string, metadata map[string]any, auth *Auth) (string, map[string]any) {
func rewriteModelForAuth(model string, auth *Auth) string {
if auth == nil || model == "" {
return model, metadata
return model
}
prefix := strings.TrimSpace(auth.Prefix)
if prefix == "" {
return model, metadata
return model
}
needle := prefix + "/"
if !strings.HasPrefix(model, needle) {
return model, metadata
return model
}
rewritten := strings.TrimPrefix(model, needle)
return rewritten, stripPrefixFromMetadata(metadata, needle)
return strings.TrimPrefix(model, needle)
}
func stripPrefixFromMetadata(metadata map[string]any, needle string) map[string]any {
if len(metadata) == 0 || needle == "" {
return metadata
}
keys := []string{
util.GeminiOriginalModelMetadataKey,
util.ModelMappingOriginalModelMetadataKey,
}
var out map[string]any
for _, key := range keys {
raw, ok := metadata[key]
if !ok {
continue
}
value, okStr := raw.(string)
if !okStr || !strings.HasPrefix(value, needle) {
continue
}
if out == nil {
out = make(map[string]any, len(metadata))
for k, v := range metadata {
out[k] = v
}
}
out[key] = strings.TrimPrefix(value, needle)
}
if out == nil {
return metadata
}
return out
}
func (m *Manager) applyAPIKeyModelMapping(auth *Auth, requestedModel string, metadata map[string]any) (string, map[string]any) {
func (m *Manager) applyAPIKeyModelMapping(auth *Auth, requestedModel string) string {
if m == nil || auth == nil {
return requestedModel, metadata
return requestedModel
}
kind, _ := auth.AccountInfo()
if !strings.EqualFold(strings.TrimSpace(kind), "api_key") {
return requestedModel, metadata
return requestedModel
}
requestedModel = strings.TrimSpace(requestedModel)
if requestedModel == "" {
return requestedModel, metadata
return requestedModel
}
// Fast path: lookup per-auth mapping table (keyed by auth.ID).
if resolved := m.lookupAPIKeyUpstreamModel(auth.ID, requestedModel); resolved != "" {
return applyUpstreamModelOverride(requestedModel, resolved, metadata)
return resolved
}
// Slow path: scan config for the matching credential entry and resolve alias.
@@ -980,8 +947,11 @@ func (m *Manager) applyAPIKeyModelMapping(auth *Auth, requestedModel string, met
upstreamModel = resolveUpstreamModelForOpenAICompatAPIKey(cfg, auth, requestedModel)
}
// applyUpstreamModelOverride lives in model_name_mappings.go.
return applyUpstreamModelOverride(requestedModel, upstreamModel, metadata)
// Return upstream model if found, otherwise return requested model.
if upstreamModel != "" {
return upstreamModel
}
return requestedModel
}
// APIKeyConfigEntry is a generic interface for API key configurations.

View File

@@ -5,7 +5,6 @@ import (
internalconfig "github.com/router-for-me/CLIProxyAPI/v6/internal/config"
"github.com/router-for-me/CLIProxyAPI/v6/internal/thinking"
"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
)
type modelMappingEntry interface {
@@ -71,31 +70,14 @@ func (m *Manager) SetOAuthModelMappings(mappings map[string][]internalconfig.Mod
m.modelNameMappings.Store(table)
}
// applyOAuthModelMapping resolves the upstream model from OAuth model mappings
// and returns the resolved model along with updated metadata. If a mapping exists,
// the returned model is the upstream model and metadata contains the original
// requested model for response translation.
func (m *Manager) applyOAuthModelMapping(auth *Auth, requestedModel string, metadata map[string]any) (string, map[string]any) {
// applyOAuthModelMapping resolves the upstream model from OAuth model mappings.
// If a mapping exists, the returned model is the upstream model.
func (m *Manager) applyOAuthModelMapping(auth *Auth, requestedModel string) string {
upstreamModel := m.resolveOAuthUpstreamModel(auth, requestedModel)
return applyUpstreamModelOverride(requestedModel, upstreamModel, metadata)
}
func applyUpstreamModelOverride(requestedModel, upstreamModel string, metadata map[string]any) (string, map[string]any) {
if upstreamModel == "" {
return requestedModel, metadata
return requestedModel
}
out := make(map[string]any, 1)
if len(metadata) > 0 {
out = make(map[string]any, len(metadata)+1)
for k, v := range metadata {
out[k] = v
}
}
// Preserve the original client model string (including any suffix) for downstream.
out[util.ModelMappingOriginalModelMetadataKey] = requestedModel
return upstreamModel, out
return upstreamModel
}
func resolveModelAliasFromConfigModels(requestedModel string, models []modelMappingEntry) string {

View File

@@ -169,19 +169,9 @@ func TestApplyOAuthModelMapping_SuffixPreservation(t *testing.T) {
mgr.SetOAuthModelMappings(mappings)
auth := &Auth{ID: "test-auth-id", Provider: "gemini-cli"}
metadata := map[string]any{"existing": "value"}
resolvedModel, resultMeta := mgr.applyOAuthModelMapping(auth, "gemini-2.5-pro(8192)", metadata)
resolvedModel := mgr.applyOAuthModelMapping(auth, "gemini-2.5-pro(8192)")
if resolvedModel != "gemini-2.5-pro-exp-03-25(8192)" {
t.Errorf("applyOAuthModelMapping() model = %q, want %q", resolvedModel, "gemini-2.5-pro-exp-03-25(8192)")
}
originalModel, ok := resultMeta["model_mapping_original_model"].(string)
if !ok || originalModel != "gemini-2.5-pro(8192)" {
t.Errorf("applyOAuthModelMapping() metadata[model_mapping_original_model] = %v, want %q", resultMeta["model_mapping_original_model"], "gemini-2.5-pro(8192)")
}
if resultMeta["existing"] != "value" {
t.Errorf("applyOAuthModelMapping() metadata[existing] = %v, want %q", resultMeta["existing"], "value")
}
}