78 lines
3.2 KiB
Go
78 lines
3.2 KiB
Go
package service
|
|
|
|
import (
|
|
"errors"
|
|
"fmt"
|
|
"github.com/bytedance/gopkg/util/gopool"
|
|
"github.com/gin-gonic/gin"
|
|
"net/http"
|
|
"one-api/common"
|
|
"one-api/logger"
|
|
"one-api/model"
|
|
relaycommon "one-api/relay/common"
|
|
"one-api/types"
|
|
)
|
|
|
|
func ReturnPreConsumedQuota(c *gin.Context, relayInfo *relaycommon.RelayInfo, preConsumedQuota int) {
|
|
if preConsumedQuota != 0 {
|
|
gopool.Go(func() {
|
|
relayInfoCopy := *relayInfo
|
|
|
|
err := PostConsumeQuota(&relayInfoCopy, -preConsumedQuota, 0, false)
|
|
if err != nil {
|
|
common.SysLog("error return pre-consumed quota: " + err.Error())
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
// PreConsumeQuota checks if the user has enough quota to pre-consume.
|
|
// It returns the pre-consumed quota if successful, or an error if not.
|
|
func PreConsumeQuota(c *gin.Context, preConsumedQuota int, relayInfo *relaycommon.RelayInfo) (int, *types.NewAPIError) {
|
|
userQuota, err := model.GetUserQuota(relayInfo.UserId, false)
|
|
if err != nil {
|
|
return 0, types.NewError(err, types.ErrorCodeQueryDataError, types.ErrOptionWithSkipRetry())
|
|
}
|
|
if userQuota <= 0 {
|
|
return 0, types.NewErrorWithStatusCode(errors.New("user quota is not enough"), types.ErrorCodeInsufficientUserQuota, http.StatusForbidden, types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog())
|
|
}
|
|
if userQuota-preConsumedQuota < 0 {
|
|
return 0, types.NewErrorWithStatusCode(fmt.Errorf("预扣费额度失败, 用户剩余额度: %s, 需要预扣费额度: %s", logger.FormatQuota(userQuota), logger.FormatQuota(preConsumedQuota)), types.ErrorCodeInsufficientUserQuota, http.StatusForbidden, types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog())
|
|
}
|
|
|
|
trustQuota := common.GetTrustQuota()
|
|
|
|
relayInfo.UserQuota = userQuota
|
|
if userQuota > trustQuota {
|
|
// 用户额度充足，判断令牌额度是否充足
|
|
if !relayInfo.TokenUnlimited {
|
|
// 非无限令牌，判断令牌额度是否充足
|
|
tokenQuota := c.GetInt("token_quota")
|
|
if tokenQuota > trustQuota {
|
|
// 令牌额度充足，信任令牌
|
|
preConsumedQuota = 0
|
|
logger.LogInfo(c, fmt.Sprintf("user %d quota %s and token %d quota %d are enough, trusted and no need to pre-consume", relayInfo.UserId, logger.FormatQuota(userQuota), relayInfo.TokenId, tokenQuota))
|
|
}
|
|
} else {
|
|
// in this case, we do not pre-consume quota
|
|
// because the user has enough quota
|
|
preConsumedQuota = 0
|
|
logger.LogInfo(c, fmt.Sprintf("user %d with unlimited token has enough quota %s, trusted and no need to pre-consume", relayInfo.UserId, logger.FormatQuota(userQuota)))
|
|
}
|
|
}
|
|
|
|
if preConsumedQuota > 0 {
|
|
err := PreConsumeTokenQuota(relayInfo, preConsumedQuota)
|
|
if err != nil {
|
|
return 0, types.NewErrorWithStatusCode(err, types.ErrorCodePreConsumeTokenQuotaFailed, http.StatusForbidden, types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog())
|
|
}
|
|
err = model.DecreaseUserQuota(relayInfo.UserId, preConsumedQuota)
|
|
if err != nil {
|
|
return 0, types.NewError(err, types.ErrorCodeUpdateDataError, types.ErrOptionWithSkipRetry())
|
|
}
|
|
logger.LogInfo(c, fmt.Sprintf("用户 %d 预扣费 %s, 预扣费后剩余额度: %s", relayInfo.UserId, logger.FormatQuota(preConsumedQuota), logger.FormatQuota(userQuota-preConsumedQuota)))
|
|
}
|
|
relayInfo.FinalPreConsumedQuota = preConsumedQuota
|
|
return preConsumedQuota, nil
|
|
}
|