Вы не можете выбрать более 25 тем Темы должны начинаться с буквы или цифры, могут содержать дефисы(-) и должны содержать не более 35 символов.
 
 
 

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