|
- 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}`))
- }
|