|
|
|
@@ -4,119 +4,46 @@ import ( |
|
|
|
"testing" |
|
|
|
|
|
|
|
"github.com/QuantumNous/new-api/setting/ratio_setting" |
|
|
|
"github.com/glebarez/sqlite" |
|
|
|
"github.com/stretchr/testify/require" |
|
|
|
"gorm.io/gorm" |
|
|
|
) |
|
|
|
|
|
|
|
const testPricingModel = "test-pricing-model-apply" |
|
|
|
const testPricingCallModel = "test-pricing-call-model" |
|
|
|
|
|
|
|
func setupPricingTest(t *testing.T) { |
|
|
|
func resetPricingRatios(t *testing.T) { |
|
|
|
t.Helper() |
|
|
|
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) |
|
|
|
require.NoError(t, err) |
|
|
|
sqlDB, _ := db.DB() |
|
|
|
sqlDB.SetMaxOpenConns(1) |
|
|
|
|
|
|
|
origDB := DB |
|
|
|
DB = db |
|
|
|
require.NoError(t, db.AutoMigrate(&ChannelPricing{})) |
|
|
|
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.UpdateModelPriceByJSONString(`{}`)) |
|
|
|
} |
|
|
|
|
|
|
|
// 全局默认定价 |
|
|
|
func TestApplyGlobalDefault_PerTokenModel(t *testing.T) { |
|
|
|
resetPricingRatios(t) |
|
|
|
require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{"`+testPricingModel+`":10}`)) |
|
|
|
require.NoError(t, ratio_setting.UpdateCompletionRatioByJSONString(`{"`+testPricingModel+`":3}`)) |
|
|
|
require.NoError(t, ratio_setting.UpdateCacheRatioByJSONString(`{"`+testPricingModel+`":0.5}`)) |
|
|
|
require.NoError(t, ratio_setting.UpdateCreateCacheRatioByJSONString(`{"`+testPricingModel+`":0.75}`)) |
|
|
|
|
|
|
|
t.Cleanup(func() { |
|
|
|
DB = origDB |
|
|
|
sqlDB.Close() |
|
|
|
ratio_setting.UpdateModelRatioByJSONString(`{}`) |
|
|
|
ratio_setting.UpdateCompletionRatioByJSONString(`{}`) |
|
|
|
ratio_setting.UpdateCacheRatioByJSONString(`{}`) |
|
|
|
ratio_setting.UpdateCreateCacheRatioByJSONString(`{}`) |
|
|
|
ratio_setting.UpdateModelPriceByJSONString(`{}`) |
|
|
|
}) |
|
|
|
} |
|
|
|
|
|
|
|
func TestApplyBestChannelPricing_NoChannelPricing(t *testing.T) { |
|
|
|
setupPricingTest(t) |
|
|
|
|
|
|
|
p := &Pricing{} |
|
|
|
applyBestChannelPricing(p, nil, testPricingModel) |
|
|
|
|
|
|
|
require.Equal(t, 0, p.QuotaType, "应使用全局按量计费") |
|
|
|
require.Equal(t, 10.0, p.ModelRatio, "应使用全局 model_ratio") |
|
|
|
require.Equal(t, 3.0, p.CompletionRatio, "应使用全局 completion_ratio") |
|
|
|
require.Equal(t, 0.5, p.CacheRatio, "应使用全局 cache_ratio") |
|
|
|
require.Equal(t, 0.75, p.CacheCreationRatio, "应使用全局 cache_creation_ratio") |
|
|
|
} |
|
|
|
|
|
|
|
func TestApplyBestChannelPricing_PerTokenCheapest(t *testing.T) { |
|
|
|
setupPricingTest(t) |
|
|
|
|
|
|
|
// 渠道A: model_ratio=8(更便宜) |
|
|
|
require.NoError(t, (&ChannelPricing{ |
|
|
|
ModelName: testPricingModel, ChannelId: 1, |
|
|
|
QuotaType: QuotaTypeByTokens, ModelRatio: 8, CompletionRatio: 2, |
|
|
|
CacheRatio: 0.3, CacheCreationRatio: 0.6, |
|
|
|
}).Insert()) |
|
|
|
// 渠道B: model_ratio=12(更贵) |
|
|
|
require.NoError(t, (&ChannelPricing{ |
|
|
|
ModelName: testPricingModel, ChannelId: 2, |
|
|
|
QuotaType: QuotaTypeByTokens, ModelRatio: 12, CompletionRatio: 4, |
|
|
|
CacheRatio: 0.8, CacheCreationRatio: 1.0, |
|
|
|
}).Insert()) |
|
|
|
|
|
|
|
cps := []ChannelPricing{ |
|
|
|
{ModelRatio: 8, CompletionRatio: 2, CacheRatio: 0.3, CacheCreationRatio: 0.6}, |
|
|
|
{ModelRatio: 12, CompletionRatio: 4, CacheRatio: 0.8, CacheCreationRatio: 1.0}, |
|
|
|
} |
|
|
|
|
|
|
|
p := &Pricing{} |
|
|
|
applyBestChannelPricing(p, cps, testPricingModel) |
|
|
|
|
|
|
|
require.Equal(t, 0, p.QuotaType) |
|
|
|
require.Equal(t, 8.0, p.ModelRatio, "应选最便宜的渠道A") |
|
|
|
require.Equal(t, 2.0, p.CompletionRatio) |
|
|
|
require.Equal(t, 0.3, p.CacheRatio) |
|
|
|
require.Equal(t, 0.6, p.CacheCreationRatio) |
|
|
|
} |
|
|
|
|
|
|
|
func TestApplyBestChannelPricing_ExtendedRatioZeroFallback(t *testing.T) { |
|
|
|
setupPricingTest(t) |
|
|
|
|
|
|
|
// 渠道定价中扩展比率为 0,应回退到全局值 |
|
|
|
require.NoError(t, (&ChannelPricing{ |
|
|
|
ModelName: testPricingModel, ChannelId: 1, |
|
|
|
QuotaType: QuotaTypeByTokens, ModelRatio: 5, CompletionRatio: 1, |
|
|
|
CacheRatio: 0, CacheCreationRatio: 0, // 0 = 未设置 |
|
|
|
}).Insert()) |
|
|
|
|
|
|
|
cps := []ChannelPricing{ |
|
|
|
{ModelRatio: 5, CompletionRatio: 1, CacheRatio: 0, CacheCreationRatio: 0}, |
|
|
|
} |
|
|
|
|
|
|
|
p := &Pricing{} |
|
|
|
applyBestChannelPricing(p, cps, testPricingModel) |
|
|
|
applyGlobalDefault(p, testPricingModel) |
|
|
|
|
|
|
|
require.Equal(t, 0.5, p.CacheRatio, "cache_ratio=0 应回退全局 0.5") |
|
|
|
require.Equal(t, 0.75, p.CacheCreationRatio, "cache_creation_ratio=0 应回退全局 0.75") |
|
|
|
require.Equal(t, QuotaTypeByTokens, p.QuotaType) |
|
|
|
require.Equal(t, 10.0, p.ModelRatio) |
|
|
|
require.Equal(t, 3.0, p.CompletionRatio) |
|
|
|
require.Equal(t, 0.5, p.CacheRatio) |
|
|
|
require.Equal(t, 0.75, p.CacheCreationRatio) |
|
|
|
} |
|
|
|
|
|
|
|
func TestApplyBestChannelPricing_PerCallCheapest(t *testing.T) { |
|
|
|
setupPricingTest(t) |
|
|
|
require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(`{"`+testPricingModel+`":0.8}`)) |
|
|
|
|
|
|
|
// 只有按次计费的渠道 |
|
|
|
cps := []ChannelPricing{ |
|
|
|
{QuotaType: QuotaTypeByCall, ModelPrice: 0.3}, |
|
|
|
{QuotaType: QuotaTypeByCall, ModelPrice: 0.5}, |
|
|
|
} |
|
|
|
func TestApplyGlobalDefault_PerCallModel(t *testing.T) { |
|
|
|
resetPricingRatios(t) |
|
|
|
require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(`{"`+testPricingCallModel+`":0.3}`)) |
|
|
|
|
|
|
|
p := &Pricing{} |
|
|
|
applyBestChannelPricing(p, cps, testPricingModel) |
|
|
|
applyGlobalDefault(p, testPricingCallModel) |
|
|
|
|
|
|
|
require.Equal(t, 1, p.QuotaType) |
|
|
|
require.Equal(t, 0.3, p.ModelPrice, "应选最便宜的按次渠道") |
|
|
|
require.Equal(t, QuotaTypeByCall, p.QuotaType) |
|
|
|
require.Equal(t, 0.3, p.ModelPrice) |
|
|
|
require.Equal(t, 0.0, p.ModelRatio) |
|
|
|
} |