From a3687b72f87a35024e093c5ef5ca83d4b458989e Mon Sep 17 00:00:00 2001 From: "1808837298@qq.com" <1808837298@qq.com> Date: Thu, 29 Feb 2024 16:39:52 +0800 Subject: [PATCH] fix: fix preConsumeQuota error --- relay/relay-text.go | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/relay/relay-text.go b/relay/relay-text.go index d864c84..63b1ff6 100644 --- a/relay/relay-text.go +++ b/relay/relay-text.go @@ -113,7 +113,7 @@ func TextHelper(c *gin.Context) *dto.OpenAIErrorWithStatusCode { } // pre-consume quota 预消耗配额 - userQuota, openaiErr := preConsumeQuota(c, preConsumedQuota, relayInfo) + preConsumedQuota, userQuota, openaiErr := preConsumeQuota(c, preConsumedQuota, relayInfo) if err != nil { return openaiErr } @@ -153,7 +153,6 @@ func TextHelper(c *gin.Context) *dto.OpenAIErrorWithStatusCode { if openaiErr != nil { return openaiErr } - postConsumeQuota(c, relayInfo, *textRequest, usage, ratio, preConsumedQuota, userQuota, modelRatio, groupRatio, modelPrice) return nil } @@ -178,17 +177,17 @@ func getPromptTokens(textRequest *dto.GeneralOpenAIRequest, info *relaycommon.Re } // 预扣费并返回用户剩余配额 -func preConsumeQuota(c *gin.Context, preConsumedQuota int, relayInfo *relaycommon.RelayInfo) (int, *dto.OpenAIErrorWithStatusCode) { +func preConsumeQuota(c *gin.Context, preConsumedQuota int, relayInfo *relaycommon.RelayInfo) (int, int, *dto.OpenAIErrorWithStatusCode) { userQuota, err := model.CacheGetUserQuota(relayInfo.UserId) if err != nil { - return 0, service.OpenAIErrorWrapper(err, "get_user_quota_failed", http.StatusInternalServerError) + return 0, 0, service.OpenAIErrorWrapper(err, "get_user_quota_failed", http.StatusInternalServerError) } if userQuota < 0 || userQuota-preConsumedQuota < 0 { - return 0, service.OpenAIErrorWrapper(errors.New("user quota is not enough"), "insufficient_user_quota", http.StatusForbidden) + return 0, 0, service.OpenAIErrorWrapper(errors.New("user quota is not enough"), "insufficient_user_quota", http.StatusForbidden) } err = model.CacheDecreaseUserQuota(relayInfo.UserId, preConsumedQuota) if err != nil { - return 0, service.OpenAIErrorWrapper(err, "decrease_user_quota_failed", http.StatusInternalServerError) + return 0, 0, service.OpenAIErrorWrapper(err, "decrease_user_quota_failed", http.StatusInternalServerError) } if userQuota > 100*preConsumedQuota { // 用户额度充足，判断令牌额度是否充足 @@ -210,10 +209,10 @@ func preConsumeQuota(c *gin.Context, preConsumedQuota int, relayInfo *relaycommo if preConsumedQuota > 0 { userQuota, err = model.PreConsumeTokenQuota(relayInfo.TokenId, preConsumedQuota) if err != nil { - return 0, service.OpenAIErrorWrapper(err, "pre_consume_token_quota_failed", http.StatusForbidden) + return 0, 0, service.OpenAIErrorWrapper(err, "pre_consume_token_quota_failed", http.StatusForbidden) } } - return userQuota, nil + return preConsumedQuota, userQuota, nil } func postConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, textRequest dto.GeneralOpenAIRequest, usage *dto.Usage, ratio float64, preConsumedQuota int, userQuota int, modelRatio float64, groupRatio float64, modelPrice float64) {