From 64e9e9cc20037faf1d975520a2c1de59f6aca3bf Mon Sep 17 00:00:00 2001
From: CaIon <1808837298@qq.com>
Date: Thu, 11 Jan 2024 14:12:38 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20=E5=AE=8C=E5=96=84=E4=BB=A4=E7=89=8C?=
 =?UTF-8?q?=E9=A2=84=E6=89=A3=E8=B4=B9=E9=80=BB=E8=BE=91?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 controller/relay-text.go | 5 +++--
 middleware/auth.go       | 3 +--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/controller/relay-text.go b/controller/relay-text.go
index f33ebfb..c68966e 100644
--- a/controller/relay-text.go
+++ b/controller/relay-text.go
@@ -52,7 +52,7 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
 	tokenId := c.GetInt("token_id")
 	userId := c.GetInt("id")
 	group := c.GetString("group")
-	tokenQuota := c.GetInt("token_quota")
+	tokenUnlimited := c.GetBool("token_unlimited_quota")
 	startTime := time.Now()
 	var textRequest GeneralOpenAIRequest
 
@@ -263,8 +263,9 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
 	}
 	if userQuota > 100*preConsumedQuota {
 		// 用户额度充足，判断令牌额度是否充足
-		if tokenQuota != -1 {
+		if !tokenUnlimited {
 			// 非无限令牌，判断令牌额度是否充足
+			tokenQuota := c.GetInt("token_quota")
 			if tokenQuota > 100*preConsumedQuota {
 				// 令牌额度充足，信任令牌
 				preConsumedQuota = 0
diff --git a/middleware/auth.go b/middleware/auth.go
index 1b97f9f..ef774f6 100644
--- a/middleware/auth.go
+++ b/middleware/auth.go
@@ -115,10 +115,9 @@ func TokenAuth() func(c *gin.Context) {
 		c.Set("id", token.UserId)
 		c.Set("token_id", token.Id)
 		c.Set("token_name", token.Name)
+		c.Set("token_unlimited_quota", token.UnlimitedQuota)
 		if !token.UnlimitedQuota {
 			c.Set("token_quota", token.RemainQuota)
-		} else {
-			c.Set("token_quota", -1)
 		}
 		if token.ModelLimitsEnabled {
 			c.Set("token_model_limit_enabled", true)