您最多选择25个主题 主题必须以字母或数字开头,可以包含连字符 (-),并且长度不得超过35个字符
 
 
 

615 行
22 KiB

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