diff --git a/relay/channel/openai/relay-openai.go b/relay/channel/openai/relay-openai.go index 60d09a0..bb19684 100644 --- a/relay/channel/openai/relay-openai.go +++ b/relay/channel/openai/relay-openai.go @@ -576,8 +576,5 @@ func preConsumeUsage(ctx *gin.Context, info *relaycommon.RelayInfo, usage *dto.R totalUsage.OutputTokenDetails.AudioTokens += usage.OutputTokenDetails.AudioTokens // clear usage err := service.PreWssConsumeQuota(ctx, info, usage) - if err == nil { - common.LogInfo(ctx, "realtime streaming consume usage success") - } return err } diff --git a/service/quota.go b/service/quota.go index 09c2fd5..9a3c542 100644 --- a/service/quota.go +++ b/service/quota.go @@ -1,6 +1,7 @@ package service import ( + "errors" "fmt" "github.com/gin-gonic/gin" "math" @@ -16,6 +17,10 @@ func PreWssConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usag if relayInfo.UsePrice { return nil } + userQuota, err := model.GetUserQuota(relayInfo.UserId) + if err != nil { + return err + } modelName := relayInfo.UpstreamModelName textInputTokens := usage.InputTokenDetails.TextTokens textOutTokens := usage.OutputTokenDetails.TextTokens @@ -38,10 +43,15 @@ func PreWssConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usag 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 { return err } + common.LogInfo(ctx, "realtime streaming consume quota success, quota: "+fmt.Sprintf("%d", quota)) err = model.CacheUpdateUserQuota(relayInfo.UserId) if err != nil { return err