feat: publish sanitized custom.97 source

service/automatic_models_test.go

@@ -5,10 +5,13 @@ import (
 	"github.com/QuantumNous/new-api/common"
 	"github.com/QuantumNous/new-api/constant"
 	"github.com/QuantumNous/new-api/model"
+	perfmetrics "github.com/QuantumNous/new-api/pkg/perf_metrics"
 	"github.com/QuantumNous/new-api/setting"
+	"github.com/QuantumNous/new-api/setting/model_setting"
 	"github.com/QuantumNous/new-api/setting/ratio_setting"
 	"github.com/gin-gonic/gin"
 	"github.com/glebarez/sqlite"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"gorm.io/driver/mysql"
 	"gorm.io/driver/postgres"
@@ -45,9 +48,9 @@ func TestAutomaticPoolModelsDatabaseMatrix(t *testing.T) {
 			model.DB = db
 			common.RedisEnabled = false
 			common.MemoryCacheEnabled = false
-			require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.Ability{}))
+			require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.Ability{}, &model.Token{}))
 			t.Cleanup(func() {
-				require.NoError(t, db.Migrator().DropTable(&model.Ability{}, &model.Channel{}))
+				require.NoError(t, db.Migrator().DropTable(&model.Token{}, &model.Ability{}, &model.Channel{}))
 				model.DB = oldDB
 				common.RedisEnabled = oldRedis
 				common.MemoryCacheEnabled = oldMemory
@@ -92,6 +95,86 @@ func TestAutomaticPoolModelsDatabaseMatrix(t *testing.T) {
 			c.Request = httptest.NewRequest("GET", "/v1/models", nil)
 			common.SetContextKey(c, constant.ContextKeyUserGroup, "auto-a")
 			require.Equal(t, []string{"gpt-6-astra", "model-a", "new-model", "shared"}, AutomaticPoolModels(c))
+
+			cfg := model_setting.GetGlobalSettings()
+			oldImageGroup, oldImageModels := cfg.ImageGenerationGroup, cfg.ImageGenerationModels
+			cfg.ImageGenerationGroup, cfg.ImageGenerationModels = "auto-b", []string{"blacklist-image"}
+			t.Cleanup(func() {
+				cfg.ImageGenerationGroup, cfg.ImageGenerationModels = oldImageGroup, oldImageModels
+				automaticMetricCache.Purge()
+				_ = poolCooldowns().Purge()
+			})
+			createChannelSelectAutoGroupsChannel(t, db, 571, "auto-a", "blacklist-model")
+			createChannelSelectAutoGroupsChannel(t, db, 572, "auto-b", "blacklist-model")
+			createChannelSelectAutoGroupsChannel(t, db, 573, "auto-a", "blacklist-image")
+			createChannelSelectAutoGroupsChannel(t, db, 574, "auto-b", "blacklist-image")
+			common.MemoryCacheEnabled = true
+			model.InitChannelCache()
+			automaticMetricCache.Set("blacklist-model", perfmetrics.QueryResult{})
+			automaticMetricCache.Set("blacklist-image", perfmetrics.QueryResult{})
+			makeContext := func(policy *model.TokenRoutingPolicy, tokenID int) *gin.Context {
+				ctx, _ := gin.CreateTestContext(httptest.NewRecorder())
+				ctx.Request = httptest.NewRequest("POST", "/v1/responses", nil)
+				ctx.Set("token_id", tokenID)
+				ctx.Set("id", 7)
+				common.SetContextKey(ctx, constant.ContextKeyUserGroup, "auto-a")
+				ctx.Set(model.TokenRoutingContextKey, policy)
+				return ctx
+			}
+			for _, preference := range []string{"cost", "speed", "stable"} {
+				for _, blocked := range [][]string{nil, {"auto-a"}, {"auto-a", "auto-b"}} {
+					policy := &model.TokenRoutingPolicy{Version: 1, Mode: "smart", AutoPreference: preference, ExcludedGroups: blocked}
+					body, err := common.Marshal(policy)
+					require.NoError(t, err)
+					token := model.Token{Key: fmt.Sprintf("blacklist-%s-%d", preference, len(blocked)), UserId: 7, RoutingPolicy: string(body)}
+					require.NoError(t, token.Insert())
+					var stored model.Token
+					require.NoError(t, db.First(&stored, token.Id).Error)
+					loaded, err := stored.GetRoutingPolicy()
+					require.NoError(t, err)
+					assert.Equal(t, blocked, loaded.ExcludedGroups)
+					ctx := makeContext(loaded, stored.Id)
+					pool := ResolveRequestPool(ctx, "blacklist-model")
+					require.Len(t, pool, 2-len(blocked))
+					assert.Equal(t, max(0, 2*len(pool)-1), PoolRetryBudget(ctx, "blacklist-model"))
+					for _, entry := range pool {
+						assert.False(t, loaded.ExcludesGroup(entry.Group))
+					}
+					channel, group, selectErr := SelectChannelForRequest(ctx, "blacklist-model", &RetryParam{Ctx: ctx, TokenGroup: "auto", ModelName: "blacklist-model"})
+					if len(pool) == 0 {
+						assert.Nil(t, channel)
+						require.NotNil(t, selectErr)
+						assert.True(t, selectErr.NoAvailableChannel)
+					} else {
+						require.Nil(t, selectErr)
+						require.NotNil(t, channel)
+						assert.False(t, loaded.ExcludesGroup(group))
+						RecordPoolRetryableFailure(ctx, "blacklist-model", group)
+						retryChannel, retryGroup, err := selectTokenRoutingPool(&RetryParam{Ctx: ctx, TokenGroup: "auto", ModelName: "blacklist-model", Retry: common.GetPointer(1)})
+						require.NoError(t, err)
+						require.NotNil(t, retryChannel)
+						assert.False(t, loaded.ExcludesGroup(retryGroup))
+					}
+					loaded.ExcludedGroups = nil
+					body, err = common.Marshal(loaded)
+					require.NoError(t, err)
+					stored.RoutingPolicy = string(body)
+					require.NoError(t, stored.Update())
+					require.NoError(t, db.First(&stored, token.Id).Error)
+					cleared, err := stored.GetRoutingPolicy()
+					require.NoError(t, err)
+					assert.Empty(t, cleared.ExcludedGroups)
+					assert.Len(t, ResolveRequestPool(makeContext(cleared, stored.Id), "blacklist-model"), 2)
+				}
+			}
+			policy := &model.TokenRoutingPolicy{Version: 1, Mode: "smart", ExcludedGroups: []string{"auto-b"}}
+			ctx := makeContext(policy, 900)
+			assert.Empty(t, ResolveRequestPool(ctx, "blacklist-image"), "the dedicated image group cannot bypass the blacklist or fall back")
+			assert.NotContains(t, AutomaticPoolModels(ctx), "gpt-6-astra")
+			assert.Contains(t, AutomaticPoolModels(ctx), "shared")
+			assert.Contains(t, AutomaticPoolModels(makeContext(&model.TokenRoutingPolicy{Version: 1, Mode: "smart"}, 901)), "gpt-6-astra", "another key is unaffected")
+			manual := &model.TokenRoutingPolicy{Version: 1, Mode: "custom", Groups: []model.TokenGroupEntry{{Group: "auto-b", Priority: 0}}, ExcludedGroups: []string{"auto-b"}}
+			assert.Equal(t, manual.Groups, ResolveRequestPool(makeContext(manual, 902), "blacklist-model"))
 		})
 	}
 }