|
- package model
-
- import (
- "testing"
- "time"
-
- "github.com/QuantumNous/new-api/common"
- "github.com/QuantumNous/new-api/constant"
- "github.com/QuantumNous/new-api/setting/ratio_setting"
- "github.com/glebarez/sqlite"
- "github.com/stretchr/testify/require"
- "gorm.io/gorm"
- )
-
- const testPricingApplyModel = "test-pricing-model-apply"
- const testPricingCallModel = "test-pricing-call-model"
- const testPricingGlobalModel = "test-pricing-global-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 setupPricingTest(t *testing.T) {
- t.Helper()
-
- resetPricingRatios(t)
-
- 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 := DB
- origUsingSQLite := common.UsingSQLite
- origPricingMap := pricingMap
- origPricingByModel := pricingByModel
- origVendorsList := vendorsList
- origSupportedEndpointMap := supportedEndpointMap
- origModelSupportEndpointTypes := modelSupportEndpointTypes
- origLastGetPricingTime := lastGetPricingTime
-
- DB = db
- common.UsingSQLite = true
- pricingMap = nil
- pricingByModel = nil
- vendorsList = nil
- supportedEndpointMap = nil
- modelSupportEndpointTypes = make(map[string][]constant.EndpointType)
- lastGetPricingTime = time.Time{}
-
- require.NoError(t, db.AutoMigrate(&Model{}, &Vendor{}, &Channel{}, &Ability{}))
-
- t.Cleanup(func() {
- DB = origDB
- common.UsingSQLite = origUsingSQLite
- pricingMap = origPricingMap
- pricingByModel = origPricingByModel
- vendorsList = origVendorsList
- supportedEndpointMap = origSupportedEndpointMap
- modelSupportEndpointTypes = origModelSupportEndpointTypes
- lastGetPricingTime = origLastGetPricingTime
- require.NoError(t, sqlDB.Close())
- resetPricingRatios(t)
- })
- }
-
- func TestApplyGlobalDefault_PerTokenModel(t *testing.T) {
- resetPricingRatios(t)
- t.Cleanup(func() { resetPricingRatios(t) })
-
- require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{"`+testPricingApplyModel+`":10}`))
- require.NoError(t, ratio_setting.UpdateCompletionRatioByJSONString(`{"`+testPricingApplyModel+`":3}`))
- require.NoError(t, ratio_setting.UpdateCacheRatioByJSONString(`{"`+testPricingApplyModel+`":0.5}`))
- require.NoError(t, ratio_setting.UpdateCreateCacheRatioByJSONString(`{"`+testPricingApplyModel+`":0.75}`))
-
- p := &Pricing{}
- applyGlobalDefault(p, testPricingApplyModel)
-
- 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)
- t.Cleanup(func() { 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)
- }
-
- func TestGetPricingByModel_UsesGlobalDefaultOnly(t *testing.T) {
- setupPricingTest(t)
-
- require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{"`+testPricingGlobalModel+`":10}`))
- require.NoError(t, ratio_setting.UpdateCompletionRatioByJSONString(`{"`+testPricingGlobalModel+`":3}`))
- require.NoError(t, ratio_setting.UpdateCacheRatioByJSONString(`{"`+testPricingGlobalModel+`":0.5}`))
- require.NoError(t, ratio_setting.UpdateCreateCacheRatioByJSONString(`{"`+testPricingGlobalModel+`":0.75}`))
-
- require.NoError(t, (&Model{
- ModelName: testPricingGlobalModel,
- Status: 1,
- }).Insert())
- require.NoError(t, (&Vendor{
- Id: 1,
- Name: "test-vendor",
- SortOrder: 1,
- }).Insert())
- require.NoError(t, (&Channel{
- Id: 1,
- Type: constant.ChannelTypeOpenAI,
- Key: "k",
- Name: "channel-a",
- Status: common.ChannelStatusEnabled,
- Models: testPricingGlobalModel,
- Group: "default",
- }).Insert())
-
- RefreshPricing()
- pricing := GetPricingByModel(testPricingGlobalModel)
-
- require.NotNil(t, pricing)
- require.Equal(t, QuotaTypeByTokens, pricing.QuotaType)
- require.Equal(t, 10.0, pricing.ModelRatio)
- require.Equal(t, 3.0, pricing.CompletionRatio)
- require.Equal(t, 0.5, pricing.CacheRatio)
- require.Equal(t, 0.75, pricing.CacheCreationRatio)
- }
-
- func TestParseModelEndpointTypes_SupportsObjectAndArray(t *testing.T) {
- objectEndpoints := parseModelEndpointTypes(`{
- "anthropic": {"path": "/v1/messages", "method": "POST"},
- "openai": "/v1/chat/completions"
- }`)
- require.ElementsMatch(t, []string{"anthropic", "openai"}, objectEndpoints)
-
- arrayEndpoints := parseModelEndpointTypes(`["anthropic", "openai", "anthropic"]`)
- require.Equal(t, []string{"anthropic", "openai"}, arrayEndpoints)
- }
-
- func TestParseModelEndpointTypes_IgnoresInvalidJSON(t *testing.T) {
- require.Nil(t, parseModelEndpointTypes(`123`))
- require.Nil(t, parseModelEndpointTypes(`{"anthropic": true}`))
- require.Nil(t, parseModelEndpointTypes(``))
-
- info, ok := common.GetDefaultEndpointInfo("anthropic")
- require.True(t, ok)
- require.Equal(t, "/v1/messages", info.Path)
- }
|