xinyin025/new-api
1
0
Fork 0
mirror of https://github.com/linux-do/new-api.git synced 2025-09-22 02:06:38 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

feat: realtime扣费时检测令牌额度

Browse Source 
(cherry picked from commit 91511b8b64fc0d28dbf657cb97e12b7d1e50070d)
This commit is contained in:
CalciumIon 2024-11-07 17:28:21 +08:00
parent 3b53a2a5ce
commit 4131183378
1 changed files with 11 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
service/quota.go
View File

@ -9,6 +9,7 @@ import (
"one-api/dto"
"one-api/model"
relaycommon "one-api/relay/common"
"strings"
"time"
)
@ -20,6 +21,12 @@ func PreWssConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usag
if err != nil {
return err
}
token, err := model.CacheGetTokenByKey(strings.TrimLeft(relayInfo.ApiKey, "sk-"))
if err != nil {
return err
}
modelName := relayInfo.UpstreamModelName
textInputTokens := usage.InputTokenDetails.TextTokens
textOutTokens := usage.OutputTokenDetails.TextTokens
@ -46,6 +53,10 @@ func PreWssConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usag
return errors.New(fmt.Sprintf("用户额度不足，剩余额度为 %d", userQuota))
}
if token.RemainQuota < quota {
return errors.New(fmt.Sprintf("令牌额度不足，剩余额度为 %d", token.RemainQuota))
}
err = model.PostConsumeTokenQuota(relayInfo, 0, quota, 0, false)
if err != nil {
return err