No puede seleccionar más de 25 temas Los temas deben comenzar con una letra o número, pueden incluir guiones ('-') y pueden tener hasta 35 caracteres de largo.
 
 
 

165 líneas
5.3 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. t.Cleanup(func() { resetPricingRatios(t) })
  64. require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{"`+testPricingApplyModel+`":10}`))
  65. require.NoError(t, ratio_setting.UpdateCompletionRatioByJSONString(`{"`+testPricingApplyModel+`":3}`))
  66. require.NoError(t, ratio_setting.UpdateCacheRatioByJSONString(`{"`+testPricingApplyModel+`":0.5}`))
  67. require.NoError(t, ratio_setting.UpdateCreateCacheRatioByJSONString(`{"`+testPricingApplyModel+`":0.75}`))
  68. p := &Pricing{}
  69. applyGlobalDefault(p, testPricingApplyModel)
  70. require.Equal(t, QuotaTypeByTokens, p.QuotaType)
  71. require.Equal(t, 10.0, p.ModelRatio)
  72. require.Equal(t, 3.0, p.CompletionRatio)
  73. require.Equal(t, 0.5, p.CacheRatio)
  74. require.Equal(t, 0.75, p.CacheCreationRatio)
  75. }
  76. func TestApplyGlobalDefault_PerCallModel(t *testing.T) {
  77. resetPricingRatios(t)
  78. t.Cleanup(func() { resetPricingRatios(t) })
  79. require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(`{"`+testPricingCallModel+`":0.3}`))
  80. p := &Pricing{}
  81. applyGlobalDefault(p, testPricingCallModel)
  82. require.Equal(t, QuotaTypeByCall, p.QuotaType)
  83. require.Equal(t, 0.3, p.ModelPrice)
  84. require.Equal(t, 0.0, p.ModelRatio)
  85. }
  86. func TestGetPricingByModel_UsesGlobalDefaultOnly(t *testing.T) {
  87. setupPricingTest(t)
  88. require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{"`+testPricingGlobalModel+`":10}`))
  89. require.NoError(t, ratio_setting.UpdateCompletionRatioByJSONString(`{"`+testPricingGlobalModel+`":3}`))
  90. require.NoError(t, ratio_setting.UpdateCacheRatioByJSONString(`{"`+testPricingGlobalModel+`":0.5}`))
  91. require.NoError(t, ratio_setting.UpdateCreateCacheRatioByJSONString(`{"`+testPricingGlobalModel+`":0.75}`))
  92. require.NoError(t, (&Model{
  93. ModelName: testPricingGlobalModel,
  94. Status: 1,
  95. }).Insert())
  96. require.NoError(t, (&Vendor{
  97. Id: 1,
  98. Name: "test-vendor",
  99. SortOrder: 1,
  100. }).Insert())
  101. require.NoError(t, (&Channel{
  102. Id: 1,
  103. Type: constant.ChannelTypeOpenAI,
  104. Key: "k",
  105. Name: "channel-a",
  106. Status: common.ChannelStatusEnabled,
  107. Models: testPricingGlobalModel,
  108. Group: "default",
  109. }).Insert())
  110. RefreshPricing()
  111. pricing := GetPricingByModel(testPricingGlobalModel)
  112. require.NotNil(t, pricing)
  113. require.Equal(t, QuotaTypeByTokens, pricing.QuotaType)
  114. require.Equal(t, 10.0, pricing.ModelRatio)
  115. require.Equal(t, 3.0, pricing.CompletionRatio)
  116. require.Equal(t, 0.5, pricing.CacheRatio)
  117. require.Equal(t, 0.75, pricing.CacheCreationRatio)
  118. }
  119. func TestParseModelEndpointTypes_SupportsObjectAndArray(t *testing.T) {
  120. objectEndpoints := parseModelEndpointTypes(`{
  121. "anthropic": {"path": "/v1/messages", "method": "POST"},
  122. "openai": "/v1/chat/completions"
  123. }`)
  124. require.ElementsMatch(t, []string{"anthropic", "openai"}, objectEndpoints)
  125. arrayEndpoints := parseModelEndpointTypes(`["anthropic", "openai", "anthropic"]`)
  126. require.Equal(t, []string{"anthropic", "openai"}, arrayEndpoints)
  127. }
  128. func TestParseModelEndpointTypes_IgnoresInvalidJSON(t *testing.T) {
  129. require.Nil(t, parseModelEndpointTypes(`123`))
  130. require.Nil(t, parseModelEndpointTypes(`{"anthropic": true}`))
  131. require.Nil(t, parseModelEndpointTypes(``))
  132. info, ok := common.GetDefaultEndpointInfo("anthropic")
  133. require.True(t, ok)
  134. require.Equal(t, "/v1/messages", info.Path)
  135. }