quota.go 22 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609
  1. package service
  2. import (
  3. "errors"
  4. "fmt"
  5. "log"
  6. "math"
  7. "strings"
  8. "time"
  9. "github.com/QuantumNous/new-api/common"
  10. "github.com/QuantumNous/new-api/constant"
  11. "github.com/QuantumNous/new-api/dto"
  12. "github.com/QuantumNous/new-api/logger"
  13. "github.com/QuantumNous/new-api/model"
  14. relaycommon "github.com/QuantumNous/new-api/relay/common"
  15. "github.com/QuantumNous/new-api/setting/ratio_setting"
  16. "github.com/QuantumNous/new-api/setting/system_setting"
  17. "github.com/QuantumNous/new-api/types"
  18. "github.com/bytedance/gopkg/util/gopool"
  19. "github.com/gin-gonic/gin"
  20. "github.com/shopspring/decimal"
  21. )
  22. type TokenDetails struct {
  23. TextTokens int
  24. AudioTokens int
  25. }
  26. type QuotaInfo struct {
  27. InputDetails TokenDetails
  28. OutputDetails TokenDetails
  29. ModelName string
  30. UsePrice bool
  31. ModelPrice float64
  32. ModelRatio float64
  33. GroupRatio float64
  34. }
  35. func hasCustomModelRatio(modelName string, currentRatio float64) bool {
  36. defaultRatio, exists := ratio_setting.GetDefaultModelRatioMap()[modelName]
  37. if !exists {
  38. return true
  39. }
  40. return currentRatio != defaultRatio
  41. }
  42. func calculateAudioQuota(info QuotaInfo) int {
  43. if info.UsePrice {
  44. modelPrice := decimal.NewFromFloat(info.ModelPrice)
  45. quotaPerUnit := decimal.NewFromFloat(common.QuotaPerUnit)
  46. groupRatio := decimal.NewFromFloat(info.GroupRatio)
  47. quota := modelPrice.Mul(quotaPerUnit).Mul(groupRatio)
  48. return int(quota.IntPart())
  49. }
  50. completionRatio := decimal.NewFromFloat(ratio_setting.GetCompletionRatio(info.ModelName))
  51. audioRatio := decimal.NewFromFloat(ratio_setting.GetAudioRatio(info.ModelName))
  52. audioCompletionRatio := decimal.NewFromFloat(ratio_setting.GetAudioCompletionRatio(info.ModelName))
  53. groupRatio := decimal.NewFromFloat(info.GroupRatio)
  54. modelRatio := decimal.NewFromFloat(info.ModelRatio)
  55. ratio := groupRatio.Mul(modelRatio)
  56. inputTextTokens := decimal.NewFromInt(int64(info.InputDetails.TextTokens))
  57. outputTextTokens := decimal.NewFromInt(int64(info.OutputDetails.TextTokens))
  58. inputAudioTokens := decimal.NewFromInt(int64(info.InputDetails.AudioTokens))
  59. outputAudioTokens := decimal.NewFromInt(int64(info.OutputDetails.AudioTokens))
  60. quota := decimal.Zero
  61. quota = quota.Add(inputTextTokens)
  62. quota = quota.Add(outputTextTokens.Mul(completionRatio))
  63. quota = quota.Add(inputAudioTokens.Mul(audioRatio))
  64. quota = quota.Add(outputAudioTokens.Mul(audioRatio).Mul(audioCompletionRatio))
  65. quota = quota.Mul(ratio)
  66. // If ratio is not zero and quota is less than or equal to zero, set quota to 1
  67. if !ratio.IsZero() && quota.LessThanOrEqual(decimal.Zero) {
  68. quota = decimal.NewFromInt(1)
  69. }
  70. return int(quota.Round(0).IntPart())
  71. }
  72. func PreWssConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usage *dto.RealtimeUsage) error {
  73. if relayInfo.UsePrice {
  74. return nil
  75. }
  76. userQuota, err := model.GetUserQuota(relayInfo.UserId, false)
  77. if err != nil {
  78. return err
  79. }
  80. token, err := model.GetTokenByKey(strings.TrimPrefix(relayInfo.TokenKey, "sk-"), false)
  81. if err != nil {
  82. return err
  83. }
  84. modelName := relayInfo.OriginModelName
  85. textInputTokens := usage.InputTokenDetails.TextTokens
  86. textOutTokens := usage.OutputTokenDetails.TextTokens
  87. audioInputTokens := usage.InputTokenDetails.AudioTokens
  88. audioOutTokens := usage.OutputTokenDetails.AudioTokens
  89. groupRatio := ratio_setting.GetGroupRatio(relayInfo.UsingGroup)
  90. modelRatio, _, _ := ratio_setting.GetModelRatio(modelName)
  91. autoGroup, exists := common.GetContextKey(ctx, constant.ContextKeyAutoGroup)
  92. if exists {
  93. groupRatio = ratio_setting.GetGroupRatio(autoGroup.(string))
  94. log.Printf("final group ratio: %f", groupRatio)
  95. relayInfo.UsingGroup = autoGroup.(string)
  96. }
  97. actualGroupRatio := groupRatio
  98. userGroupRatio, ok := ratio_setting.GetGroupGroupRatio(relayInfo.UserGroup, relayInfo.UsingGroup)
  99. if ok {
  100. actualGroupRatio = userGroupRatio
  101. }
  102. quotaInfo := QuotaInfo{
  103. InputDetails: TokenDetails{
  104. TextTokens: textInputTokens,
  105. AudioTokens: audioInputTokens,
  106. },
  107. OutputDetails: TokenDetails{
  108. TextTokens: textOutTokens,
  109. AudioTokens: audioOutTokens,
  110. },
  111. ModelName: modelName,
  112. UsePrice: relayInfo.UsePrice,
  113. ModelRatio: modelRatio,
  114. GroupRatio: actualGroupRatio,
  115. }
  116. quota := calculateAudioQuota(quotaInfo)
  117. if userQuota < quota {
  118. return fmt.Errorf("user quota is not enough, user quota: %s, need quota: %s", logger.FormatQuota(userQuota), logger.FormatQuota(quota))
  119. }
  120. if !token.UnlimitedQuota && token.RemainQuota < quota {
  121. return fmt.Errorf("token quota is not enough, token remain quota: %s, need quota: %s", logger.FormatQuota(token.RemainQuota), logger.FormatQuota(quota))
  122. }
  123. err = PostConsumeQuota(relayInfo, quota, 0, false)
  124. if err != nil {
  125. return err
  126. }
  127. logger.LogInfo(ctx, "realtime streaming consume quota success, quota: "+fmt.Sprintf("%d", quota))
  128. return nil
  129. }
  130. func PostWssConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, modelName string,
  131. usage *dto.RealtimeUsage, extraContent string) {
  132. useTimeSeconds := time.Now().Unix() - relayInfo.StartTime.Unix()
  133. textInputTokens := usage.InputTokenDetails.TextTokens
  134. textOutTokens := usage.OutputTokenDetails.TextTokens
  135. audioInputTokens := usage.InputTokenDetails.AudioTokens
  136. audioOutTokens := usage.OutputTokenDetails.AudioTokens
  137. tokenName := ctx.GetString("token_name")
  138. completionRatio := decimal.NewFromFloat(ratio_setting.GetCompletionRatio(modelName))
  139. audioRatio := decimal.NewFromFloat(ratio_setting.GetAudioRatio(relayInfo.OriginModelName))
  140. audioCompletionRatio := decimal.NewFromFloat(ratio_setting.GetAudioCompletionRatio(modelName))
  141. modelRatio := relayInfo.PriceData.ModelRatio
  142. groupRatio := relayInfo.PriceData.GroupRatioInfo.GroupRatio
  143. modelPrice := relayInfo.PriceData.ModelPrice
  144. usePrice := relayInfo.PriceData.UsePrice
  145. quotaInfo := QuotaInfo{
  146. InputDetails: TokenDetails{
  147. TextTokens: textInputTokens,
  148. AudioTokens: audioInputTokens,
  149. },
  150. OutputDetails: TokenDetails{
  151. TextTokens: textOutTokens,
  152. AudioTokens: audioOutTokens,
  153. },
  154. ModelName: modelName,
  155. UsePrice: usePrice,
  156. ModelRatio: modelRatio,
  157. GroupRatio: groupRatio,
  158. }
  159. quota := calculateAudioQuota(quotaInfo)
  160. totalTokens := usage.TotalTokens
  161. var logContent string
  162. if !usePrice {
  163. logContent = fmt.Sprintf("模型倍率 %.2f,补全倍率 %.2f,音频倍率 %.2f,音频补全倍率 %.2f,分组倍率 %.2f",
  164. modelRatio, completionRatio.InexactFloat64(), audioRatio.InexactFloat64(), audioCompletionRatio.InexactFloat64(), groupRatio)
  165. } else {
  166. logContent = fmt.Sprintf("模型价格 %.2f,分组倍率 %.2f", modelPrice, groupRatio)
  167. }
  168. // record all the consume log even if quota is 0
  169. if totalTokens == 0 {
  170. // in this case, must be some error happened
  171. // we cannot just return, because we may have to return the pre-consumed quota
  172. quota = 0
  173. logContent += fmt.Sprintf("(可能是上游超时)")
  174. logger.LogError(ctx, fmt.Sprintf("total tokens is 0, cannot consume quota, userId %d, channelId %d, "+
  175. "tokenId %d, model %s, pre-consumed quota %d", relayInfo.UserId, relayInfo.ChannelId, relayInfo.TokenId, modelName, relayInfo.FinalPreConsumedQuota))
  176. } else {
  177. model.UpdateUserUsedQuotaAndRequestCount(relayInfo.UserId, quota)
  178. model.UpdateChannelUsedQuota(relayInfo.ChannelId, quota)
  179. }
  180. logModel := modelName
  181. if extraContent != "" {
  182. logContent += ", " + extraContent
  183. }
  184. other := GenerateWssOtherInfo(ctx, relayInfo, usage, modelRatio, groupRatio,
  185. completionRatio.InexactFloat64(), audioRatio.InexactFloat64(), audioCompletionRatio.InexactFloat64(), modelPrice, relayInfo.PriceData.GroupRatioInfo.GroupSpecialRatio)
  186. model.RecordConsumeLog(ctx, relayInfo.UserId, model.RecordConsumeLogParams{
  187. ChannelId: relayInfo.ChannelId,
  188. PromptTokens: usage.InputTokens,
  189. CompletionTokens: usage.OutputTokens,
  190. ModelName: logModel,
  191. TokenName: tokenName,
  192. Quota: quota,
  193. Content: logContent,
  194. TokenId: relayInfo.TokenId,
  195. UseTimeSeconds: int(useTimeSeconds),
  196. IsStream: relayInfo.IsStream,
  197. Group: relayInfo.UsingGroup,
  198. Other: other,
  199. })
  200. }
  201. func PostClaudeConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usage *dto.Usage) {
  202. if usage != nil {
  203. ObserveChannelAffinityUsageCacheByRelayFormat(ctx, usage, relayInfo.GetFinalRequestRelayFormat())
  204. }
  205. useTimeSeconds := time.Now().Unix() - relayInfo.StartTime.Unix()
  206. promptTokens := usage.PromptTokens
  207. completionTokens := usage.CompletionTokens
  208. modelName := relayInfo.OriginModelName
  209. tokenName := ctx.GetString("token_name")
  210. completionRatio := relayInfo.PriceData.CompletionRatio
  211. modelRatio := relayInfo.PriceData.ModelRatio
  212. groupRatio := relayInfo.PriceData.GroupRatioInfo.GroupRatio
  213. modelPrice := relayInfo.PriceData.ModelPrice
  214. cacheRatio := relayInfo.PriceData.CacheRatio
  215. cacheTokens := usage.PromptTokensDetails.CachedTokens
  216. cacheCreationRatio := relayInfo.PriceData.CacheCreationRatio
  217. cacheCreationRatio5m := relayInfo.PriceData.CacheCreation5mRatio
  218. cacheCreationRatio1h := relayInfo.PriceData.CacheCreation1hRatio
  219. cacheCreationTokens := usage.PromptTokensDetails.CachedCreationTokens
  220. cacheCreationTokens5m := usage.ClaudeCacheCreation5mTokens
  221. cacheCreationTokens1h := usage.ClaudeCacheCreation1hTokens
  222. if relayInfo.ChannelType == constant.ChannelTypeOpenRouter {
  223. promptTokens -= cacheTokens
  224. isUsingCustomSettings := relayInfo.PriceData.UsePrice || hasCustomModelRatio(modelName, relayInfo.PriceData.ModelRatio)
  225. if cacheCreationTokens == 0 && relayInfo.PriceData.CacheCreationRatio != 1 && usage.Cost != 0 && !isUsingCustomSettings {
  226. maybeCacheCreationTokens := CalcOpenRouterCacheCreateTokens(*usage, relayInfo.PriceData)
  227. if maybeCacheCreationTokens >= 0 && promptTokens >= maybeCacheCreationTokens {
  228. cacheCreationTokens = maybeCacheCreationTokens
  229. }
  230. }
  231. promptTokens -= cacheCreationTokens
  232. }
  233. calculateQuota := 0.0
  234. if !relayInfo.PriceData.UsePrice {
  235. calculateQuota = float64(promptTokens)
  236. calculateQuota += float64(cacheTokens) * cacheRatio
  237. calculateQuota += float64(cacheCreationTokens5m) * cacheCreationRatio5m
  238. calculateQuota += float64(cacheCreationTokens1h) * cacheCreationRatio1h
  239. remainingCacheCreationTokens := cacheCreationTokens - cacheCreationTokens5m - cacheCreationTokens1h
  240. if remainingCacheCreationTokens > 0 {
  241. calculateQuota += float64(remainingCacheCreationTokens) * cacheCreationRatio
  242. }
  243. calculateQuota += float64(completionTokens) * completionRatio
  244. calculateQuota = calculateQuota * groupRatio * modelRatio
  245. } else {
  246. calculateQuota = modelPrice * common.QuotaPerUnit * groupRatio
  247. }
  248. if modelRatio != 0 && calculateQuota <= 0 {
  249. calculateQuota = 1
  250. }
  251. quota := int(calculateQuota)
  252. totalTokens := promptTokens + completionTokens
  253. var logContent string
  254. // record all the consume log even if quota is 0
  255. if totalTokens == 0 {
  256. // in this case, must be some error happened
  257. // we cannot just return, because we may have to return the pre-consumed quota
  258. quota = 0
  259. logContent += fmt.Sprintf("(可能是上游出错)")
  260. logger.LogError(ctx, fmt.Sprintf("total tokens is 0, cannot consume quota, userId %d, channelId %d, "+
  261. "tokenId %d, model %s, pre-consumed quota %d", relayInfo.UserId, relayInfo.ChannelId, relayInfo.TokenId, modelName, relayInfo.FinalPreConsumedQuota))
  262. } else {
  263. model.UpdateUserUsedQuotaAndRequestCount(relayInfo.UserId, quota)
  264. model.UpdateChannelUsedQuota(relayInfo.ChannelId, quota)
  265. }
  266. if err := SettleBilling(ctx, relayInfo, quota); err != nil {
  267. logger.LogError(ctx, "error settling billing: "+err.Error())
  268. }
  269. other := GenerateClaudeOtherInfo(ctx, relayInfo, modelRatio, groupRatio, completionRatio,
  270. cacheTokens, cacheRatio,
  271. cacheCreationTokens, cacheCreationRatio,
  272. cacheCreationTokens5m, cacheCreationRatio5m,
  273. cacheCreationTokens1h, cacheCreationRatio1h,
  274. modelPrice, relayInfo.PriceData.GroupRatioInfo.GroupSpecialRatio)
  275. model.RecordConsumeLog(ctx, relayInfo.UserId, model.RecordConsumeLogParams{
  276. ChannelId: relayInfo.ChannelId,
  277. PromptTokens: promptTokens,
  278. CompletionTokens: completionTokens,
  279. ModelName: modelName,
  280. TokenName: tokenName,
  281. Quota: quota,
  282. Content: logContent,
  283. TokenId: relayInfo.TokenId,
  284. UseTimeSeconds: int(useTimeSeconds),
  285. IsStream: relayInfo.IsStream,
  286. Group: relayInfo.UsingGroup,
  287. Other: other,
  288. })
  289. }
  290. func CalcOpenRouterCacheCreateTokens(usage dto.Usage, priceData types.PriceData) int {
  291. if priceData.CacheCreationRatio == 1 {
  292. return 0
  293. }
  294. quotaPrice := priceData.ModelRatio / common.QuotaPerUnit
  295. promptCacheCreatePrice := quotaPrice * priceData.CacheCreationRatio
  296. promptCacheReadPrice := quotaPrice * priceData.CacheRatio
  297. completionPrice := quotaPrice * priceData.CompletionRatio
  298. cost, _ := usage.Cost.(float64)
  299. totalPromptTokens := float64(usage.PromptTokens)
  300. completionTokens := float64(usage.CompletionTokens)
  301. promptCacheReadTokens := float64(usage.PromptTokensDetails.CachedTokens)
  302. return int(math.Round((cost -
  303. totalPromptTokens*quotaPrice +
  304. promptCacheReadTokens*(quotaPrice-promptCacheReadPrice) -
  305. completionTokens*completionPrice) /
  306. (promptCacheCreatePrice - quotaPrice)))
  307. }
  308. func PostAudioConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usage *dto.Usage, extraContent string) {
  309. useTimeSeconds := time.Now().Unix() - relayInfo.StartTime.Unix()
  310. textInputTokens := usage.PromptTokensDetails.TextTokens
  311. textOutTokens := usage.CompletionTokenDetails.TextTokens
  312. audioInputTokens := usage.PromptTokensDetails.AudioTokens
  313. audioOutTokens := usage.CompletionTokenDetails.AudioTokens
  314. tokenName := ctx.GetString("token_name")
  315. completionRatio := decimal.NewFromFloat(ratio_setting.GetCompletionRatio(relayInfo.OriginModelName))
  316. audioRatio := decimal.NewFromFloat(ratio_setting.GetAudioRatio(relayInfo.OriginModelName))
  317. audioCompletionRatio := decimal.NewFromFloat(ratio_setting.GetAudioCompletionRatio(relayInfo.OriginModelName))
  318. modelRatio := relayInfo.PriceData.ModelRatio
  319. groupRatio := relayInfo.PriceData.GroupRatioInfo.GroupRatio
  320. modelPrice := relayInfo.PriceData.ModelPrice
  321. usePrice := relayInfo.PriceData.UsePrice
  322. quotaInfo := QuotaInfo{
  323. InputDetails: TokenDetails{
  324. TextTokens: textInputTokens,
  325. AudioTokens: audioInputTokens,
  326. },
  327. OutputDetails: TokenDetails{
  328. TextTokens: textOutTokens,
  329. AudioTokens: audioOutTokens,
  330. },
  331. ModelName: relayInfo.OriginModelName,
  332. UsePrice: usePrice,
  333. ModelRatio: modelRatio,
  334. GroupRatio: groupRatio,
  335. }
  336. quota := calculateAudioQuota(quotaInfo)
  337. totalTokens := usage.TotalTokens
  338. var logContent string
  339. if !usePrice {
  340. logContent = fmt.Sprintf("模型倍率 %.2f,补全倍率 %.2f,音频倍率 %.2f,音频补全倍率 %.2f,分组倍率 %.2f",
  341. modelRatio, completionRatio.InexactFloat64(), audioRatio.InexactFloat64(), audioCompletionRatio.InexactFloat64(), groupRatio)
  342. } else {
  343. logContent = fmt.Sprintf("模型价格 %.2f,分组倍率 %.2f", modelPrice, groupRatio)
  344. }
  345. // record all the consume log even if quota is 0
  346. if totalTokens == 0 {
  347. // in this case, must be some error happened
  348. // we cannot just return, because we may have to return the pre-consumed quota
  349. quota = 0
  350. logContent += fmt.Sprintf("(可能是上游超时)")
  351. logger.LogError(ctx, fmt.Sprintf("total tokens is 0, cannot consume quota, userId %d, channelId %d, "+
  352. "tokenId %d, model %s, pre-consumed quota %d", relayInfo.UserId, relayInfo.ChannelId, relayInfo.TokenId, relayInfo.OriginModelName, relayInfo.FinalPreConsumedQuota))
  353. } else {
  354. model.UpdateUserUsedQuotaAndRequestCount(relayInfo.UserId, quota)
  355. model.UpdateChannelUsedQuota(relayInfo.ChannelId, quota)
  356. }
  357. if err := SettleBilling(ctx, relayInfo, quota); err != nil {
  358. logger.LogError(ctx, "error settling billing: "+err.Error())
  359. }
  360. logModel := relayInfo.OriginModelName
  361. if extraContent != "" {
  362. logContent += ", " + extraContent
  363. }
  364. other := GenerateAudioOtherInfo(ctx, relayInfo, usage, modelRatio, groupRatio,
  365. completionRatio.InexactFloat64(), audioRatio.InexactFloat64(), audioCompletionRatio.InexactFloat64(), modelPrice, relayInfo.PriceData.GroupRatioInfo.GroupSpecialRatio)
  366. model.RecordConsumeLog(ctx, relayInfo.UserId, model.RecordConsumeLogParams{
  367. ChannelId: relayInfo.ChannelId,
  368. PromptTokens: usage.PromptTokens,
  369. CompletionTokens: usage.CompletionTokens,
  370. ModelName: logModel,
  371. TokenName: tokenName,
  372. Quota: quota,
  373. Content: logContent,
  374. TokenId: relayInfo.TokenId,
  375. UseTimeSeconds: int(useTimeSeconds),
  376. IsStream: relayInfo.IsStream,
  377. Group: relayInfo.UsingGroup,
  378. Other: other,
  379. })
  380. }
  381. func PreConsumeTokenQuota(relayInfo *relaycommon.RelayInfo, quota int) error {
  382. if quota < 0 {
  383. return errors.New("quota 不能为负数!")
  384. }
  385. if relayInfo.IsPlayground {
  386. return nil
  387. }
  388. //if relayInfo.TokenUnlimited {
  389. // return nil
  390. //}
  391. token, err := model.GetTokenByKey(relayInfo.TokenKey, false)
  392. if err != nil {
  393. return err
  394. }
  395. if !relayInfo.TokenUnlimited && token.RemainQuota < quota {
  396. return fmt.Errorf("token quota is not enough, token remain quota: %s, need quota: %s", logger.FormatQuota(token.RemainQuota), logger.FormatQuota(quota))
  397. }
  398. err = model.DecreaseTokenQuota(relayInfo.TokenId, relayInfo.TokenKey, quota)
  399. if err != nil {
  400. return err
  401. }
  402. return nil
  403. }
  404. func PostConsumeQuota(relayInfo *relaycommon.RelayInfo, quota int, preConsumedQuota int, sendEmail bool) (err error) {
  405. // 1) Consume from wallet quota OR subscription item
  406. if relayInfo != nil && relayInfo.BillingSource == BillingSourceSubscription {
  407. if relayInfo.SubscriptionId == 0 {
  408. return errors.New("subscription id is missing")
  409. }
  410. delta := int64(quota)
  411. if delta != 0 {
  412. if err := model.PostConsumeUserSubscriptionDelta(relayInfo.SubscriptionId, delta); err != nil {
  413. return err
  414. }
  415. relayInfo.SubscriptionPostDelta += delta
  416. }
  417. } else {
  418. // Wallet
  419. if quota > 0 {
  420. err = model.DecreaseUserQuota(relayInfo.UserId, quota)
  421. } else {
  422. err = model.IncreaseUserQuota(relayInfo.UserId, -quota, false)
  423. }
  424. if err != nil {
  425. return err
  426. }
  427. }
  428. if !relayInfo.IsPlayground {
  429. if quota > 0 {
  430. err = model.DecreaseTokenQuota(relayInfo.TokenId, relayInfo.TokenKey, quota)
  431. } else {
  432. err = model.IncreaseTokenQuota(relayInfo.TokenId, relayInfo.TokenKey, -quota)
  433. }
  434. if err != nil {
  435. return err
  436. }
  437. }
  438. if sendEmail {
  439. if (quota + preConsumedQuota) != 0 {
  440. checkAndSendQuotaNotify(relayInfo, quota, preConsumedQuota)
  441. }
  442. }
  443. return nil
  444. }
  445. func checkAndSendQuotaNotify(relayInfo *relaycommon.RelayInfo, quota int, preConsumedQuota int) {
  446. gopool.Go(func() {
  447. userSetting := relayInfo.UserSetting
  448. threshold := common.QuotaRemindThreshold
  449. if userSetting.QuotaWarningThreshold != 0 {
  450. threshold = int(userSetting.QuotaWarningThreshold)
  451. }
  452. //noMoreQuota := userCache.Quota-(quota+preConsumedQuota) <= 0
  453. quotaTooLow := false
  454. consumeQuota := quota + preConsumedQuota
  455. if relayInfo.UserQuota-consumeQuota < threshold {
  456. quotaTooLow = true
  457. }
  458. if quotaTooLow {
  459. prompt := "您的额度即将用尽"
  460. topUpLink := fmt.Sprintf("%s/console/topup", system_setting.ServerAddress)
  461. // 根据通知方式生成不同的内容格式
  462. var content string
  463. var values []interface{}
  464. notifyType := userSetting.NotifyType
  465. if notifyType == "" {
  466. notifyType = dto.NotifyTypeEmail
  467. }
  468. if notifyType == dto.NotifyTypeBark {
  469. // Bark推送使用简短文本,不支持HTML
  470. content = "{{value}},剩余额度:{{value}},请及时充值"
  471. values = []interface{}{prompt, logger.FormatQuota(relayInfo.UserQuota)}
  472. } else if notifyType == dto.NotifyTypeGotify {
  473. content = "{{value}},当前剩余额度为 {{value}},请及时充值。"
  474. values = []interface{}{prompt, logger.FormatQuota(relayInfo.UserQuota)}
  475. } else {
  476. // 默认内容格式,适用于Email和Webhook(支持HTML)
  477. content = "{{value}},当前剩余额度为 {{value}},为了不影响您的使用,请及时充值。<br/>充值链接:<a href='{{value}}'>{{value}}</a>"
  478. values = []interface{}{prompt, logger.FormatQuota(relayInfo.UserQuota), topUpLink, topUpLink}
  479. }
  480. err := NotifyUser(relayInfo.UserId, relayInfo.UserEmail, relayInfo.UserSetting, dto.NewNotify(dto.NotifyTypeQuotaExceed, prompt, content, values))
  481. if err != nil {
  482. common.SysError(fmt.Sprintf("failed to send quota notify to user %d: %s", relayInfo.UserId, err.Error()))
  483. }
  484. }
  485. })
  486. }
  487. func checkAndSendSubscriptionQuotaNotify(relayInfo *relaycommon.RelayInfo) {
  488. gopool.Go(func() {
  489. if relayInfo == nil {
  490. return
  491. }
  492. if relayInfo.SubscriptionId == 0 || relayInfo.SubscriptionAmountTotal <= 0 {
  493. return
  494. }
  495. userSetting := relayInfo.UserSetting
  496. threshold := common.QuotaRemindThreshold
  497. if userSetting.QuotaWarningThreshold != 0 {
  498. threshold = int(userSetting.QuotaWarningThreshold)
  499. }
  500. usedAfter := relayInfo.SubscriptionAmountUsedAfterPreConsume + relayInfo.SubscriptionPostDelta
  501. remaining := relayInfo.SubscriptionAmountTotal - usedAfter
  502. if remaining >= int64(threshold) {
  503. return
  504. }
  505. prompt := "您的订阅额度即将用尽"
  506. topUpLink := fmt.Sprintf("%s/console/topup", system_setting.ServerAddress)
  507. var content string
  508. var values []interface{}
  509. notifyType := userSetting.NotifyType
  510. if notifyType == "" {
  511. notifyType = dto.NotifyTypeEmail
  512. }
  513. if notifyType == dto.NotifyTypeBark {
  514. content = "{{value}},剩余额度:{{value}},请及时充值"
  515. values = []interface{}{prompt, logger.FormatQuota(int(remaining))}
  516. } else if notifyType == dto.NotifyTypeGotify {
  517. content = "{{value}},当前剩余额度为 {{value}},请及时充值。"
  518. values = []interface{}{prompt, logger.FormatQuota(int(remaining))}
  519. } else {
  520. content = "{{value}},当前剩余额度为 {{value}},为了不影响您的使用,请及时充值。<br/>充值链接:<a href='{{value}}'>{{value}}</a>"
  521. values = []interface{}{prompt, logger.FormatQuota(int(remaining)), topUpLink, topUpLink}
  522. }
  523. if err := NotifyUser(relayInfo.UserId, relayInfo.UserEmail, relayInfo.UserSetting, dto.NewNotify(dto.NotifyTypeQuotaExceed, prompt, content, values)); err != nil {
  524. common.SysError(fmt.Sprintf("failed to send subscription quota notify to user %d: %s", relayInfo.UserId, err.Error()))
  525. }
  526. })
  527. }