mirror of
https://github.com/router-for-me/CLIProxyAPI.git
synced 2026-02-03 13:00:52 +08:00
ed5ec5b55c
This commit introduces several improvements to the AMP (Advanced Model Proxy) module: - **Model Mapping Logic:** The `FallbackHandler` now uses a more robust approach for model mapping. It includes the extraction and preservation of dynamic "thinking suffixes" (e.g., `(xhigh)`) during mapping, ensuring that these configurations are correctly applied to the mapped model. A new `resolveMappedModel` function centralizes this logic for cleaner code. - **ModelMapper Verification:** The `ModelMapper` in `model_mapping.go` now verifies that the target model of a mapping has available providers *after* normalizing it. This prevents mappings to non-existent or unresolvable models. - **Gemini Thinking Configuration Cleanup:** In `gemini_thinking.go`, unnecessary `generationConfig.thinkingConfig.include_thoughts` and `generationConfig.thinkingConfig.thinkingBudget` fields are now deleted from the request body when applying Gemini thinking levels. This prevents potential conflicts or redundant configurations. - **Testing:** A new test case `TestModelMapper_MapModel_TargetWithThinkingSuffix` has been added to `model_mapping_test.go` to specifically cover the preservation of thinking suffixes during model mapping.
114 lines
3.5 KiB
Go
114 lines
3.5 KiB
Go
// 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 (
|
|
"strings"
|
|
"sync"
|
|
|
|
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
|
|
"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 // from -> to (normalized lowercase keys)
|
|
}
|
|
|
|
// NewModelMapper creates a new model mapper with the given initial mappings.
|
|
func NewModelMapper(mappings []config.AmpModelMapping) *DefaultModelMapper {
|
|
m := &DefaultModelMapper{
|
|
mappings: make(map[string]string),
|
|
}
|
|
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.
|
|
func (m *DefaultModelMapper) MapModel(requestedModel string) string {
|
|
if requestedModel == "" {
|
|
return ""
|
|
}
|
|
|
|
m.mu.RLock()
|
|
defer m.mu.RUnlock()
|
|
|
|
// Normalize the requested model for lookup
|
|
normalizedRequest := strings.ToLower(strings.TrimSpace(requestedModel))
|
|
|
|
// Check for direct mapping
|
|
targetModel, exists := m.mappings[normalizedRequest]
|
|
if !exists {
|
|
return ""
|
|
}
|
|
|
|
// Verify target model has available providers
|
|
normalizedTarget, _ := util.NormalizeThinkingModel(targetModel)
|
|
providers := util.GetProviderName(normalizedTarget)
|
|
if len(providers) == 0 {
|
|
log.Debugf("amp model mapping: target model %s has no available providers, skipping mapping", targetModel)
|
|
return ""
|
|
}
|
|
|
|
// 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))
|
|
|
|
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
|
|
}
|
|
|
|
// 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))
|
|
}
|
|
}
|
|
|
|
// 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
|
|
}
|