Nelze vybrat více než 25 témat Téma musí začínat písmenem nebo číslem, může obsahovat pomlčky („-“) a může být dlouhé až 35 znaků.
 
 
 

140 řádky
4.4 KiB

  1. package model
  2. import (
  3. "testing"
  4. "time"
  5. "github.com/QuantumNous/new-api/common"
  6. "github.com/QuantumNous/new-api/constant"
  7. "github.com/QuantumNous/new-api/setting/ratio_setting"
  8. "github.com/glebarez/sqlite"
  9. "github.com/stretchr/testify/require"
  10. "gorm.io/gorm"
  11. )
  12. const testPricingApplyModel = "test-pricing-model-apply"
  13. const testPricingCallModel = "test-pricing-call-model"
  14. const testPricingGlobalModel = "test-pricing-global-model"
  15. func resetPricingRatios(t *testing.T) {
  16. t.Helper()
  17. require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{}`))
  18. require.NoError(t, ratio_setting.UpdateCompletionRatioByJSONString(`{}`))
  19. require.NoError(t, ratio_setting.UpdateCacheRatioByJSONString(`{}`))
  20. require.NoError(t, ratio_setting.UpdateCreateCacheRatioByJSONString(`{}`))
  21. require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(`{}`))
  22. }
  23. func setupPricingTest(t *testing.T) {
  24. t.Helper()
  25. resetPricingRatios(t)
  26. db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
  27. require.NoError(t, err)
  28. sqlDB, err := db.DB()
  29. require.NoError(t, err)
  30. sqlDB.SetMaxOpenConns(1)
  31. origDB := DB
  32. origUsingSQLite := common.UsingSQLite
  33. origPricingMap := pricingMap
  34. origPricingByModel := pricingByModel
  35. origVendorsList := vendorsList
  36. origSupportedEndpointMap := supportedEndpointMap
  37. origModelSupportEndpointTypes := modelSupportEndpointTypes
  38. origLastGetPricingTime := lastGetPricingTime
  39. DB = db
  40. common.UsingSQLite = true
  41. pricingMap = nil
  42. pricingByModel = nil
  43. vendorsList = nil
  44. supportedEndpointMap = nil
  45. modelSupportEndpointTypes = make(map[string][]constant.EndpointType)
  46. lastGetPricingTime = time.Time{}
  47. require.NoError(t, db.AutoMigrate(&Model{}, &Vendor{}, &Channel{}, &Ability{}))
  48. t.Cleanup(func() {
  49. DB = origDB
  50. common.UsingSQLite = origUsingSQLite
  51. pricingMap = origPricingMap
  52. pricingByModel = origPricingByModel
  53. vendorsList = origVendorsList
  54. supportedEndpointMap = origSupportedEndpointMap
  55. modelSupportEndpointTypes = origModelSupportEndpointTypes
  56. lastGetPricingTime = origLastGetPricingTime
  57. require.NoError(t, sqlDB.Close())
  58. resetPricingRatios(t)
  59. })
  60. }
  61. func TestApplyGlobalDefault_PerTokenModel(t *testing.T) {
  62. resetPricingRatios(t)
  63. require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{"`+testPricingApplyModel+`":10}`))
  64. require.NoError(t, ratio_setting.UpdateCompletionRatioByJSONString(`{"`+testPricingApplyModel+`":3}`))
  65. require.NoError(t, ratio_setting.UpdateCacheRatioByJSONString(`{"`+testPricingApplyModel+`":0.5}`))
  66. require.NoError(t, ratio_setting.UpdateCreateCacheRatioByJSONString(`{"`+testPricingApplyModel+`":0.75}`))
  67. p := &Pricing{}
  68. applyGlobalDefault(p, testPricingApplyModel)
  69. require.Equal(t, QuotaTypeByTokens, p.QuotaType)
  70. require.Equal(t, 10.0, p.ModelRatio)
  71. require.Equal(t, 3.0, p.CompletionRatio)
  72. require.Equal(t, 0.5, p.CacheRatio)
  73. require.Equal(t, 0.75, p.CacheCreationRatio)
  74. }
  75. func TestApplyGlobalDefault_PerCallModel(t *testing.T) {
  76. resetPricingRatios(t)
  77. require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(`{"`+testPricingCallModel+`":0.3}`))
  78. p := &Pricing{}
  79. applyGlobalDefault(p, testPricingCallModel)
  80. require.Equal(t, QuotaTypeByCall, p.QuotaType)
  81. require.Equal(t, 0.3, p.ModelPrice)
  82. require.Equal(t, 0.0, p.ModelRatio)
  83. }
  84. func TestGetPricingByModel_UsesGlobalDefaultOnly(t *testing.T) {
  85. setupPricingTest(t)
  86. require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{"`+testPricingGlobalModel+`":10}`))
  87. require.NoError(t, ratio_setting.UpdateCompletionRatioByJSONString(`{"`+testPricingGlobalModel+`":3}`))
  88. require.NoError(t, ratio_setting.UpdateCacheRatioByJSONString(`{"`+testPricingGlobalModel+`":0.5}`))
  89. require.NoError(t, ratio_setting.UpdateCreateCacheRatioByJSONString(`{"`+testPricingGlobalModel+`":0.75}`))
  90. require.NoError(t, (&Model{
  91. ModelName: testPricingGlobalModel,
  92. Status: 1,
  93. }).Insert())
  94. require.NoError(t, (&Vendor{
  95. Id: 1,
  96. Name: "test-vendor",
  97. SortOrder: 1,
  98. }).Insert())
  99. require.NoError(t, (&Channel{
  100. Id: 1,
  101. Type: constant.ChannelTypeOpenAI,
  102. Key: "k",
  103. Name: "channel-a",
  104. Status: common.ChannelStatusEnabled,
  105. Models: testPricingGlobalModel,
  106. Group: "default",
  107. }).Insert())
  108. RefreshPricing()
  109. pricing := GetPricingByModel(testPricingGlobalModel)
  110. require.NotNil(t, pricing)
  111. require.Equal(t, QuotaTypeByTokens, pricing.QuotaType)
  112. require.Equal(t, 10.0, pricing.ModelRatio)
  113. require.Equal(t, 3.0, pricing.CompletionRatio)
  114. require.Equal(t, 0.5, pricing.CacheRatio)
  115. require.Equal(t, 0.75, pricing.CacheCreationRatio)
  116. }