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

try to fix email

Browse Source
This commit is contained in:
CaIon 2023-11-15 18:27:13 +08:00
parent b4bd9a19d9
commit 16ad764f9b
5 changed files with 21 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
controller/relay-audio.go
View File

@ -99,7 +99,7 @@ func relayAudioHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode
go func() { go func() {
quota := countTokenText(audioResponse.Text, audioModel) quota := countTokenText(audioResponse.Text, audioModel)
quotaDelta := quota - preConsumedQuota quotaDelta := quota - preConsumedQuota
err := model.PostConsumeTokenQuota(tokenId, userQuota, quotaDelta, preConsumedQuota) err := model.PostConsumeTokenQuota(tokenId, userQuota, quotaDelta, preConsumedQuota, true)
if err != nil { if err != nil {
common.SysError("error consuming token remain quota: " + err.Error()) common.SysError("error consuming token remain quota: " + err.Error())
} }

2
controller/relay-image.go
View File

@ -147,7 +147,7 @@ func relayImageHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode
var textResponse ImageResponse var textResponse ImageResponse
defer func(ctx context.Context) { defer func(ctx context.Context) {
if consumeQuota { if consumeQuota {
err := model.PostConsumeTokenQuota(tokenId, userId, quota, 0) err := model.PostConsumeTokenQuota(tokenId, userId, quota, 0, true)
if err != nil { if err != nil {
common.SysError("error consuming token remain quota: " + err.Error()) common.SysError("error consuming token remain quota: " + err.Error())
} }

2
controller/relay-mj.go
View File

@ -359,7 +359,7 @@ func relayMidjourneySubmit(c *gin.Context, relayMode int) *MidjourneyResponse {
defer func(ctx context.Context) { defer func(ctx context.Context) {
if consumeQuota { if consumeQuota {
err := model.PostConsumeTokenQuota(tokenId, userQuota, quota, 0) err := model.PostConsumeTokenQuota(tokenId, userQuota, quota, 0, true)
if err != nil { if err != nil {
common.SysError("error consuming token remain quota: " + err.Error()) common.SysError("error consuming token remain quota: " + err.Error())
} }

4
controller/relay-text.go
View File

@ -400,7 +400,7 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
if preConsumedQuota != 0 { if preConsumedQuota != 0 {
go func(ctx context.Context) { go func(ctx context.Context) {
// return pre-consumed quota // return pre-consumed quota
err := model.PostConsumeTokenQuota(tokenId, userQuota, -preConsumedQuota, 0) err := model.PostConsumeTokenQuota(tokenId, userQuota, -preConsumedQuota, 0, false)
if err != nil { if err != nil {
common.LogError(ctx, "error return pre-consumed quota: "+err.Error()) common.LogError(ctx, "error return pre-consumed quota: "+err.Error())
} }
@ -434,7 +434,7 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
quota = 0 quota = 0
} }
quotaDelta := quota - preConsumedQuota quotaDelta := quota - preConsumedQuota
err := model.PostConsumeTokenQuota(tokenId, userQuota, quotaDelta, preConsumedQuota) err := model.PostConsumeTokenQuota(tokenId, userQuota, quotaDelta, preConsumedQuota, true)
if err != nil { if err != nil {
common.LogError(ctx, "error consuming token remain quota: "+err.Error()) common.LogError(ctx, "error consuming token remain quota: "+err.Error())
} }

26
model/token.go
View File

@ -5,6 +5,7 @@ import (
"fmt" "fmt"
"gorm.io/gorm" "gorm.io/gorm"
"one-api/common" "one-api/common"
"strconv"
"strings" "strings"
) )
@ -194,22 +195,31 @@ func PreConsumeTokenQuota(tokenId int, quota int) (userQuota int, err error) {
return 0, err return 0, err
} }
if userQuota < quota { if userQuota < quota {
return userQuota, errors.New(fmt.Sprintf("用户额度不足，剩余额度为 %d", userQuota)) return 0, errors.New(fmt.Sprintf("用户额度不足，剩余额度为 %d", userQuota))
} }
if !token.UnlimitedQuota { if !token.UnlimitedQuota {
err = DecreaseTokenQuota(tokenId, quota) err = DecreaseTokenQuota(tokenId, quota)
if err != nil { if err != nil {
return userQuota, err return 0, err
} }
} }
err = DecreaseUserQuota(token.UserId, quota) err = DecreaseUserQuota(token.UserId, quota)
return userQuota, err return userQuota - quota, err
} }
func PostConsumeTokenQuota(tokenId int, userQuota int, quota int, preConsumedQuota int) (err error) { func PostConsumeTokenQuota(tokenId int, userQuota int, quota int, preConsumedQuota int, sendEmail bool) (err error) {
token, err := GetTokenById(tokenId) token, err := GetTokenById(tokenId)
if quota > 0 { if quota > 0 {
err = DecreaseUserQuota(token.UserId, quota)
} else {
err = IncreaseUserQuota(token.UserId, -quota)
}
if err != nil {
return err
}
if sendEmail {
quotaTooLow := userQuota >= common.QuotaRemindThreshold && userQuota-(quota+preConsumedQuota) < common.QuotaRemindThreshold quotaTooLow := userQuota >= common.QuotaRemindThreshold && userQuota-(quota+preConsumedQuota) < common.QuotaRemindThreshold
noMoreQuota := userQuota-(quota+preConsumedQuota) <= 0 noMoreQuota := userQuota-(quota+preConsumedQuota) <= 0
if quotaTooLow || noMoreQuota { if quotaTooLow || noMoreQuota {
@ -229,16 +239,12 @@ func PostConsumeTokenQuota(tokenId int, userQuota int, quota int, preConsumedQuo
if err != nil { if err != nil {
common.SysError("failed to send email" + err.Error()) common.SysError("failed to send email" + err.Error())
} }
common.SysLog("user quota is low, consumed quota: " + strconv.Itoa(quota) + ", user quota: " + strconv.Itoa(userQuota))
} }
}() }()
} }
err = DecreaseUserQuota(token.UserId, quota)
} else {
err = IncreaseUserQuota(token.UserId, -quota)
}
if err != nil {
return err
} }
if !token.UnlimitedQuota { if !token.UnlimitedQuota {
if quota > 0 { if quota > 0 {
err = DecreaseTokenQuota(tokenId, quota) err = DecreaseTokenQuota(tokenId, quota)