package model import ( "testing" "github.com/QuantumNous/new-api/setting/ratio_setting" "github.com/stretchr/testify/require" ) const testPricingModel = "test-pricing-model-apply" const testPricingCallModel = "test-pricing-call-model" func resetPricingRatios(t *testing.T) { t.Helper() 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}`)) p := &Pricing{} applyGlobalDefault(p, testPricingModel) 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 TestApplyGlobalDefault_PerCallModel(t *testing.T) { resetPricingRatios(t) require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(`{"`+testPricingCallModel+`":0.3}`)) p := &Pricing{} applyGlobalDefault(p, testPricingCallModel) require.Equal(t, QuotaTypeByCall, p.QuotaType) require.Equal(t, 0.3, p.ModelPrice) require.Equal(t, 0.0, p.ModelRatio) }