181 lines
9.0 KiB
Go
181 lines
9.0 KiB
Go
package service
|
|
|
|
import (
|
|
"fmt"
|
|
"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"
|
|
"gorm.io/gorm"
|
|
"gorm.io/gorm/schema"
|
|
"net/http/httptest"
|
|
"os"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestAutomaticPoolModelsDatabaseMatrix(t *testing.T) {
|
|
for _, dialect := range []string{"sqlite", "mysql", "postgres"} {
|
|
t.Run(dialect, func(t *testing.T) {
|
|
var driver gorm.Dialector
|
|
switch dialect {
|
|
case "sqlite":
|
|
driver = sqlite.Open(":memory:")
|
|
case "mysql":
|
|
if os.Getenv("TEST_MYSQL_DSN") == "" {
|
|
t.Skip("no isolated MySQL database")
|
|
}
|
|
driver = mysql.Open(os.Getenv("TEST_MYSQL_DSN"))
|
|
case "postgres":
|
|
if os.Getenv("TEST_POSTGRES_DSN") == "" {
|
|
t.Skip("no isolated PostgreSQL database")
|
|
}
|
|
driver = postgres.Open(os.Getenv("TEST_POSTGRES_DSN"))
|
|
}
|
|
db, err := gorm.Open(driver, &gorm.Config{NamingStrategy: schema.NamingStrategy{TablePrefix: fmt.Sprintf("automodel_%d_", time.Now().UnixNano())}})
|
|
require.NoError(t, err)
|
|
oldDB, oldRedis, oldMemory := model.DB, common.RedisEnabled, common.MemoryCacheEnabled
|
|
oldGroups, oldRatios := setting.UserUsableGroups2JSONString(), ratio_setting.GroupRatio2JSONString()
|
|
model.DB = db
|
|
common.RedisEnabled = false
|
|
common.MemoryCacheEnabled = false
|
|
require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.Ability{}, &model.Token{}))
|
|
t.Cleanup(func() {
|
|
require.NoError(t, db.Migrator().DropTable(&model.Token{}, &model.Ability{}, &model.Channel{}))
|
|
model.DB = oldDB
|
|
common.RedisEnabled = oldRedis
|
|
common.MemoryCacheEnabled = oldMemory
|
|
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(oldGroups))
|
|
require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(oldRatios))
|
|
conn, e := db.DB()
|
|
require.NoError(t, e)
|
|
require.NoError(t, conn.Close())
|
|
})
|
|
var version string
|
|
query := "select version()"
|
|
if dialect == "sqlite" {
|
|
query = "select sqlite_version()"
|
|
}
|
|
require.NoError(t, db.Raw(query).Scan(&version).Error)
|
|
t.Log("database version:", version)
|
|
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"auto-a":"","auto-b":""}`))
|
|
require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"auto-a":1,"auto-b":1,"private":1}`))
|
|
require.NoError(t, db.Create(&[]model.Channel{{Id: 561, Status: 1}, {Id: 562, Status: 1}, {Id: 563, Status: 2}, {Id: 564, Status: 1}}).Error)
|
|
require.NoError(t, db.Create(&[]model.Ability{
|
|
{ChannelId: 561, Group: "auto-a", Model: "model-a", Enabled: true},
|
|
{ChannelId: 561, Group: "auto-a", Model: "shared", Enabled: true},
|
|
{ChannelId: 562, Group: "auto-b", Model: "shared", Enabled: true},
|
|
{ChannelId: 562, Group: "auto-b", Model: "gpt-6-astra", Enabled: true},
|
|
{ChannelId: 563, Group: "auto-b", Model: "disabled-channel", Enabled: true},
|
|
{ChannelId: 562, Group: "auto-b", Model: "disabled-ability", Enabled: false},
|
|
{ChannelId: 564, Group: "private", Model: "unauthorized", Enabled: true},
|
|
}).Error)
|
|
for _, preference := range []string{"cost", "speed", "stable"} {
|
|
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
|
c.Request = httptest.NewRequest("GET", "/v1/models", nil)
|
|
common.SetContextKey(c, constant.ContextKeyUserGroup, "auto-a")
|
|
c.Set(model.TokenRoutingContextKey, &model.TokenRoutingPolicy{Version: 1, Mode: "smart", AutoPreference: preference})
|
|
c.Set("pool_failures", map[string]int{"auto-a": 2, "auto-b": 2})
|
|
require.Equal(t, []string{"gpt-6-astra", "model-a", "shared"}, AutomaticPoolModels(c))
|
|
}
|
|
require.NoError(t, db.Create(&model.Channel{Id: 565, Status: 1}).Error)
|
|
require.NoError(t, db.Create(&model.Ability{ChannelId: 565, Group: "new-group", Model: "new-model", Enabled: true}).Error)
|
|
require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"auto-a":"","auto-b":"","new-group":""}`))
|
|
require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"auto-a":1,"auto-b":1,"new-group":1}`))
|
|
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
|
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"))
|
|
})
|
|
}
|
|
}
|