diff --git a/common/redis.go b/common/redis.go index 1a05721c..12c477b8 100644 --- a/common/redis.go +++ b/common/redis.go @@ -61,3 +61,8 @@ func RedisDel(key string) error { ctx := context.Background() return RDB.Del(ctx, key).Err() } + +func RedisDecrease(key string, value int64) error { + ctx := context.Background() + return RDB.DecrBy(ctx, key, value).Err() +} diff --git a/controller/relay-text.go b/controller/relay-text.go index e8dab514..761ca86f 100644 --- a/controller/relay-text.go +++ b/controller/relay-text.go @@ -194,7 +194,11 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode { if err != nil { return errorWrapper(err, "get_user_quota_failed", http.StatusInternalServerError) } - if userQuota > 10*preConsumedQuota { + err = model.CacheDecreaseUserQuota(userId, preConsumedQuota) + if err != nil { + return errorWrapper(err, "decrease_user_quota_failed", http.StatusInternalServerError) + } + if userQuota > 100*preConsumedQuota { // in this case, we do not pre-consume quota // because the user has enough quota preConsumedQuota = 0 diff --git a/model/cache.go b/model/cache.go index 64666c86..55fbba9b 100644 --- a/model/cache.go +++ b/model/cache.go @@ -95,6 +95,14 @@ func CacheUpdateUserQuota(id int) error { return err } +func CacheDecreaseUserQuota(id int, quota int) error { + if !common.RedisEnabled { + return nil + } + err := common.RedisDecrease(fmt.Sprintf("user_quota:%d", id), int64(quota)) + return err +} + func CacheIsUserEnabled(userId int) bool { if !common.RedisEnabled { return IsUserEnabled(userId)