Просмотр исходного кода

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

(cherry picked from commit 91511b8b64fc0d28dbf657cb97e12b7d1e50070d)
CalciumIon 1 год назад
Родитель
Сommit
4131183378
1 измененных файлов с 11 добавлено и 0 удалено
  1. 11 0
      service/quota.go

+ 11 - 0
service/quota.go

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