quota.go 25 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688
  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. "github.com/QuantumNous/new-api/pkg/billingexpr"
  15. relaycommon "github.com/QuantumNous/new-api/relay/common"
  16. "github.com/QuantumNous/new-api/setting/ratio_setting"
  17. "github.com/QuantumNous/new-api/setting/system_setting"
  18. "github.com/QuantumNous/new-api/types"
  19. "github.com/bytedance/gopkg/util/gopool"
  20. "github.com/gin-gonic/gin"
  21. "github.com/shopspring/decimal"
  22. )
  23. type TokenDetails struct {
  24. TextTokens int
  25. AudioTokens int
  26. }
  27. type QuotaInfo struct {
  28. InputDetails TokenDetails
  29. OutputDetails TokenDetails
  30. ModelName string
  31. UsePrice bool
  32. ModelPrice float64
  33. ModelRatio float64
  34. GroupRatio 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)
  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)
  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. }
  117. quota := calculateAudioQuota(quotaInfo)
  118. if userQuota < quota {
  119. return fmt.Errorf("user quota is not enough, user quota: %s, need quota: %s", logger.FormatQuota(userQuota), logger.FormatQuota(quota))
  120. }
  121. if !token.UnlimitedQuota && token.RemainQuota < quota {
  122. return fmt.Errorf("token quota is not enough, token remain quota: %s, need quota: %s", logger.FormatQuota(token.RemainQuota), logger.FormatQuota(quota))
  123. }
  124. err = PostConsumeQuota(relayInfo, quota, 0, false)
  125. if err != nil {
  126. return err
  127. }
  128. logger.LogInfo(ctx, "realtime streaming consume quota success, quota: "+fmt.Sprintf("%d", quota))
  129. return nil
  130. }
  131. func PostWssConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, modelName string,
  132. usage *dto.RealtimeUsage, extraContent string) {
  133. // Tiered billing early return
  134. if ok, tieredQuota, tieredResult := TryTieredSettle(relayInfo, billingexpr.TokenParams{
  135. P: float64(usage.InputTokens),
  136. C: float64(usage.OutputTokens),
  137. }); ok {
  138. postConsumeQuotaTieredService(ctx, relayInfo, modelName, usage.InputTokens, usage.OutputTokens, usage.TotalTokens, tieredQuota, tieredResult, extraContent)
  139. return
  140. }
  141. useTimeSeconds := time.Now().Unix() - relayInfo.StartTime.Unix()
  142. textInputTokens := usage.InputTokenDetails.TextTokens
  143. textOutTokens := usage.OutputTokenDetails.TextTokens
  144. audioInputTokens := usage.InputTokenDetails.AudioTokens
  145. audioOutTokens := usage.OutputTokenDetails.AudioTokens
  146. tokenName := ctx.GetString("token_name")
  147. completionRatio := decimal.NewFromFloat(ratio_setting.GetCompletionRatio(modelName))
  148. audioRatio := decimal.NewFromFloat(ratio_setting.GetAudioRatio(relayInfo.OriginModelName))
  149. audioCompletionRatio := decimal.NewFromFloat(ratio_setting.GetAudioCompletionRatio(modelName))
  150. modelRatio := relayInfo.PriceData.ModelRatio
  151. groupRatio := relayInfo.PriceData.GroupRatioInfo.GroupRatio
  152. modelPrice := relayInfo.PriceData.ModelPrice
  153. usePrice := relayInfo.PriceData.UsePrice
  154. quotaInfo := QuotaInfo{
  155. InputDetails: TokenDetails{
  156. TextTokens: textInputTokens,
  157. AudioTokens: audioInputTokens,
  158. },
  159. OutputDetails: TokenDetails{
  160. TextTokens: textOutTokens,
  161. AudioTokens: audioOutTokens,
  162. },
  163. ModelName: modelName,
  164. UsePrice: usePrice,
  165. ModelRatio: modelRatio,
  166. GroupRatio: groupRatio,
  167. }
  168. quota := calculateAudioQuota(quotaInfo)
  169. totalTokens := usage.TotalTokens
  170. var logContent string
  171. if !usePrice {
  172. logContent = fmt.Sprintf("模型倍率 %.2f,补全倍率 %.2f,音频倍率 %.2f,音频补全倍率 %.2f,分组倍率 %.2f",
  173. modelRatio, completionRatio.InexactFloat64(), audioRatio.InexactFloat64(), audioCompletionRatio.InexactFloat64(), groupRatio)
  174. } else {
  175. logContent = fmt.Sprintf("模型价格 %.2f,分组倍率 %.2f", modelPrice, groupRatio)
  176. }
  177. // record all the consume log even if quota is 0
  178. if totalTokens == 0 {
  179. // in this case, must be some error happened
  180. // we cannot just return, because we may have to return the pre-consumed quota
  181. quota = 0
  182. logContent += fmt.Sprintf("(可能是上游超时)")
  183. logger.LogError(ctx, fmt.Sprintf("total tokens is 0, cannot consume quota, userId %d, channelId %d, "+
  184. "tokenId %d, model %s, pre-consumed quota %d", relayInfo.UserId, relayInfo.ChannelId, relayInfo.TokenId, modelName, relayInfo.FinalPreConsumedQuota))
  185. } else {
  186. model.UpdateUserUsedQuotaAndRequestCount(relayInfo.UserId, quota)
  187. model.UpdateChannelUsedQuota(relayInfo.ChannelId, quota)
  188. }
  189. logModel := modelName
  190. if extraContent != "" {
  191. logContent += ", " + extraContent
  192. }
  193. other := GenerateWssOtherInfo(ctx, relayInfo, usage, modelRatio, groupRatio,
  194. completionRatio.InexactFloat64(), audioRatio.InexactFloat64(), audioCompletionRatio.InexactFloat64(), modelPrice, relayInfo.PriceData.GroupRatioInfo.GroupSpecialRatio)
  195. model.RecordConsumeLog(ctx, relayInfo.UserId, model.RecordConsumeLogParams{
  196. ChannelId: relayInfo.ChannelId,
  197. PromptTokens: usage.InputTokens,
  198. CompletionTokens: usage.OutputTokens,
  199. ModelName: logModel,
  200. TokenName: tokenName,
  201. Quota: quota,
  202. Content: logContent,
  203. TokenId: relayInfo.TokenId,
  204. UseTimeSeconds: int(useTimeSeconds),
  205. IsStream: relayInfo.IsStream,
  206. Group: relayInfo.UsingGroup,
  207. Other: other,
  208. })
  209. }
  210. func PostClaudeConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usage *dto.Usage) {
  211. if usage != nil {
  212. ObserveChannelAffinityUsageCacheByRelayFormat(ctx, usage, relayInfo.GetFinalRequestRelayFormat())
  213. }
  214. // Tiered billing early return
  215. if ok, tieredQuota, tieredResult := TryTieredSettle(relayInfo, billingexpr.TokenParams{
  216. P: float64(usage.PromptTokens),
  217. C: float64(usage.CompletionTokens),
  218. CR: float64(usage.PromptTokensDetails.CachedTokens),
  219. CC: float64(usage.PromptTokensDetails.CachedCreationTokens - usage.ClaudeCacheCreation1hTokens),
  220. CC1h: float64(usage.ClaudeCacheCreation1hTokens),
  221. }); ok {
  222. postConsumeQuotaTieredService(ctx, relayInfo, relayInfo.OriginModelName, usage.PromptTokens, usage.CompletionTokens, usage.PromptTokens+usage.CompletionTokens, tieredQuota, tieredResult, "")
  223. return
  224. }
  225. useTimeSeconds := time.Now().Unix() - relayInfo.StartTime.Unix()
  226. promptTokens := usage.PromptTokens
  227. completionTokens := usage.CompletionTokens
  228. modelName := relayInfo.OriginModelName
  229. tokenName := ctx.GetString("token_name")
  230. completionRatio := relayInfo.PriceData.CompletionRatio
  231. modelRatio := relayInfo.PriceData.ModelRatio
  232. groupRatio := relayInfo.PriceData.GroupRatioInfo.GroupRatio
  233. modelPrice := relayInfo.PriceData.ModelPrice
  234. cacheRatio := relayInfo.PriceData.CacheRatio
  235. cacheTokens := usage.PromptTokensDetails.CachedTokens
  236. cacheCreationRatio := relayInfo.PriceData.CacheCreationRatio
  237. cacheCreationRatio5m := relayInfo.PriceData.CacheCreation5mRatio
  238. cacheCreationRatio1h := relayInfo.PriceData.CacheCreation1hRatio
  239. cacheCreationTokens := usage.PromptTokensDetails.CachedCreationTokens
  240. cacheCreationTokens5m := usage.ClaudeCacheCreation5mTokens
  241. cacheCreationTokens1h := usage.ClaudeCacheCreation1hTokens
  242. if relayInfo.ChannelType == constant.ChannelTypeOpenRouter {
  243. promptTokens -= cacheTokens
  244. isUsingCustomSettings := relayInfo.PriceData.UsePrice || hasCustomModelRatio(modelName, relayInfo.PriceData.ModelRatio)
  245. if cacheCreationTokens == 0 && relayInfo.PriceData.CacheCreationRatio != 1 && usage.Cost != 0 && !isUsingCustomSettings {
  246. maybeCacheCreationTokens := CalcOpenRouterCacheCreateTokens(*usage, relayInfo.PriceData)
  247. if maybeCacheCreationTokens >= 0 && promptTokens >= maybeCacheCreationTokens {
  248. cacheCreationTokens = maybeCacheCreationTokens
  249. }
  250. }
  251. promptTokens -= cacheCreationTokens
  252. }
  253. calculateQuota := 0.0
  254. if !relayInfo.PriceData.UsePrice {
  255. calculateQuota = float64(promptTokens)
  256. calculateQuota += float64(cacheTokens) * cacheRatio
  257. calculateQuota += float64(cacheCreationTokens5m) * cacheCreationRatio5m
  258. calculateQuota += float64(cacheCreationTokens1h) * cacheCreationRatio1h
  259. remainingCacheCreationTokens := cacheCreationTokens - cacheCreationTokens5m - cacheCreationTokens1h
  260. if remainingCacheCreationTokens > 0 {
  261. calculateQuota += float64(remainingCacheCreationTokens) * cacheCreationRatio
  262. }
  263. calculateQuota += float64(completionTokens) * completionRatio
  264. calculateQuota = calculateQuota * groupRatio * modelRatio
  265. } else {
  266. calculateQuota = modelPrice * common.QuotaPerUnit * groupRatio
  267. }
  268. if modelRatio != 0 && calculateQuota <= 0 {
  269. calculateQuota = 1
  270. }
  271. quota := int(calculateQuota)
  272. totalTokens := promptTokens + completionTokens
  273. var logContent string
  274. // record all the consume log even if quota is 0
  275. if totalTokens == 0 {
  276. // in this case, must be some error happened
  277. // we cannot just return, because we may have to return the pre-consumed quota
  278. quota = 0
  279. logContent += fmt.Sprintf("(可能是上游出错)")
  280. logger.LogError(ctx, fmt.Sprintf("total tokens is 0, cannot consume quota, userId %d, channelId %d, "+
  281. "tokenId %d, model %s, pre-consumed quota %d", relayInfo.UserId, relayInfo.ChannelId, relayInfo.TokenId, modelName, relayInfo.FinalPreConsumedQuota))
  282. } else {
  283. model.UpdateUserUsedQuotaAndRequestCount(relayInfo.UserId, quota)
  284. model.UpdateChannelUsedQuota(relayInfo.ChannelId, quota)
  285. }
  286. if err := SettleBilling(ctx, relayInfo, quota); err != nil {
  287. logger.LogError(ctx, "error settling billing: "+err.Error())
  288. }
  289. other := GenerateClaudeOtherInfo(ctx, relayInfo, modelRatio, groupRatio, completionRatio,
  290. cacheTokens, cacheRatio,
  291. cacheCreationTokens, cacheCreationRatio,
  292. cacheCreationTokens5m, cacheCreationRatio5m,
  293. cacheCreationTokens1h, cacheCreationRatio1h,
  294. modelPrice, relayInfo.PriceData.GroupRatioInfo.GroupSpecialRatio)
  295. model.RecordConsumeLog(ctx, relayInfo.UserId, model.RecordConsumeLogParams{
  296. ChannelId: relayInfo.ChannelId,
  297. PromptTokens: promptTokens,
  298. CompletionTokens: completionTokens,
  299. ModelName: modelName,
  300. TokenName: tokenName,
  301. Quota: quota,
  302. Content: logContent,
  303. TokenId: relayInfo.TokenId,
  304. UseTimeSeconds: int(useTimeSeconds),
  305. IsStream: relayInfo.IsStream,
  306. Group: relayInfo.UsingGroup,
  307. Other: other,
  308. })
  309. }
  310. func CalcOpenRouterCacheCreateTokens(usage dto.Usage, priceData types.PriceData) int {
  311. if priceData.CacheCreationRatio == 1 {
  312. return 0
  313. }
  314. quotaPrice := priceData.ModelRatio / common.QuotaPerUnit
  315. promptCacheCreatePrice := quotaPrice * priceData.CacheCreationRatio
  316. promptCacheReadPrice := quotaPrice * priceData.CacheRatio
  317. completionPrice := quotaPrice * priceData.CompletionRatio
  318. cost, _ := usage.Cost.(float64)
  319. totalPromptTokens := float64(usage.PromptTokens)
  320. completionTokens := float64(usage.CompletionTokens)
  321. promptCacheReadTokens := float64(usage.PromptTokensDetails.CachedTokens)
  322. return int(math.Round((cost -
  323. totalPromptTokens*quotaPrice +
  324. promptCacheReadTokens*(quotaPrice-promptCacheReadPrice) -
  325. completionTokens*completionPrice) /
  326. (promptCacheCreatePrice - quotaPrice)))
  327. }
  328. func PostAudioConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usage *dto.Usage, extraContent string) {
  329. // Tiered billing early return
  330. if ok, tieredQuota, tieredResult := TryTieredSettle(relayInfo, billingexpr.TokenParams{
  331. P: float64(usage.PromptTokens),
  332. C: float64(usage.CompletionTokens),
  333. CR: float64(usage.PromptTokensDetails.CachedTokens),
  334. }); ok {
  335. postConsumeQuotaTieredService(ctx, relayInfo, relayInfo.OriginModelName, usage.PromptTokens, usage.CompletionTokens, usage.TotalTokens, tieredQuota, tieredResult, extraContent)
  336. return
  337. }
  338. useTimeSeconds := time.Now().Unix() - relayInfo.StartTime.Unix()
  339. textInputTokens := usage.PromptTokensDetails.TextTokens
  340. textOutTokens := usage.CompletionTokenDetails.TextTokens
  341. audioInputTokens := usage.PromptTokensDetails.AudioTokens
  342. audioOutTokens := usage.CompletionTokenDetails.AudioTokens
  343. tokenName := ctx.GetString("token_name")
  344. completionRatio := decimal.NewFromFloat(ratio_setting.GetCompletionRatio(relayInfo.OriginModelName))
  345. audioRatio := decimal.NewFromFloat(ratio_setting.GetAudioRatio(relayInfo.OriginModelName))
  346. audioCompletionRatio := decimal.NewFromFloat(ratio_setting.GetAudioCompletionRatio(relayInfo.OriginModelName))
  347. modelRatio := relayInfo.PriceData.ModelRatio
  348. groupRatio := relayInfo.PriceData.GroupRatioInfo.GroupRatio
  349. modelPrice := relayInfo.PriceData.ModelPrice
  350. usePrice := relayInfo.PriceData.UsePrice
  351. quotaInfo := QuotaInfo{
  352. InputDetails: TokenDetails{
  353. TextTokens: textInputTokens,
  354. AudioTokens: audioInputTokens,
  355. },
  356. OutputDetails: TokenDetails{
  357. TextTokens: textOutTokens,
  358. AudioTokens: audioOutTokens,
  359. },
  360. ModelName: relayInfo.OriginModelName,
  361. UsePrice: usePrice,
  362. ModelRatio: modelRatio,
  363. GroupRatio: groupRatio,
  364. }
  365. quota := calculateAudioQuota(quotaInfo)
  366. totalTokens := usage.TotalTokens
  367. var logContent string
  368. if !usePrice {
  369. logContent = fmt.Sprintf("模型倍率 %.2f,补全倍率 %.2f,音频倍率 %.2f,音频补全倍率 %.2f,分组倍率 %.2f",
  370. modelRatio, completionRatio.InexactFloat64(), audioRatio.InexactFloat64(), audioCompletionRatio.InexactFloat64(), groupRatio)
  371. } else {
  372. logContent = fmt.Sprintf("模型价格 %.2f,分组倍率 %.2f", modelPrice, groupRatio)
  373. }
  374. // record all the consume log even if quota is 0
  375. if totalTokens == 0 {
  376. // in this case, must be some error happened
  377. // we cannot just return, because we may have to return the pre-consumed quota
  378. quota = 0
  379. logContent += fmt.Sprintf("(可能是上游超时)")
  380. logger.LogError(ctx, fmt.Sprintf("total tokens is 0, cannot consume quota, userId %d, channelId %d, "+
  381. "tokenId %d, model %s, pre-consumed quota %d", relayInfo.UserId, relayInfo.ChannelId, relayInfo.TokenId, relayInfo.OriginModelName, relayInfo.FinalPreConsumedQuota))
  382. } else {
  383. model.UpdateUserUsedQuotaAndRequestCount(relayInfo.UserId, quota)
  384. model.UpdateChannelUsedQuota(relayInfo.ChannelId, quota)
  385. }
  386. if err := SettleBilling(ctx, relayInfo, quota); err != nil {
  387. logger.LogError(ctx, "error settling billing: "+err.Error())
  388. }
  389. logModel := relayInfo.OriginModelName
  390. if extraContent != "" {
  391. logContent += ", " + extraContent
  392. }
  393. other := GenerateAudioOtherInfo(ctx, relayInfo, usage, modelRatio, groupRatio,
  394. completionRatio.InexactFloat64(), audioRatio.InexactFloat64(), audioCompletionRatio.InexactFloat64(), modelPrice, relayInfo.PriceData.GroupRatioInfo.GroupSpecialRatio)
  395. model.RecordConsumeLog(ctx, relayInfo.UserId, model.RecordConsumeLogParams{
  396. ChannelId: relayInfo.ChannelId,
  397. PromptTokens: usage.PromptTokens,
  398. CompletionTokens: usage.CompletionTokens,
  399. ModelName: logModel,
  400. TokenName: tokenName,
  401. Quota: quota,
  402. Content: logContent,
  403. TokenId: relayInfo.TokenId,
  404. UseTimeSeconds: int(useTimeSeconds),
  405. IsStream: relayInfo.IsStream,
  406. Group: relayInfo.UsingGroup,
  407. Other: other,
  408. })
  409. }
  410. func PreConsumeTokenQuota(relayInfo *relaycommon.RelayInfo, quota int) error {
  411. if quota < 0 {
  412. return errors.New("quota 不能为负数!")
  413. }
  414. if relayInfo.IsPlayground {
  415. return nil
  416. }
  417. //if relayInfo.TokenUnlimited {
  418. // return nil
  419. //}
  420. token, err := model.GetTokenByKey(relayInfo.TokenKey, false)
  421. if err != nil {
  422. return err
  423. }
  424. if !relayInfo.TokenUnlimited && token.RemainQuota < quota {
  425. return fmt.Errorf("token quota is not enough, token remain quota: %s, need quota: %s", logger.FormatQuota(token.RemainQuota), logger.FormatQuota(quota))
  426. }
  427. err = model.DecreaseTokenQuota(relayInfo.TokenId, relayInfo.TokenKey, quota)
  428. if err != nil {
  429. return err
  430. }
  431. return nil
  432. }
  433. func PostConsumeQuota(relayInfo *relaycommon.RelayInfo, quota int, preConsumedQuota int, sendEmail bool) (err error) {
  434. // 1) Consume from wallet quota OR subscription item
  435. if relayInfo != nil && relayInfo.BillingSource == BillingSourceSubscription {
  436. if relayInfo.SubscriptionId == 0 {
  437. return errors.New("subscription id is missing")
  438. }
  439. delta := int64(quota)
  440. if delta != 0 {
  441. if err := model.PostConsumeUserSubscriptionDelta(relayInfo.SubscriptionId, delta); err != nil {
  442. return err
  443. }
  444. relayInfo.SubscriptionPostDelta += delta
  445. }
  446. } else {
  447. // Wallet
  448. if quota > 0 {
  449. err = model.DecreaseUserQuota(relayInfo.UserId, quota)
  450. } else {
  451. err = model.IncreaseUserQuota(relayInfo.UserId, -quota, false)
  452. }
  453. if err != nil {
  454. return err
  455. }
  456. }
  457. if !relayInfo.IsPlayground {
  458. if quota > 0 {
  459. err = model.DecreaseTokenQuota(relayInfo.TokenId, relayInfo.TokenKey, quota)
  460. } else {
  461. err = model.IncreaseTokenQuota(relayInfo.TokenId, relayInfo.TokenKey, -quota)
  462. }
  463. if err != nil {
  464. return err
  465. }
  466. }
  467. if sendEmail {
  468. if (quota + preConsumedQuota) != 0 {
  469. checkAndSendQuotaNotify(relayInfo, quota, preConsumedQuota)
  470. }
  471. }
  472. return nil
  473. }
  474. func checkAndSendQuotaNotify(relayInfo *relaycommon.RelayInfo, quota int, preConsumedQuota int) {
  475. gopool.Go(func() {
  476. userSetting := relayInfo.UserSetting
  477. threshold := common.QuotaRemindThreshold
  478. if userSetting.QuotaWarningThreshold != 0 {
  479. threshold = int(userSetting.QuotaWarningThreshold)
  480. }
  481. //noMoreQuota := userCache.Quota-(quota+preConsumedQuota) <= 0
  482. quotaTooLow := false
  483. consumeQuota := quota + preConsumedQuota
  484. if relayInfo.UserQuota-consumeQuota < threshold {
  485. quotaTooLow = true
  486. }
  487. if quotaTooLow {
  488. prompt := "您的额度即将用尽"
  489. topUpLink := fmt.Sprintf("%s/console/topup", system_setting.ServerAddress)
  490. // 根据通知方式生成不同的内容格式
  491. var content string
  492. var values []interface{}
  493. notifyType := userSetting.NotifyType
  494. if notifyType == "" {
  495. notifyType = dto.NotifyTypeEmail
  496. }
  497. if notifyType == dto.NotifyTypeBark {
  498. // Bark推送使用简短文本,不支持HTML
  499. content = "{{value}},剩余额度:{{value}},请及时充值"
  500. values = []interface{}{prompt, logger.FormatQuota(relayInfo.UserQuota)}
  501. } else if notifyType == dto.NotifyTypeGotify {
  502. content = "{{value}},当前剩余额度为 {{value}},请及时充值。"
  503. values = []interface{}{prompt, logger.FormatQuota(relayInfo.UserQuota)}
  504. } else {
  505. // 默认内容格式,适用于Email和Webhook(支持HTML)
  506. content = "{{value}},当前剩余额度为 {{value}},为了不影响您的使用,请及时充值。<br/>充值链接:<a href='{{value}}'>{{value}}</a>"
  507. values = []interface{}{prompt, logger.FormatQuota(relayInfo.UserQuota), topUpLink, topUpLink}
  508. }
  509. err := NotifyUser(relayInfo.UserId, relayInfo.UserEmail, relayInfo.UserSetting, dto.NewNotify(dto.NotifyTypeQuotaExceed, prompt, content, values))
  510. if err != nil {
  511. common.SysError(fmt.Sprintf("failed to send quota notify to user %d: %s", relayInfo.UserId, err.Error()))
  512. }
  513. }
  514. })
  515. }
  516. func checkAndSendSubscriptionQuotaNotify(relayInfo *relaycommon.RelayInfo) {
  517. gopool.Go(func() {
  518. if relayInfo == nil {
  519. return
  520. }
  521. if relayInfo.SubscriptionId == 0 || relayInfo.SubscriptionAmountTotal <= 0 {
  522. return
  523. }
  524. userSetting := relayInfo.UserSetting
  525. threshold := common.QuotaRemindThreshold
  526. if userSetting.QuotaWarningThreshold != 0 {
  527. threshold = int(userSetting.QuotaWarningThreshold)
  528. }
  529. usedAfter := relayInfo.SubscriptionAmountUsedAfterPreConsume + relayInfo.SubscriptionPostDelta
  530. remaining := relayInfo.SubscriptionAmountTotal - usedAfter
  531. if remaining >= int64(threshold) {
  532. return
  533. }
  534. prompt := "您的订阅额度即将用尽"
  535. topUpLink := fmt.Sprintf("%s/console/topup", system_setting.ServerAddress)
  536. var content string
  537. var values []interface{}
  538. notifyType := userSetting.NotifyType
  539. if notifyType == "" {
  540. notifyType = dto.NotifyTypeEmail
  541. }
  542. if notifyType == dto.NotifyTypeBark {
  543. content = "{{value}},剩余额度:{{value}},请及时充值"
  544. values = []interface{}{prompt, logger.FormatQuota(int(remaining))}
  545. } else if notifyType == dto.NotifyTypeGotify {
  546. content = "{{value}},当前剩余额度为 {{value}},请及时充值。"
  547. values = []interface{}{prompt, logger.FormatQuota(int(remaining))}
  548. } else {
  549. content = "{{value}},当前剩余额度为 {{value}},为了不影响您的使用,请及时充值。<br/>充值链接:<a href='{{value}}'>{{value}}</a>"
  550. values = []interface{}{prompt, logger.FormatQuota(int(remaining)), topUpLink, topUpLink}
  551. }
  552. if err := NotifyUser(relayInfo.UserId, relayInfo.UserEmail, relayInfo.UserSetting, dto.NewNotify(dto.NotifyTypeQuotaExceed, prompt, content, values)); err != nil {
  553. common.SysError(fmt.Sprintf("failed to send subscription quota notify to user %d: %s", relayInfo.UserId, err.Error()))
  554. }
  555. })
  556. }
  557. func postConsumeQuotaTieredService(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, modelName string,
  558. promptTokens, completionTokens, totalTokens, quota int, tieredResult *TieredResultWrapper, extraContent string) {
  559. useTimeSeconds := time.Now().Unix() - relayInfo.StartTime.Unix()
  560. tokenName := ctx.GetString("token_name")
  561. groupRatio := relayInfo.PriceData.GroupRatioInfo.GroupRatio
  562. var logContent string
  563. if totalTokens == 0 {
  564. quota = 0
  565. logContent = "上游没有返回计费信息(可能是上游超时)"
  566. logger.LogError(ctx, fmt.Sprintf("tiered billing: total tokens is 0, userId %d, channelId %d, tokenId %d, model %s, pre-consumed %d",
  567. relayInfo.UserId, relayInfo.ChannelId, relayInfo.TokenId, modelName, relayInfo.FinalPreConsumedQuota))
  568. } else {
  569. if groupRatio != 0 && quota == 0 {
  570. quota = 1
  571. }
  572. model.UpdateUserUsedQuotaAndRequestCount(relayInfo.UserId, quota)
  573. model.UpdateChannelUsedQuota(relayInfo.ChannelId, quota)
  574. }
  575. if err := SettleBilling(ctx, relayInfo, quota); err != nil {
  576. logger.LogError(ctx, "error settling tiered billing: "+err.Error())
  577. }
  578. if extraContent != "" {
  579. logContent += extraContent
  580. }
  581. other := GenerateTieredOtherInfo(ctx, relayInfo, tieredResult)
  582. model.RecordConsumeLog(ctx, relayInfo.UserId, model.RecordConsumeLogParams{
  583. ChannelId: relayInfo.ChannelId,
  584. PromptTokens: promptTokens,
  585. CompletionTokens: completionTokens,
  586. ModelName: modelName,
  587. TokenName: tokenName,
  588. Quota: quota,
  589. Content: logContent,
  590. TokenId: relayInfo.TokenId,
  591. UseTimeSeconds: int(useTimeSeconds),
  592. IsStream: relayInfo.IsStream,
  593. Group: relayInfo.UsingGroup,
  594. Other: other,
  595. })
  596. }