// Package amp provides model mapping functionality for routing Amp CLI requests // to alternative models when the requested model is not available locally. package amp import ( "regexp" "strings" "sync" "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" log "github.com/sirupsen/logrus" ) // ModelMapper provides model name mapping/aliasing for Amp CLI requests. // When an Amp request comes in for a model that isn't available locally, // this mapper can redirect it to an alternative model that IS available. type ModelMapper interface { // MapModel returns the target model name if a mapping exists and the target // model has available providers. Returns empty string if no mapping applies. MapModel(requestedModel string) string // UpdateMappings refreshes the mapping configuration (for hot-reload). UpdateMappings(mappings []config.AmpModelMapping) } // DefaultModelMapper implements ModelMapper with thread-safe mapping storage. type DefaultModelMapper struct { mu sync.RWMutex mappings map[string]string // exact: from -> to (normalized lowercase keys) regexps []regexMapping // regex rules evaluated in order } // NewModelMapper creates a new model mapper with the given initial mappings. func NewModelMapper(mappings []config.AmpModelMapping) *DefaultModelMapper { m := &DefaultModelMapper{ mappings: make(map[string]string), regexps: nil, } m.UpdateMappings(mappings) return m } // MapModel checks if a mapping exists for the requested model and if the // target model has available local providers. Returns the mapped model name // or empty string if no valid mapping exists. // // If the requested model contains a thinking suffix (e.g., "g25p(8192)"), // the suffix is preserved in the returned model name (e.g., "gemini-2.5-pro(8192)"). // However, if the mapping target already contains a suffix, the config suffix // takes priority over the user's suffix. func (m *DefaultModelMapper) MapModel(requestedModel string) string { if requestedModel == "" { return "" } m.mu.RLock() defer m.mu.RUnlock() // Extract thinking suffix from requested model using ParseSuffix requestResult := thinking.ParseSuffix(requestedModel) baseModel := requestResult.ModelName // Normalize the base model for lookup (case-insensitive) normalizedBase := strings.ToLower(strings.TrimSpace(baseModel)) // Check for direct mapping using base model name targetModel, exists := m.mappings[normalizedBase] if !exists { // Try regex mappings in order using base model only // (suffix is handled separately via ParseSuffix) for _, rm := range m.regexps { if rm.re.MatchString(baseModel) { targetModel = rm.to exists = true break } } if !exists { return "" } } // Check if target model already has a thinking suffix (config priority) targetResult := thinking.ParseSuffix(targetModel) // Verify target model has available providers (use base model for lookup) providers := util.GetProviderName(targetResult.ModelName) if len(providers) == 0 { log.Debugf("amp model mapping: target model %s has no available providers, skipping mapping", targetModel) return "" } // Suffix handling: config suffix takes priority, otherwise preserve user suffix if targetResult.HasSuffix { // Config's "to" already contains a suffix - use it as-is (config priority) return targetModel } // Preserve user's thinking suffix on the mapped model // (skip empty suffixes to avoid returning "model()") if requestResult.HasSuffix && requestResult.RawSuffix != "" { return targetModel + "(" + requestResult.RawSuffix + ")" } // Note: Detailed routing log is handled by logAmpRouting in fallback_handlers.go return targetModel } // UpdateMappings refreshes the mapping configuration from config. // This is called during initialization and on config hot-reload. func (m *DefaultModelMapper) UpdateMappings(mappings []config.AmpModelMapping) { m.mu.Lock() defer m.mu.Unlock() // Clear and rebuild mappings m.mappings = make(map[string]string, len(mappings)) m.regexps = make([]regexMapping, 0, len(mappings)) for _, mapping := range mappings { from := strings.TrimSpace(mapping.From) to := strings.TrimSpace(mapping.To) if from == "" || to == "" { log.Warnf("amp model mapping: skipping invalid mapping (from=%q, to=%q)", from, to) continue } if mapping.Regex { // Compile case-insensitive regex; wrap with (?i) to match behavior of exact lookups pattern := "(?i)" + from re, err := regexp.Compile(pattern) if err != nil { log.Warnf("amp model mapping: invalid regex %q: %v", from, err) continue } m.regexps = append(m.regexps, regexMapping{re: re, to: to}) log.Debugf("amp model regex mapping registered: /%s/ -> %s", from, to) } else { // Store with normalized lowercase key for case-insensitive lookup normalizedFrom := strings.ToLower(from) m.mappings[normalizedFrom] = to log.Debugf("amp model mapping registered: %s -> %s", from, to) } } if len(m.mappings) > 0 { log.Infof("amp model mapping: loaded %d mapping(s)", len(m.mappings)) } if n := len(m.regexps); n > 0 { log.Infof("amp model mapping: loaded %d regex mapping(s)", n) } } // GetMappings returns a copy of current mappings (for debugging/status). func (m *DefaultModelMapper) GetMappings() map[string]string { m.mu.RLock() defer m.mu.RUnlock() result := make(map[string]string, len(m.mappings)) for k, v := range m.mappings { result[k] = v } return result } type regexMapping struct { re *regexp.Regexp to string }