diff --git a/controller/compact_filter_test.go b/controller/compact_filter_test.go new file mode 100644 index 000000000000..185104066c1d --- /dev/null +++ b/controller/compact_filter_test.go @@ -0,0 +1,22 @@ +package controller + +import ( + "testing" + + "github.com/QuantumNous/new-api/model" + "github.com/stretchr/testify/require" +) + +func TestFilterPricingHidesCompactModels(t *testing.T) { + pricing := []model.Pricing{ + {ModelName: "gpt-5.4", EnableGroup: []string{"default"}}, + {ModelName: "gpt-5.4-openai-compact", EnableGroup: []string{"default"}}, + {ModelName: "gpt-5.4-mini", EnableGroup: []string{"default"}}, + } + + filtered := filterPricingByUsableGroups(pricing, map[string]string{"default": "default"}) + + require.Len(t, filtered, 2) + require.Equal(t, "gpt-5.4", filtered[0].ModelName) + require.Equal(t, "gpt-5.4-mini", filtered[1].ModelName) +} diff --git a/controller/model.go b/controller/model.go index 4dbd45838dd8..56475ff0b184 100644 --- a/controller/model.go +++ b/controller/model.go @@ -3,6 +3,7 @@ package controller import ( "fmt" "net/http" + "strings" "time" "github.com/QuantumNous/new-api/common" @@ -18,6 +19,7 @@ import ( "github.com/QuantumNous/new-api/relay/helper" "github.com/QuantumNous/new-api/service" "github.com/QuantumNous/new-api/setting/operation_setting" + "github.com/QuantumNous/new-api/setting/ratio_setting" "github.com/QuantumNous/new-api/types" "github.com/gin-gonic/gin" "github.com/samber/lo" @@ -132,7 +134,10 @@ func ListModels(c *gin.Context, modelType int) { } else { tokenModelLimit = map[string]bool{} } - for allowModel, _ := range tokenModelLimit { + for allowModel := range tokenModelLimit { + if isCompactModelName(allowModel) { + continue + } if !acceptUnsetRatioModel { if !helper.HasModelBillingConfig(allowModel) { continue @@ -180,6 +185,9 @@ func ListModels(c *gin.Context, modelType int) { models = model.GetGroupEnabledModels(group) } for _, modelName := range models { + if isCompactModelName(modelName) { + continue + } if !acceptUnsetRatioModel { if !helper.HasModelBillingConfig(modelName) { continue @@ -241,24 +249,52 @@ func ListModels(c *gin.Context, modelType int) { func ChannelListModels(c *gin.Context) { c.JSON(200, gin.H{ "success": true, - "data": openAIModels, + "data": filterCompactOpenAIModels(openAIModels), }) } func DashboardListModels(c *gin.Context) { + filtered := make(map[int][]string, len(channelId2Models)) + for channelId, models := range channelId2Models { + filtered[channelId] = filterCompactModelNames(models) + } c.JSON(200, gin.H{ "success": true, - "data": channelId2Models, + "data": filtered, }) } func EnabledListModels(c *gin.Context) { c.JSON(200, gin.H{ "success": true, - "data": model.GetEnabledModels(), + "data": filterCompactModelNames(model.GetEnabledModels()), }) } +func isCompactModelName(modelName string) bool { + return strings.HasSuffix(modelName, ratio_setting.CompactModelSuffix) +} + +func filterCompactModelNames(modelNames []string) []string { + filtered := make([]string, 0, len(modelNames)) + for _, modelName := range modelNames { + if !isCompactModelName(modelName) { + filtered = append(filtered, modelName) + } + } + return filtered +} + +func filterCompactOpenAIModels(models []dto.OpenAIModels) []dto.OpenAIModels { + filtered := make([]dto.OpenAIModels, 0, len(models)) + for _, item := range models { + if !isCompactModelName(item.Id) { + filtered = append(filtered, item) + } + } + return filtered +} + func RetrieveModel(c *gin.Context, modelType int) { modelId := c.Param("model") if aiModel, ok := openAIModelsMap[modelId]; ok { diff --git a/controller/pricing.go b/controller/pricing.go index 8252327244c4..6e7ebc940093 100644 --- a/controller/pricing.go +++ b/controller/pricing.go @@ -19,6 +19,9 @@ func filterPricingByUsableGroups(pricing []model.Pricing, usableGroup map[string filtered := make([]model.Pricing, 0, len(pricing)) for _, item := range pricing { + if isCompactModelName(item.ModelName) { + continue + } if common.StringsContains(item.EnableGroup, "all") { filtered = append(filtered, item) continue diff --git a/setting/ratio_setting/cache_ratio.go b/setting/ratio_setting/cache_ratio.go index fe6e3b3262a4..54f0a552d3b3 100644 --- a/setting/ratio_setting/cache_ratio.go +++ b/setting/ratio_setting/cache_ratio.go @@ -142,6 +142,11 @@ func UpdateCreateCacheRatioByJSONString(jsonStr string) error { func GetCacheRatio(name string) (float64, bool) { ratio, ok := cacheRatioMap.Get(name) if !ok { + if baseModelName, isCompact := CompactBaseModelName(name); isCompact { + if ratio, ok := cacheRatioMap.Get(baseModelName); ok { + return ratio, true + } + } return 1, false // Default to 1 if not found } return ratio, true @@ -150,6 +155,11 @@ func GetCacheRatio(name string) (float64, bool) { func GetCreateCacheRatio(name string) (float64, bool) { ratio, ok := createCacheRatioMap.Get(name) if !ok { + if baseModelName, isCompact := CompactBaseModelName(name); isCompact { + if ratio, ok := createCacheRatioMap.Get(baseModelName); ok { + return ratio, true + } + } return 1.25, false // Default to 1.25 if not found } return ratio, true diff --git a/setting/ratio_setting/compact_suffix.go b/setting/ratio_setting/compact_suffix.go index 2d2fe3c34bb9..b5747e33d356 100644 --- a/setting/ratio_setting/compact_suffix.go +++ b/setting/ratio_setting/compact_suffix.go @@ -5,6 +5,13 @@ import "strings" const CompactModelSuffix = "-openai-compact" const CompactWildcardModelKey = "*" + CompactModelSuffix +func CompactBaseModelName(modelName string) (string, bool) { + if !strings.HasSuffix(modelName, CompactModelSuffix) { + return modelName, false + } + return strings.TrimSuffix(modelName, CompactModelSuffix), true +} + func WithCompactModelSuffix(modelName string) string { if strings.HasSuffix(modelName, CompactModelSuffix) { return modelName diff --git a/setting/ratio_setting/compact_suffix_test.go b/setting/ratio_setting/compact_suffix_test.go new file mode 100644 index 000000000000..9f3905b645d4 --- /dev/null +++ b/setting/ratio_setting/compact_suffix_test.go @@ -0,0 +1,77 @@ +package ratio_setting + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestCompactModelBillingInheritsBaseModel(t *testing.T) { + restoreRatioSettings(t) + + require.NoError(t, UpdateModelPriceByJSONString(`{"gpt-5.4":0.42,"*-openai-compact":9.9}`)) + require.NoError(t, UpdateModelRatioByJSONString(`{"gpt-5.4":3.14,"*-openai-compact":8.8}`)) + require.NoError(t, UpdateCompletionRatioByJSONString(`{"gpt-5.4":6.28}`)) + require.NoError(t, UpdateCacheRatioByJSONString(`{"gpt-5.4":0.5}`)) + require.NoError(t, UpdateCreateCacheRatioByJSONString(`{"gpt-5.4":1.5}`)) + + compact := "gpt-5.4-openai-compact" + price, ok := GetModelPrice(compact, false) + require.True(t, ok) + require.Equal(t, 0.42, price) + + ratio, ok, matchName := GetModelRatio(compact) + require.True(t, ok) + require.Equal(t, 3.14, ratio) + require.Equal(t, "gpt-5.4", matchName) + + require.Equal(t, 6.28, GetCompletionRatio(compact)) + + cacheRatio, ok := GetCacheRatio(compact) + require.True(t, ok) + require.Equal(t, 0.5, cacheRatio) + + createCacheRatio, ok := GetCreateCacheRatio(compact) + require.True(t, ok) + require.Equal(t, 1.5, createCacheRatio) +} + +func TestCompactModelBillingFallsBackToWildcard(t *testing.T) { + restoreRatioSettings(t) + + require.NoError(t, UpdateModelPriceByJSONString(`{"*-openai-compact":9.9}`)) + require.NoError(t, UpdateModelRatioByJSONString(`{"*-openai-compact":8.8}`)) + + compact := "gpt-5.4-openai-compact" + price, ok := GetModelPrice(compact, false) + require.True(t, ok) + require.Equal(t, 9.9, price) + + ratio, ok, matchName := GetModelRatio(compact) + require.True(t, ok) + require.Equal(t, 8.8, ratio) + require.Equal(t, CompactWildcardModelKey, matchName) +} + +func TestWithCompactModelSuffixDoesNotDuplicateSuffix(t *testing.T) { + require.Equal(t, "gpt-5.4-openai-compact", WithCompactModelSuffix("gpt-5.4")) + require.Equal(t, "gpt-5.4-openai-compact", WithCompactModelSuffix("gpt-5.4-openai-compact")) +} + +func restoreRatioSettings(t *testing.T) { + t.Helper() + + modelPriceJSON := ModelPrice2JSONString() + modelRatioJSON := ModelRatio2JSONString() + completionRatioJSON := CompletionRatio2JSONString() + cacheRatioJSON := CacheRatio2JSONString() + createCacheRatioJSON := CreateCacheRatio2JSONString() + + t.Cleanup(func() { + require.NoError(t, UpdateModelPriceByJSONString(modelPriceJSON)) + require.NoError(t, UpdateModelRatioByJSONString(modelRatioJSON)) + require.NoError(t, UpdateCompletionRatioByJSONString(completionRatioJSON)) + require.NoError(t, UpdateCacheRatioByJSONString(cacheRatioJSON)) + require.NoError(t, UpdateCreateCacheRatioByJSONString(createCacheRatioJSON)) + }) +} diff --git a/setting/ratio_setting/model_ratio.go b/setting/ratio_setting/model_ratio.go index 80702ee42ad2..a022c4700ae6 100644 --- a/setting/ratio_setting/model_ratio.go +++ b/setting/ratio_setting/model_ratio.go @@ -371,7 +371,10 @@ func GetModelPrice(name string, printErr bool) (float64, bool) { return price, true } - if strings.HasSuffix(name, CompactModelSuffix) { + if baseModelName, isCompact := CompactBaseModelName(name); isCompact { + if price, ok := modelPriceMap.Get(baseModelName); ok { + return price, true + } price, ok := modelPriceMap.Get(CompactWildcardModelKey) if !ok { if printErr { @@ -405,9 +408,12 @@ func GetModelRatio(name string) (float64, bool, string) { ratio, ok := modelRatioMap.Get(name) if !ok { - if strings.HasSuffix(name, CompactModelSuffix) { + if baseModelName, isCompact := CompactBaseModelName(name); isCompact { + if baseRatio, ok := modelRatioMap.Get(baseModelName); ok { + return baseRatio, true, baseModelName + } if wildcardRatio, ok := modelRatioMap.Get(CompactWildcardModelKey); ok { - return wildcardRatio, true, name + return wildcardRatio, true, CompactWildcardModelKey } //return 0, true, name } @@ -448,6 +454,11 @@ func GetCompletionRatio(name string) float64 { return ratio } } + if baseModelName, isCompact := CompactBaseModelName(name); isCompact { + if ratio, ok := completionRatioMap.Get(baseModelName); ok { + return ratio + } + } hardCodedRatio, contain := getHardcodedCompletionModelRatio(name) if contain { return hardCodedRatio @@ -475,6 +486,15 @@ func GetCompletionRatioInfo(name string) CompletionRatioInfo { } } + if baseModelName, isCompact := CompactBaseModelName(name); isCompact { + if ratio, ok := completionRatioMap.Get(baseModelName); ok { + return CompletionRatioInfo{ + Ratio: ratio, + Locked: false, + } + } + } + hardCodedRatio, locked := getHardcodedCompletionModelRatio(name) if locked { return CompletionRatioInfo{