From 16ad764f9b7162b9032435b2e6163f7157281b52 Mon Sep 17 00:00:00 2001
From: CaIon <1808837298@qq.com>
Date: Wed, 15 Nov 2023 18:27:13 +0800
Subject: [PATCH] try to fix email

---
 controller/relay-audio.go |  2 +-
 controller/relay-image.go |  2 +-
 controller/relay-mj.go    |  2 +-
 controller/relay-text.go  |  4 ++--
 model/token.go            | 26 ++++++++++++++++----------
 5 files changed, 21 insertions(+), 15 deletions(-)

diff --git a/controller/relay-audio.go b/controller/relay-audio.go
index fe91dbc..13d9c9f 100644
--- a/controller/relay-audio.go
+++ b/controller/relay-audio.go
@@ -99,7 +99,7 @@ func relayAudioHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode
 		go func() {
 			quota := countTokenText(audioResponse.Text, audioModel)
 			quotaDelta := quota - preConsumedQuota
-			err := model.PostConsumeTokenQuota(tokenId, userQuota, quotaDelta, preConsumedQuota)
+			err := model.PostConsumeTokenQuota(tokenId, userQuota, quotaDelta, preConsumedQuota, true)
 			if err != nil {
 				common.SysError("error consuming token remain quota: " + err.Error())
 			}
diff --git a/controller/relay-image.go b/controller/relay-image.go
index 5cebcdb..8c16ec1 100644
--- a/controller/relay-image.go
+++ b/controller/relay-image.go
@@ -147,7 +147,7 @@ func relayImageHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode
 	var textResponse ImageResponse
 	defer func(ctx context.Context) {
 		if consumeQuota {
-			err := model.PostConsumeTokenQuota(tokenId, userId, quota, 0)
+			err := model.PostConsumeTokenQuota(tokenId, userId, quota, 0, true)
 			if err != nil {
 				common.SysError("error consuming token remain quota: " + err.Error())
 			}
diff --git a/controller/relay-mj.go b/controller/relay-mj.go
index 948c57c..89b0f0c 100644
--- a/controller/relay-mj.go
+++ b/controller/relay-mj.go
@@ -359,7 +359,7 @@ func relayMidjourneySubmit(c *gin.Context, relayMode int) *MidjourneyResponse {
 
 	defer func(ctx context.Context) {
 		if consumeQuota {
-			err := model.PostConsumeTokenQuota(tokenId, userQuota, quota, 0)
+			err := model.PostConsumeTokenQuota(tokenId, userQuota, quota, 0, true)
 			if err != nil {
 				common.SysError("error consuming token remain quota: " + err.Error())
 			}
diff --git a/controller/relay-text.go b/controller/relay-text.go
index 6f56be8..2729650 100644
--- a/controller/relay-text.go
+++ b/controller/relay-text.go
@@ -400,7 +400,7 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
 			if preConsumedQuota != 0 {
 				go func(ctx context.Context) {
 					// return pre-consumed quota
-					err := model.PostConsumeTokenQuota(tokenId, userQuota, -preConsumedQuota, 0)
+					err := model.PostConsumeTokenQuota(tokenId, userQuota, -preConsumedQuota, 0, false)
 					if err != nil {
 						common.LogError(ctx, "error return pre-consumed quota: "+err.Error())
 					}
@@ -434,7 +434,7 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
 					quota = 0
 				}
 				quotaDelta := quota - preConsumedQuota
-				err := model.PostConsumeTokenQuota(tokenId, userQuota, quotaDelta, preConsumedQuota)
+				err := model.PostConsumeTokenQuota(tokenId, userQuota, quotaDelta, preConsumedQuota, true)
 				if err != nil {
 					common.LogError(ctx, "error consuming token remain quota: "+err.Error())
 				}
diff --git a/model/token.go b/model/token.go
index 06b9775..5c4bc55 100644
--- a/model/token.go
+++ b/model/token.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"gorm.io/gorm"
 	"one-api/common"
+	"strconv"
 	"strings"
 )
 
@@ -194,22 +195,31 @@ func PreConsumeTokenQuota(tokenId int, quota int) (userQuota int, err error) {
 		return 0, err
 	}
 	if userQuota < quota {
-		return userQuota, errors.New(fmt.Sprintf("用户额度不足，剩余额度为 %d", userQuota))
+		return 0, errors.New(fmt.Sprintf("用户额度不足，剩余额度为 %d", userQuota))
 	}
 	if !token.UnlimitedQuota {
 		err = DecreaseTokenQuota(tokenId, quota)
 		if err != nil {
-			return userQuota, err
+			return 0, err
 		}
 	}
 	err = DecreaseUserQuota(token.UserId, quota)
-	return userQuota, err
+	return userQuota - quota, err
 }
 
-func PostConsumeTokenQuota(tokenId int, userQuota int, quota int, preConsumedQuota int) (err error) {
+func PostConsumeTokenQuota(tokenId int, userQuota int, quota int, preConsumedQuota int, sendEmail bool) (err error) {
 	token, err := GetTokenById(tokenId)
 
 	if quota > 0 {
+		err = DecreaseUserQuota(token.UserId, quota)
+	} else {
+		err = IncreaseUserQuota(token.UserId, -quota)
+	}
+	if err != nil {
+		return err
+	}
+
+	if sendEmail {
 		quotaTooLow := userQuota >= common.QuotaRemindThreshold && userQuota-(quota+preConsumedQuota) < common.QuotaRemindThreshold
 		noMoreQuota := userQuota-(quota+preConsumedQuota) <= 0
 		if quotaTooLow || noMoreQuota {
@@ -229,16 +239,12 @@ func PostConsumeTokenQuota(tokenId int, userQuota int, quota int, preConsumedQuo
 					if err != nil {
 						common.SysError("failed to send email" + err.Error())
 					}
+					common.SysLog("user quota is low, consumed quota: " + strconv.Itoa(quota) + ", user quota: " + strconv.Itoa(userQuota))
 				}
 			}()
 		}
-		err = DecreaseUserQuota(token.UserId, quota)
-	} else {
-		err = IncreaseUserQuota(token.UserId, -quota)
-	}
-	if err != nil {
-		return err
 	}
+
 	if !token.UnlimitedQuota {
 		if quota > 0 {
 			err = DecreaseTokenQuota(tokenId, quota)