package helper import ( "net/http" "net/http/httptest" "testing" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/dto" "github.com/QuantumNous/new-api/model" relaycommon "github.com/QuantumNous/new-api/relay/common" "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/glebarez/sqlite" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "gorm.io/gorm" ) const testModelUsePriceSwitch = "test-useprice-switch-model" func setupUsePriceSwitchTest(t *testing.T) { t.Helper() db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) require.NoError(t, err) sqlDB, err := db.DB() require.NoError(t, err) sqlDB.SetMaxOpenConns(1) origDB := model.DB origUsingSQLite := common.UsingSQLite origRedisEnabled := common.RedisEnabled model.DB = db common.UsingSQLite = true common.RedisEnabled = false require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(`{"`+testModelUsePriceSwitch+`":0.5}`)) require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{"`+testModelUsePriceSwitch+`":15}`)) require.NoError(t, ratio_setting.UpdateCompletionRatioByJSONString(`{"`+testModelUsePriceSwitch+`":3}`)) require.NoError(t, ratio_setting.UpdateCacheRatioByJSONString(`{"`+testModelUsePriceSwitch+`":0.1}`)) require.NoError(t, ratio_setting.UpdateCreateCacheRatioByJSONString(`{"`+testModelUsePriceSwitch+`":1.25}`)) require.NoError(t, ratio_setting.UpdateImageRatioByJSONString(`{"`+testModelUsePriceSwitch+`":2.0}`)) require.NoError(t, ratio_setting.UpdateAudioRatioByJSONString(`{"`+testModelUsePriceSwitch+`":5.0}`)) require.NoError(t, ratio_setting.UpdateAudioCompletionRatioByJSONString(`{"`+testModelUsePriceSwitch+`":3.0}`)) t.Cleanup(func() { model.DB = origDB common.UsingSQLite = origUsingSQLite common.RedisEnabled = origRedisEnabled require.NoError(t, sqlDB.Close()) require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(`{}`)) require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{}`)) require.NoError(t, ratio_setting.UpdateCompletionRatioByJSONString(`{}`)) require.NoError(t, ratio_setting.UpdateCacheRatioByJSONString(`{}`)) require.NoError(t, ratio_setting.UpdateCreateCacheRatioByJSONString(`{}`)) require.NoError(t, ratio_setting.UpdateImageRatioByJSONString(`{}`)) require.NoError(t, ratio_setting.UpdateAudioRatioByJSONString(`{}`)) require.NoError(t, ratio_setting.UpdateAudioCompletionRatioByJSONString(`{}`)) }) } func buildTestContext(t *testing.T) *gin.Context { t.Helper() gin.SetMode(gin.TestMode) w := httptest.NewRecorder() c, _ := gin.CreateTestContext(w) req, err := http.NewRequest(http.MethodPost, "/v1/chat/completions", nil) require.NoError(t, err) c.Request = req return c } func TestGlobalRatiosFilledWhenUsePriceTrue(t *testing.T) { setupUsePriceSwitchTest(t) c := buildTestContext(t) info := &relaycommon.RelayInfo{ OriginModelName: testModelUsePriceSwitch, UserSetting: dto.UserSetting{}, } priceData, err := ModelPriceHelper(c, info, 100, &types.TokenCountMeta{}) require.NoError(t, err) assert.True(t, priceData.UsePrice) assert.Equal(t, 0.5, priceData.ModelPrice) assert.Equal(t, 0.1, priceData.CacheRatio) assert.Equal(t, 1.25, priceData.CacheCreationRatio) assert.Equal(t, 2.0, priceData.ImageRatio) assert.Equal(t, 5.0, priceData.AudioRatio) assert.Equal(t, 3.0, priceData.AudioCompletionRatio) } func TestGlobalRatiosFilledWhenUseRatioMode(t *testing.T) { setupUsePriceSwitchTest(t) c := buildTestContext(t) require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(`{}`)) info := &relaycommon.RelayInfo{ OriginModelName: testModelUsePriceSwitch, UserSetting: dto.UserSetting{}, } priceData, err := ModelPriceHelper(c, info, 100, &types.TokenCountMeta{}) require.NoError(t, err) assert.False(t, priceData.UsePrice) assert.Equal(t, 15.0, priceData.ModelRatio) assert.Equal(t, 3.0, priceData.CompletionRatio) assert.Equal(t, 0.1, priceData.CacheRatio) assert.Equal(t, 2.0, priceData.ImageRatio) } func TestHandleGroupRatio_DefaultGroup(t *testing.T) { setupUsePriceSwitchTest(t) c := buildTestContext(t) info := &relaycommon.RelayInfo{ UsingGroup: "default", UserGroup: "default", } groupRatioInfo := HandleGroupRatio(c, info) assert.Equal(t, 1.0, groupRatioInfo.GroupRatio) assert.False(t, groupRatioInfo.HasSpecialRatio) assert.Equal(t, -1.0, groupRatioInfo.GroupSpecialRatio) } func TestHandleGroupRatio_CustomGroupRatio(t *testing.T) { setupUsePriceSwitchTest(t) require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1,"svip":1,"premium":0.7}`)) c := buildTestContext(t) info := &relaycommon.RelayInfo{ UsingGroup: "premium", UserGroup: "default", } groupRatioInfo := HandleGroupRatio(c, info) assert.Equal(t, 0.7, groupRatioInfo.GroupRatio) assert.False(t, groupRatioInfo.HasSpecialRatio) // cleanup require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1,"svip":1}`)) } func TestHandleGroupRatio_SpecialGroupGroupRatio(t *testing.T) { setupUsePriceSwitchTest(t) require.NoError(t, ratio_setting.UpdateGroupGroupRatioByJSONString(`{"vip":{"premium":0.8}}`)) c := buildTestContext(t) info := &relaycommon.RelayInfo{ UsingGroup: "premium", UserGroup: "vip", } groupRatioInfo := HandleGroupRatio(c, info) assert.Equal(t, 0.8, groupRatioInfo.GroupRatio) assert.Equal(t, 0.8, groupRatioInfo.GroupSpecialRatio) assert.True(t, groupRatioInfo.HasSpecialRatio) // cleanup require.NoError(t, ratio_setting.UpdateGroupGroupRatioByJSONString(`{}`)) } func TestHandleGroupRatio_AutoGroupOverridesUsingGroup(t *testing.T) { setupUsePriceSwitchTest(t) require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1,"svip":1,"vip-channel":0.6}`)) c := buildTestContext(t) // Simulate: distributor set auto_group = "vip-channel" in context c.Set("auto_group", "vip-channel") info := &relaycommon.RelayInfo{ UsingGroup: "default", // will be overridden UserGroup: "default", } groupRatioInfo := HandleGroupRatio(c, info) // UsingGroup should be updated to "vip-channel" assert.Equal(t, "vip-channel", info.UsingGroup) // GroupRatio should reflect vip-channel's ratio assert.Equal(t, 0.6, groupRatioInfo.GroupRatio) // cleanup require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1,"svip":1}`)) } func TestHandleGroupRatio_AutoGroupWithSpecialRatio(t *testing.T) { setupUsePriceSwitchTest(t) require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1,"svip":1,"premium":1.0}`)) require.NoError(t, ratio_setting.UpdateGroupGroupRatioByJSONString(`{"vip":{"premium":0.5}}`)) c := buildTestContext(t) c.Set("auto_group", "premium") info := &relaycommon.RelayInfo{ UsingGroup: "default", UserGroup: "vip", } groupRatioInfo := HandleGroupRatio(c, info) assert.Equal(t, "premium", info.UsingGroup) // Should use the special ratio (0.5) from vip->premium, not the generic group ratio (1.0) assert.Equal(t, 0.5, groupRatioInfo.GroupRatio) assert.True(t, groupRatioInfo.HasSpecialRatio) // cleanup require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1,"svip":1}`)) require.NoError(t, ratio_setting.UpdateGroupGroupRatioByJSONString(`{}`)) } // --- ModelPriceHelper with group ratio --- func TestModelPriceHelper_GroupRatioAffectsPreConsume(t *testing.T) { setupUsePriceSwitchTest(t) require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1,"svip":1,"discount":0.5}`)) c := buildTestContext(t) info := &relaycommon.RelayInfo{ OriginModelName: testModelUsePriceSwitch, UsingGroup: "discount", UserGroup: "default", UserSetting: dto.UserSetting{}, } priceData, err := ModelPriceHelper(c, info, 100, &types.TokenCountMeta{}) require.NoError(t, err) // model price 0.5, group ratio 0.5 => preConsume = 0.5 * QuotaPerUnit * 0.5 assert.Equal(t, 0.5, priceData.GroupRatioInfo.GroupRatio) assert.True(t, priceData.UsePrice) // 0.5 (modelPrice) * 500000 (QuotaPerUnit) * 0.5 (groupRatio) = 125000 expectedQuota := int(0.5 * float64(common.QuotaPerUnit) * 0.5) assert.Equal(t, expectedQuota, priceData.QuotaToPreConsume) // cleanup require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1,"svip":1}`)) } func TestModelPriceHelper_RatioModeWithGroupRatio(t *testing.T) { setupUsePriceSwitchTest(t) // Remove model price so it falls back to ratio mode require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(`{}`)) require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1,"svip":1,"double":2.0}`)) c := buildTestContext(t) info := &relaycommon.RelayInfo{ OriginModelName: testModelUsePriceSwitch, UsingGroup: "double", UserGroup: "default", UserSetting: dto.UserSetting{}, } priceData, err := ModelPriceHelper(c, info, 100, &types.TokenCountMeta{}) require.NoError(t, err) assert.False(t, priceData.UsePrice) assert.Equal(t, 15.0, priceData.ModelRatio) assert.Equal(t, 2.0, priceData.GroupRatioInfo.GroupRatio) // ratio = modelRatio(15) * groupRatio(2.0) = 30 // preConsumedQuota = int(preConsumedTokens * 30) expectedRatio := 15.0 * 2.0 assert.Equal(t, expectedRatio, priceData.ModelRatio*priceData.GroupRatioInfo.GroupRatio) // cleanup require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1,"svip":1}`)) } func TestModelPriceHelper_FreeGroup(t *testing.T) { setupUsePriceSwitchTest(t) require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1,"svip":1,"free":0}`)) // Disable pre-consume for free models so freeModel flag is set origSetting := operation_setting.GetQuotaSetting() operation_setting.GetQuotaSetting().EnableFreeModelPreConsume = false t.Cleanup(func() { operation_setting.GetQuotaSetting().EnableFreeModelPreConsume = origSetting.EnableFreeModelPreConsume }) c := buildTestContext(t) info := &relaycommon.RelayInfo{ OriginModelName: testModelUsePriceSwitch, UsingGroup: "free", UserGroup: "default", UserSetting: dto.UserSetting{}, } priceData, err := ModelPriceHelper(c, info, 100, &types.TokenCountMeta{}) require.NoError(t, err) // group ratio = 0 means free model assert.Equal(t, 0.0, priceData.GroupRatioInfo.GroupRatio) assert.True(t, priceData.FreeModel) assert.Equal(t, 0, priceData.QuotaToPreConsume) // cleanup require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1,"svip":1}`)) }