e2037ad756
This commit introduces a major architectural refactoring to improve quota management, centralize logging, and streamline the relay handling logic. Key changes: - **Pre-consume Quota:** Implements a new mechanism to check and reserve user quota *before* making the request to the upstream provider. This ensures more accurate quota deduction and prevents users from exceeding their limits due to concurrent requests. - **Unified Relay Handlers:** Refactors the relay logic to use generic handlers (e.g., `ChatHandler`, `ImageHandler`) instead of provider-specific implementations. This significantly reduces code duplication and simplifies adding new channels. - **Centralized Logger:** A new dedicated `logger` package is introduced, and all system logging calls are migrated to use it, moving this responsibility out of the `common` package. - **Code Reorganization:** DTOs are generalized (e.g., `dalle.go` -> `openai_image.go`) and utility code is moved to more appropriate packages (e.g., `common/http.go` -> `service/http.go`) for better code structure.
73 lines
3.0 KiB
Go
73 lines
3.0 KiB
Go
package service
|
|
|
|
import (
|
|
"errors"
|
|
"fmt"
|
|
"github.com/bytedance/gopkg/util/gopool"
|
|
"github.com/gin-gonic/gin"
|
|
"net/http"
|
|
"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 {
|
|
logger.SysError("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("pre-consume quota failed, user quota: %s, need quota: %s", logger.FormatQuota(userQuota), logger.FormatQuota(preConsumedQuota)), types.ErrorCodeInsufficientUserQuota, http.StatusForbidden, types.ErrOptionWithSkipRetry(), types.ErrOptionWithNoRecordErrorLog())
|
|
}
|
|
relayInfo.UserQuota = userQuota
|
|
if userQuota > 100*preConsumedQuota {
|
|
// 用户额度充足，判断令牌额度是否充足
|
|
if !relayInfo.TokenUnlimited {
|
|
// 非无限令牌，判断令牌额度是否充足
|
|
tokenQuota := c.GetInt("token_quota")
|
|
if tokenQuota > 100*preConsumedQuota {
|
|
// 令牌额度充足，信任令牌
|
|
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())
|
|
}
|
|
}
|
|
relayInfo.FinalPreConsumedQuota = preConsumedQuota
|
|
return preConsumedQuota, nil
|
|
}
|