From 413118337839dbc0faa2b46fc193d1a4b31f2647 Mon Sep 17 00:00:00 2001
From: CalciumIon <1808837298@qq.com>
Date: Thu, 7 Nov 2024 17:28:21 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20realtime=E6=89=A3=E8=B4=B9=E6=97=B6?=
 =?UTF-8?q?=E6=A3=80=E6=B5=8B=E4=BB=A4=E7=89=8C=E9=A2=9D=E5=BA=A6?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

(cherry picked from commit 91511b8b64fc0d28dbf657cb97e12b7d1e50070d)
---
 service/quota.go | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/service/quota.go b/service/quota.go
index 695c073..0ca9145 100644
--- a/service/quota.go
+++ b/service/quota.go
@@ -9,6 +9,7 @@ import (
 	"one-api/dto"
 	"one-api/model"
 	relaycommon "one-api/relay/common"
+	"strings"
 	"time"
 )
 
@@ -20,6 +21,12 @@ func PreWssConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usag
 	if err != nil {
 		return err
 	}
+
+	token, err := model.CacheGetTokenByKey(strings.TrimLeft(relayInfo.ApiKey, "sk-"))
+	if err != nil {
+		return err
+	}
+
 	modelName := relayInfo.UpstreamModelName
 	textInputTokens := usage.InputTokenDetails.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))
 	}
 
+	if token.RemainQuota < quota {
+		return errors.New(fmt.Sprintf("令牌额度不足，剩余额度为 %d", token.RemainQuota))
+	}
+
 	err = model.PostConsumeTokenQuota(relayInfo, 0, quota, 0, false)
 	if err != nil {
 		return err