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

feat: realtime pre consume

Browse Source 
(cherry picked from commit 273d154e1640bae26b7caedddf1685e9ff21ab74)
This commit is contained in:
1808837298@qq.com 2024-10-07 19:15:14 +08:00 committed by CalciumIon
parent 24b3ed50d7
commit e5c05d77b7
2 changed files with 11 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
relay/channel/openai/relay-openai.go
View File

@ -576,8 +576,5 @@ func preConsumeUsage(ctx *gin.Context, info *relaycommon.RelayInfo, usage *dto.R
totalUsage.OutputTokenDetails.AudioTokens += usage.OutputTokenDetails.AudioTokens totalUsage.OutputTokenDetails.AudioTokens += usage.OutputTokenDetails.AudioTokens
// clear usage // clear usage
err := service.PreWssConsumeQuota(ctx, info, usage) err := service.PreWssConsumeQuota(ctx, info, usage)
if err == nil {
common.LogInfo(ctx, "realtime streaming consume usage success")
}
return err return err
} }

12
service/quota.go
View File

@ -1,6 +1,7 @@
package service package service
import ( import (
"errors"
"fmt" "fmt"
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"math" "math"
@ -16,6 +17,10 @@ func PreWssConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usag
if relayInfo.UsePrice { if relayInfo.UsePrice {
return nil return nil
} }
userQuota, err := model.GetUserQuota(relayInfo.UserId)
if err != nil {
return err
}
modelName := relayInfo.UpstreamModelName modelName := relayInfo.UpstreamModelName
textInputTokens := usage.InputTokenDetails.TextTokens textInputTokens := usage.InputTokenDetails.TextTokens
textOutTokens := usage.OutputTokenDetails.TextTokens textOutTokens := usage.OutputTokenDetails.TextTokens
@ -38,10 +43,15 @@ func PreWssConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usag
quota = 1 quota = 1
} }
err := model.PostConsumeTokenQuota(relayInfo, 0, quota, 0, false) if userQuota < quota {
return errors.New(fmt.Sprintf("用户额度不足，剩余额度为 %d", userQuota))
}
err = model.PostConsumeTokenQuota(relayInfo, 0, quota, 0, false)
if err != nil { if err != nil {
return err return err
} }
common.LogInfo(ctx, "realtime streaming consume quota success, quota: "+fmt.Sprintf("%d", quota))
err = model.CacheUpdateUserQuota(relayInfo.UserId) err = model.CacheUpdateUserQuota(relayInfo.UserId)
if err != nil { if err != nil {
return err return err