You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 

308 regels
10 KiB

  1. package helper
  2. import (
  3. "net/http"
  4. "net/http/httptest"
  5. "testing"
  6. "github.com/QuantumNous/new-api/common"
  7. "github.com/QuantumNous/new-api/dto"
  8. "github.com/QuantumNous/new-api/model"
  9. relaycommon "github.com/QuantumNous/new-api/relay/common"
  10. "github.com/QuantumNous/new-api/setting/operation_setting"
  11. "github.com/QuantumNous/new-api/setting/ratio_setting"
  12. "github.com/QuantumNous/new-api/types"
  13. "github.com/gin-gonic/gin"
  14. "github.com/glebarez/sqlite"
  15. "github.com/stretchr/testify/assert"
  16. "github.com/stretchr/testify/require"
  17. "gorm.io/gorm"
  18. )
  19. const testModelUsePriceSwitch = "test-useprice-switch-model"
  20. func setupUsePriceSwitchTest(t *testing.T) {
  21. t.Helper()
  22. db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{})
  23. require.NoError(t, err)
  24. sqlDB, err := db.DB()
  25. require.NoError(t, err)
  26. sqlDB.SetMaxOpenConns(1)
  27. origDB := model.DB
  28. origUsingSQLite := common.UsingSQLite
  29. origRedisEnabled := common.RedisEnabled
  30. model.DB = db
  31. common.UsingSQLite = true
  32. common.RedisEnabled = false
  33. require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(`{"`+testModelUsePriceSwitch+`":0.5}`))
  34. require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{"`+testModelUsePriceSwitch+`":15}`))
  35. require.NoError(t, ratio_setting.UpdateCompletionRatioByJSONString(`{"`+testModelUsePriceSwitch+`":3}`))
  36. require.NoError(t, ratio_setting.UpdateCacheRatioByJSONString(`{"`+testModelUsePriceSwitch+`":0.1}`))
  37. require.NoError(t, ratio_setting.UpdateCreateCacheRatioByJSONString(`{"`+testModelUsePriceSwitch+`":1.25}`))
  38. require.NoError(t, ratio_setting.UpdateImageRatioByJSONString(`{"`+testModelUsePriceSwitch+`":2.0}`))
  39. require.NoError(t, ratio_setting.UpdateAudioRatioByJSONString(`{"`+testModelUsePriceSwitch+`":5.0}`))
  40. require.NoError(t, ratio_setting.UpdateAudioCompletionRatioByJSONString(`{"`+testModelUsePriceSwitch+`":3.0}`))
  41. t.Cleanup(func() {
  42. model.DB = origDB
  43. common.UsingSQLite = origUsingSQLite
  44. common.RedisEnabled = origRedisEnabled
  45. require.NoError(t, sqlDB.Close())
  46. require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(`{}`))
  47. require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{}`))
  48. require.NoError(t, ratio_setting.UpdateCompletionRatioByJSONString(`{}`))
  49. require.NoError(t, ratio_setting.UpdateCacheRatioByJSONString(`{}`))
  50. require.NoError(t, ratio_setting.UpdateCreateCacheRatioByJSONString(`{}`))
  51. require.NoError(t, ratio_setting.UpdateImageRatioByJSONString(`{}`))
  52. require.NoError(t, ratio_setting.UpdateAudioRatioByJSONString(`{}`))
  53. require.NoError(t, ratio_setting.UpdateAudioCompletionRatioByJSONString(`{}`))
  54. })
  55. }
  56. func buildTestContext(t *testing.T) *gin.Context {
  57. t.Helper()
  58. gin.SetMode(gin.TestMode)
  59. w := httptest.NewRecorder()
  60. c, _ := gin.CreateTestContext(w)
  61. req, err := http.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
  62. require.NoError(t, err)
  63. c.Request = req
  64. return c
  65. }
  66. func TestGlobalRatiosFilledWhenUsePriceTrue(t *testing.T) {
  67. setupUsePriceSwitchTest(t)
  68. c := buildTestContext(t)
  69. info := &relaycommon.RelayInfo{
  70. OriginModelName: testModelUsePriceSwitch,
  71. UserSetting: dto.UserSetting{},
  72. }
  73. priceData, err := ModelPriceHelper(c, info, 100, &types.TokenCountMeta{})
  74. require.NoError(t, err)
  75. assert.True(t, priceData.UsePrice)
  76. assert.Equal(t, 0.5, priceData.ModelPrice)
  77. assert.Equal(t, 0.1, priceData.CacheRatio)
  78. assert.Equal(t, 1.25, priceData.CacheCreationRatio)
  79. assert.Equal(t, 2.0, priceData.ImageRatio)
  80. assert.Equal(t, 5.0, priceData.AudioRatio)
  81. assert.Equal(t, 3.0, priceData.AudioCompletionRatio)
  82. }
  83. func TestGlobalRatiosFilledWhenUseRatioMode(t *testing.T) {
  84. setupUsePriceSwitchTest(t)
  85. c := buildTestContext(t)
  86. require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(`{}`))
  87. info := &relaycommon.RelayInfo{
  88. OriginModelName: testModelUsePriceSwitch,
  89. UserSetting: dto.UserSetting{},
  90. }
  91. priceData, err := ModelPriceHelper(c, info, 100, &types.TokenCountMeta{})
  92. require.NoError(t, err)
  93. assert.False(t, priceData.UsePrice)
  94. assert.Equal(t, 15.0, priceData.ModelRatio)
  95. assert.Equal(t, 3.0, priceData.CompletionRatio)
  96. assert.Equal(t, 0.1, priceData.CacheRatio)
  97. assert.Equal(t, 2.0, priceData.ImageRatio)
  98. }
  99. func TestHandleGroupRatio_DefaultGroup(t *testing.T) {
  100. setupUsePriceSwitchTest(t)
  101. c := buildTestContext(t)
  102. info := &relaycommon.RelayInfo{
  103. UsingGroup: "default",
  104. UserGroup: "default",
  105. }
  106. groupRatioInfo := HandleGroupRatio(c, info)
  107. assert.Equal(t, 1.0, groupRatioInfo.GroupRatio)
  108. assert.False(t, groupRatioInfo.HasSpecialRatio)
  109. assert.Equal(t, -1.0, groupRatioInfo.GroupSpecialRatio)
  110. }
  111. func TestHandleGroupRatio_CustomGroupRatio(t *testing.T) {
  112. setupUsePriceSwitchTest(t)
  113. require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1,"svip":1,"premium":0.7}`))
  114. c := buildTestContext(t)
  115. info := &relaycommon.RelayInfo{
  116. UsingGroup: "premium",
  117. UserGroup: "default",
  118. }
  119. groupRatioInfo := HandleGroupRatio(c, info)
  120. assert.Equal(t, 0.7, groupRatioInfo.GroupRatio)
  121. assert.False(t, groupRatioInfo.HasSpecialRatio)
  122. // cleanup
  123. require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1,"svip":1}`))
  124. }
  125. func TestHandleGroupRatio_SpecialGroupGroupRatio(t *testing.T) {
  126. setupUsePriceSwitchTest(t)
  127. require.NoError(t, ratio_setting.UpdateGroupGroupRatioByJSONString(`{"vip":{"premium":0.8}}`))
  128. c := buildTestContext(t)
  129. info := &relaycommon.RelayInfo{
  130. UsingGroup: "premium",
  131. UserGroup: "vip",
  132. }
  133. groupRatioInfo := HandleGroupRatio(c, info)
  134. assert.Equal(t, 0.8, groupRatioInfo.GroupRatio)
  135. assert.Equal(t, 0.8, groupRatioInfo.GroupSpecialRatio)
  136. assert.True(t, groupRatioInfo.HasSpecialRatio)
  137. // cleanup
  138. require.NoError(t, ratio_setting.UpdateGroupGroupRatioByJSONString(`{}`))
  139. }
  140. func TestHandleGroupRatio_AutoGroupOverridesUsingGroup(t *testing.T) {
  141. setupUsePriceSwitchTest(t)
  142. require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1,"svip":1,"vip-channel":0.6}`))
  143. c := buildTestContext(t)
  144. // Simulate: distributor set auto_group = "vip-channel" in context
  145. c.Set("auto_group", "vip-channel")
  146. info := &relaycommon.RelayInfo{
  147. UsingGroup: "default", // will be overridden
  148. UserGroup: "default",
  149. }
  150. groupRatioInfo := HandleGroupRatio(c, info)
  151. // UsingGroup should be updated to "vip-channel"
  152. assert.Equal(t, "vip-channel", info.UsingGroup)
  153. // GroupRatio should reflect vip-channel's ratio
  154. assert.Equal(t, 0.6, groupRatioInfo.GroupRatio)
  155. // cleanup
  156. require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1,"svip":1}`))
  157. }
  158. func TestHandleGroupRatio_AutoGroupWithSpecialRatio(t *testing.T) {
  159. setupUsePriceSwitchTest(t)
  160. require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1,"svip":1,"premium":1.0}`))
  161. require.NoError(t, ratio_setting.UpdateGroupGroupRatioByJSONString(`{"vip":{"premium":0.5}}`))
  162. c := buildTestContext(t)
  163. c.Set("auto_group", "premium")
  164. info := &relaycommon.RelayInfo{
  165. UsingGroup: "default",
  166. UserGroup: "vip",
  167. }
  168. groupRatioInfo := HandleGroupRatio(c, info)
  169. assert.Equal(t, "premium", info.UsingGroup)
  170. // Should use the special ratio (0.5) from vip->premium, not the generic group ratio (1.0)
  171. assert.Equal(t, 0.5, groupRatioInfo.GroupRatio)
  172. assert.True(t, groupRatioInfo.HasSpecialRatio)
  173. // cleanup
  174. require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1,"svip":1}`))
  175. require.NoError(t, ratio_setting.UpdateGroupGroupRatioByJSONString(`{}`))
  176. }
  177. // --- ModelPriceHelper with group ratio ---
  178. func TestModelPriceHelper_GroupRatioAffectsPreConsume(t *testing.T) {
  179. setupUsePriceSwitchTest(t)
  180. require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1,"svip":1,"discount":0.5}`))
  181. c := buildTestContext(t)
  182. info := &relaycommon.RelayInfo{
  183. OriginModelName: testModelUsePriceSwitch,
  184. UsingGroup: "discount",
  185. UserGroup: "default",
  186. UserSetting: dto.UserSetting{},
  187. }
  188. priceData, err := ModelPriceHelper(c, info, 100, &types.TokenCountMeta{})
  189. require.NoError(t, err)
  190. // model price 0.5, group ratio 0.5 => preConsume = 0.5 * QuotaPerUnit * 0.5
  191. assert.Equal(t, 0.5, priceData.GroupRatioInfo.GroupRatio)
  192. assert.True(t, priceData.UsePrice)
  193. // 0.5 (modelPrice) * 500000 (QuotaPerUnit) * 0.5 (groupRatio) = 125000
  194. expectedQuota := int(0.5 * float64(common.QuotaPerUnit) * 0.5)
  195. assert.Equal(t, expectedQuota, priceData.QuotaToPreConsume)
  196. // cleanup
  197. require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1,"svip":1}`))
  198. }
  199. func TestModelPriceHelper_RatioModeWithGroupRatio(t *testing.T) {
  200. setupUsePriceSwitchTest(t)
  201. // Remove model price so it falls back to ratio mode
  202. require.NoError(t, ratio_setting.UpdateModelPriceByJSONString(`{}`))
  203. require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1,"svip":1,"double":2.0}`))
  204. c := buildTestContext(t)
  205. info := &relaycommon.RelayInfo{
  206. OriginModelName: testModelUsePriceSwitch,
  207. UsingGroup: "double",
  208. UserGroup: "default",
  209. UserSetting: dto.UserSetting{},
  210. }
  211. priceData, err := ModelPriceHelper(c, info, 100, &types.TokenCountMeta{})
  212. require.NoError(t, err)
  213. assert.False(t, priceData.UsePrice)
  214. assert.Equal(t, 15.0, priceData.ModelRatio)
  215. assert.Equal(t, 2.0, priceData.GroupRatioInfo.GroupRatio)
  216. // ratio = modelRatio(15) * groupRatio(2.0) = 30
  217. // preConsumedQuota = int(preConsumedTokens * 30)
  218. expectedRatio := 15.0 * 2.0
  219. assert.Equal(t, expectedRatio, priceData.ModelRatio*priceData.GroupRatioInfo.GroupRatio)
  220. // cleanup
  221. require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1,"svip":1}`))
  222. }
  223. func TestModelPriceHelper_FreeGroup(t *testing.T) {
  224. setupUsePriceSwitchTest(t)
  225. require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1,"svip":1,"free":0}`))
  226. // Disable pre-consume for free models so freeModel flag is set
  227. origSetting := operation_setting.GetQuotaSetting()
  228. operation_setting.GetQuotaSetting().EnableFreeModelPreConsume = false
  229. t.Cleanup(func() {
  230. operation_setting.GetQuotaSetting().EnableFreeModelPreConsume = origSetting.EnableFreeModelPreConsume
  231. })
  232. c := buildTestContext(t)
  233. info := &relaycommon.RelayInfo{
  234. OriginModelName: testModelUsePriceSwitch,
  235. UsingGroup: "free",
  236. UserGroup: "default",
  237. UserSetting: dto.UserSetting{},
  238. }
  239. priceData, err := ModelPriceHelper(c, info, 100, &types.TokenCountMeta{})
  240. require.NoError(t, err)
  241. // group ratio = 0 means free model
  242. assert.Equal(t, 0.0, priceData.GroupRatioInfo.GroupRatio)
  243. assert.True(t, priceData.FreeModel)
  244. assert.Equal(t, 0, priceData.QuotaToPreConsume)
  245. // cleanup
  246. require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1,"svip":1}`))
  247. }