From b3f46223a83104fd2598ef92e5dd4a5fc4647c79 Mon Sep 17 00:00:00 2001
From: CaIon <1808837298@qq.com>
Date: Sun, 21 Jan 2024 20:44:26 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E8=AF=B7=E6=B1=82=E5=87=BA=E7=8E=B0=20?=
 =?UTF-8?q?0=20token=E7=9A=84=E6=97=B6=E5=80=99=EF=BC=8C=E5=8A=A0=E5=85=A5?=
 =?UTF-8?q?=E9=94=99=E8=AF=AF=E6=8F=90=E7=A4=BA=E5=B9=B6=E6=89=93=E5=8D=B0?=
 =?UTF-8?q?=E6=97=A5=E5=BF=97?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 controller/relay-text.go | 44 ++++++++++++++++++++++------------------
 1 file changed, 24 insertions(+), 20 deletions(-)

diff --git a/controller/relay-text.go b/controller/relay-text.go
index 6538028..a86bc79 100644
--- a/controller/relay-text.go
+++ b/controller/relay-text.go
@@ -474,6 +474,7 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
 	defer func(ctx context.Context) {
 		// c.Writer.Flush()
 		go func() {
+			useTimeSeconds := time.Now().Unix() - startTime.Unix()
 			promptTokens = textResponse.Usage.PromptTokens
 			completionTokens = textResponse.Usage.CompletionTokens
 
@@ -489,38 +490,41 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
 				quota = int(modelPrice * common.QuotaPerUnit * groupRatio)
 			}
 			totalTokens := promptTokens + completionTokens
-			if totalTokens == 0 {
-				// in this case, must be some error happened
-				// we cannot just return, because we may have to return the pre-consumed quota
-				quota = 0
-			}
-			quotaDelta := quota - preConsumedQuota
-			err := model.PostConsumeTokenQuota(tokenId, userQuota, quotaDelta, preConsumedQuota, true)
-			if err != nil {
-				common.LogError(ctx, "error consuming token remain quota: "+err.Error())
-			}
-			err = model.CacheUpdateUserQuota(userId)
-			if err != nil {
-				common.LogError(ctx, "error update user quota cache: "+err.Error())
-			}
-
-			// record all the consume log even if quota is 0
-			useTimeSeconds := time.Now().Unix() - startTime.Unix()
 			var logContent string
 			if modelPrice == -1 {
 				logContent = fmt.Sprintf("模型倍率 %.2f，分组倍率 %.2f", modelRatio, groupRatio)
 			} else {
 				logContent = fmt.Sprintf("模型价格 %.2f，分组倍率 %.2f", modelPrice, groupRatio)
 			}
+
+			// record all the consume log even if quota is 0
+			if totalTokens == 0 {
+				// in this case, must be some error happened
+				// we cannot just return, because we may have to return the pre-consumed quota
+				quota = 0
+				logContent += fmt.Sprintf("（有疑问请联系管理员）")
+				common.LogError(ctx, fmt.Sprintf("total tokens is 0, cannot consume quota, userId %d, channelId %d, tokenId %d, model %s， pre-consumed quota %d", userId, channelId, tokenId, textRequest.Model, preConsumedQuota))
+			} else {
+				quotaDelta := quota - preConsumedQuota
+				err := model.PostConsumeTokenQuota(tokenId, userQuota, quotaDelta, preConsumedQuota, true)
+				if err != nil {
+					common.LogError(ctx, "error consuming token remain quota: "+err.Error())
+				}
+				err = model.CacheUpdateUserQuota(userId)
+				if err != nil {
+					common.LogError(ctx, "error update user quota cache: "+err.Error())
+				}
+				model.UpdateUserUsedQuotaAndRequestCount(userId, quota)
+				model.UpdateChannelUsedQuota(channelId, quota)
+			}
+
 			logModel := textRequest.Model
 			if strings.HasPrefix(logModel, "gpt-4-gizmo") {
 				logModel = "gpt-4-gizmo-*"
 				logContent += fmt.Sprintf("，模型 %s", textRequest.Model)
 			}
-
 			model.RecordConsumeLog(ctx, userId, channelId, promptTokens, completionTokens, logModel, tokenName, quota, logContent, tokenId, userQuota, int(useTimeSeconds), isStream)
-			model.UpdateUserUsedQuotaAndRequestCount(userId, quota)
-			model.UpdateChannelUsedQuota(channelId, quota)
+
 			//if quota != 0 {
 			//
 			//}