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

feat: realtime pre consume

(cherry picked from commit 273d154e1640bae26b7caedddf1685e9ff21ab74)
1808837298@qq.com 1 год назад
Родитель
Сommit
e5c05d77b7
2 измененных файлов с 11 добавлено и 4 удалено
  1. 0 3
      relay/channel/openai/relay-openai.go
  2. 11 1
      service/quota.go

+ 0 - 3
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
 }

+ 11 - 1
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