8374a83084
Add three billing lifecycle methods to the TaskAdaptor interface: - EstimateBilling: compute OtherRatios from user request before pricing - AdjustBillingOnSubmit: adjust ratios from upstream submit response - AdjustBillingOnComplete: determine final quota at task terminal state Introduce BaseBilling as embeddable no-op default for adaptors without custom billing. Move Sora/Ali OtherRatios logic from shared validation into per-adaptor EstimateBilling implementations. Add TaskBillingContext to persist pricing params (model_price, group_ratio, other_ratios) in task private data for async polling settlement. Extract RecalculateTaskQuota as a general-purpose delta settlement function and unify polling billing via settleTaskBillingOnComplete (adaptor-first, then token-based fallback).
645 lines
22 KiB
Go
645 lines
22 KiB
Go
package controller
|
||
|
||
import (
|
||
"errors"
|
||
"fmt"
|
||
"io"
|
||
"log"
|
||
"net/http"
|
||
"strings"
|
||
"time"
|
||
|
||
"github.com/QuantumNous/new-api/common"
|
||
"github.com/QuantumNous/new-api/constant"
|
||
"github.com/QuantumNous/new-api/dto"
|
||
"github.com/QuantumNous/new-api/logger"
|
||
"github.com/QuantumNous/new-api/middleware"
|
||
"github.com/QuantumNous/new-api/model"
|
||
"github.com/QuantumNous/new-api/relay"
|
||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||
relayconstant "github.com/QuantumNous/new-api/relay/constant"
|
||
"github.com/QuantumNous/new-api/relay/helper"
|
||
"github.com/QuantumNous/new-api/service"
|
||
"github.com/QuantumNous/new-api/setting"
|
||
"github.com/QuantumNous/new-api/setting/operation_setting"
|
||
"github.com/QuantumNous/new-api/types"
|
||
|
||
"github.com/bytedance/gopkg/util/gopool"
|
||
|
||
"github.com/gin-gonic/gin"
|
||
"github.com/gorilla/websocket"
|
||
)
|
||
|
||
func relayHandler(c *gin.Context, info *relaycommon.RelayInfo) *types.NewAPIError {
|
||
var err *types.NewAPIError
|
||
switch info.RelayMode {
|
||
case relayconstant.RelayModeImagesGenerations, relayconstant.RelayModeImagesEdits:
|
||
err = relay.ImageHelper(c, info)
|
||
case relayconstant.RelayModeAudioSpeech:
|
||
fallthrough
|
||
case relayconstant.RelayModeAudioTranslation:
|
||
fallthrough
|
||
case relayconstant.RelayModeAudioTranscription:
|
||
err = relay.AudioHelper(c, info)
|
||
case relayconstant.RelayModeRerank:
|
||
err = relay.RerankHelper(c, info)
|
||
case relayconstant.RelayModeEmbeddings:
|
||
err = relay.EmbeddingHelper(c, info)
|
||
case relayconstant.RelayModeResponses, relayconstant.RelayModeResponsesCompact:
|
||
err = relay.ResponsesHelper(c, info)
|
||
default:
|
||
err = relay.TextHelper(c, info)
|
||
}
|
||
return err
|
||
}
|
||
|
||
func geminiRelayHandler(c *gin.Context, info *relaycommon.RelayInfo) *types.NewAPIError {
|
||
var err *types.NewAPIError
|
||
if strings.Contains(c.Request.URL.Path, "embed") {
|
||
err = relay.GeminiEmbeddingHandler(c, info)
|
||
} else {
|
||
err = relay.GeminiHelper(c, info)
|
||
}
|
||
return err
|
||
}
|
||
|
||
func Relay(c *gin.Context, relayFormat types.RelayFormat) {
|
||
|
||
requestId := c.GetString(common.RequestIdKey)
|
||
//group := common.GetContextKeyString(c, constant.ContextKeyUsingGroup)
|
||
//originalModel := common.GetContextKeyString(c, constant.ContextKeyOriginalModel)
|
||
|
||
var (
|
||
newAPIError *types.NewAPIError
|
||
ws *websocket.Conn
|
||
)
|
||
|
||
if relayFormat == types.RelayFormatOpenAIRealtime {
|
||
var err error
|
||
ws, err = upgrader.Upgrade(c.Writer, c.Request, nil)
|
||
if err != nil {
|
||
helper.WssError(c, ws, types.NewError(err, types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry()).ToOpenAIError())
|
||
return
|
||
}
|
||
defer ws.Close()
|
||
}
|
||
|
||
defer func() {
|
||
if newAPIError != nil {
|
||
logger.LogError(c, fmt.Sprintf("relay error: %s", newAPIError.Error()))
|
||
newAPIError.SetMessage(common.MessageWithRequestId(newAPIError.Error(), requestId))
|
||
switch relayFormat {
|
||
case types.RelayFormatOpenAIRealtime:
|
||
helper.WssError(c, ws, newAPIError.ToOpenAIError())
|
||
case types.RelayFormatClaude:
|
||
c.JSON(newAPIError.StatusCode, gin.H{
|
||
"type": "error",
|
||
"error": newAPIError.ToClaudeError(),
|
||
})
|
||
default:
|
||
c.JSON(newAPIError.StatusCode, gin.H{
|
||
"error": newAPIError.ToOpenAIError(),
|
||
})
|
||
}
|
||
}
|
||
}()
|
||
|
||
request, err := helper.GetAndValidateRequest(c, relayFormat)
|
||
if err != nil {
|
||
// Map "request body too large" to 413 so clients can handle it correctly
|
||
if common.IsRequestBodyTooLargeError(err) || errors.Is(err, common.ErrRequestBodyTooLarge) {
|
||
newAPIError = types.NewErrorWithStatusCode(err, types.ErrorCodeReadRequestBodyFailed, http.StatusRequestEntityTooLarge, types.ErrOptionWithSkipRetry())
|
||
} else {
|
||
newAPIError = types.NewError(err, types.ErrorCodeInvalidRequest)
|
||
}
|
||
return
|
||
}
|
||
|
||
relayInfo, err := relaycommon.GenRelayInfo(c, relayFormat, request, ws)
|
||
if err != nil {
|
||
newAPIError = types.NewError(err, types.ErrorCodeGenRelayInfoFailed)
|
||
return
|
||
}
|
||
|
||
needSensitiveCheck := setting.ShouldCheckPromptSensitive()
|
||
needCountToken := constant.CountToken
|
||
// Avoid building huge CombineText (strings.Join) when token counting and sensitive check are both disabled.
|
||
var meta *types.TokenCountMeta
|
||
if needSensitiveCheck || needCountToken {
|
||
meta = request.GetTokenCountMeta()
|
||
} else {
|
||
meta = fastTokenCountMetaForPricing(request)
|
||
}
|
||
|
||
if needSensitiveCheck && meta != nil {
|
||
contains, words := service.CheckSensitiveText(meta.CombineText)
|
||
if contains {
|
||
logger.LogWarn(c, fmt.Sprintf("user sensitive words detected: %s", strings.Join(words, ", ")))
|
||
newAPIError = types.NewError(err, types.ErrorCodeSensitiveWordsDetected)
|
||
return
|
||
}
|
||
}
|
||
|
||
tokens, err := service.EstimateRequestToken(c, meta, relayInfo)
|
||
if err != nil {
|
||
newAPIError = types.NewError(err, types.ErrorCodeCountTokenFailed)
|
||
return
|
||
}
|
||
|
||
relayInfo.SetEstimatePromptTokens(tokens)
|
||
|
||
priceData, err := helper.ModelPriceHelper(c, relayInfo, tokens, meta)
|
||
if err != nil {
|
||
newAPIError = types.NewError(err, types.ErrorCodeModelPriceError)
|
||
return
|
||
}
|
||
|
||
// common.SetContextKey(c, constant.ContextKeyTokenCountMeta, meta)
|
||
|
||
if priceData.FreeModel {
|
||
logger.LogInfo(c, fmt.Sprintf("模型 %s 免费，跳过预扣费", relayInfo.OriginModelName))
|
||
} else {
|
||
newAPIError = service.PreConsumeBilling(c, priceData.QuotaToPreConsume, relayInfo)
|
||
if newAPIError != nil {
|
||
return
|
||
}
|
||
}
|
||
|
||
defer func() {
|
||
// Only return quota if downstream failed and quota was actually pre-consumed
|
||
if newAPIError != nil {
|
||
newAPIError = service.NormalizeViolationFeeError(newAPIError)
|
||
if relayInfo.Billing != nil {
|
||
relayInfo.Billing.Refund(c)
|
||
}
|
||
service.ChargeViolationFeeIfNeeded(c, relayInfo, newAPIError)
|
||
}
|
||
}()
|
||
|
||
retryParam := &service.RetryParam{
|
||
Ctx: c,
|
||
TokenGroup: relayInfo.TokenGroup,
|
||
ModelName: relayInfo.OriginModelName,
|
||
Retry: common.GetPointer(0),
|
||
}
|
||
|
||
for ; retryParam.GetRetry() <= common.RetryTimes; retryParam.IncreaseRetry() {
|
||
channel, channelErr := getChannel(c, relayInfo, retryParam)
|
||
if channelErr != nil {
|
||
logger.LogError(c, channelErr.Error())
|
||
newAPIError = channelErr
|
||
break
|
||
}
|
||
|
||
addUsedChannel(c, channel.Id)
|
||
bodyStorage, bodyErr := common.GetBodyStorage(c)
|
||
if bodyErr != nil {
|
||
// Ensure consistent 413 for oversized bodies even when error occurs later (e.g., retry path)
|
||
if common.IsRequestBodyTooLargeError(bodyErr) || errors.Is(bodyErr, common.ErrRequestBodyTooLarge) {
|
||
newAPIError = types.NewErrorWithStatusCode(bodyErr, types.ErrorCodeReadRequestBodyFailed, http.StatusRequestEntityTooLarge, types.ErrOptionWithSkipRetry())
|
||
} else {
|
||
newAPIError = types.NewErrorWithStatusCode(bodyErr, types.ErrorCodeReadRequestBodyFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry())
|
||
}
|
||
break
|
||
}
|
||
c.Request.Body = io.NopCloser(bodyStorage)
|
||
|
||
switch relayFormat {
|
||
case types.RelayFormatOpenAIRealtime:
|
||
newAPIError = relay.WssHelper(c, relayInfo)
|
||
case types.RelayFormatClaude:
|
||
newAPIError = relay.ClaudeHelper(c, relayInfo)
|
||
case types.RelayFormatGemini:
|
||
newAPIError = geminiRelayHandler(c, relayInfo)
|
||
default:
|
||
newAPIError = relayHandler(c, relayInfo)
|
||
}
|
||
|
||
if newAPIError == nil {
|
||
return
|
||
}
|
||
|
||
newAPIError = service.NormalizeViolationFeeError(newAPIError)
|
||
|
||
processChannelError(c, *types.NewChannelError(channel.Id, channel.Type, channel.Name, channel.ChannelInfo.IsMultiKey, common.GetContextKeyString(c, constant.ContextKeyChannelKey), channel.GetAutoBan()), newAPIError)
|
||
|
||
if !shouldRetry(c, newAPIError, common.RetryTimes-retryParam.GetRetry()) {
|
||
break
|
||
}
|
||
}
|
||
|
||
useChannel := c.GetStringSlice("use_channel")
|
||
if len(useChannel) > 1 {
|
||
retryLogStr := fmt.Sprintf("重试：%s", strings.Trim(strings.Join(strings.Fields(fmt.Sprint(useChannel)), "->"), "[]"))
|
||
logger.LogInfo(c, retryLogStr)
|
||
}
|
||
}
|
||
|
||
var upgrader = websocket.Upgrader{
|
||
Subprotocols: []string{"realtime"}, // WS 握手支持的协议，如果有使用 Sec-WebSocket-Protocol，则必须在此声明对应的 Protocol TODO add other protocol
|
||
CheckOrigin: func(r *http.Request) bool {
|
||
return true // 允许跨域
|
||
},
|
||
}
|
||
|
||
func addUsedChannel(c *gin.Context, channelId int) {
|
||
useChannel := c.GetStringSlice("use_channel")
|
||
useChannel = append(useChannel, fmt.Sprintf("%d", channelId))
|
||
c.Set("use_channel", useChannel)
|
||
}
|
||
|
||
func fastTokenCountMetaForPricing(request dto.Request) *types.TokenCountMeta {
|
||
if request == nil {
|
||
return &types.TokenCountMeta{}
|
||
}
|
||
meta := &types.TokenCountMeta{
|
||
TokenType: types.TokenTypeTokenizer,
|
||
}
|
||
switch r := request.(type) {
|
||
case *dto.GeneralOpenAIRequest:
|
||
if r.MaxCompletionTokens > r.MaxTokens {
|
||
meta.MaxTokens = int(r.MaxCompletionTokens)
|
||
} else {
|
||
meta.MaxTokens = int(r.MaxTokens)
|
||
}
|
||
case *dto.OpenAIResponsesRequest:
|
||
meta.MaxTokens = int(r.MaxOutputTokens)
|
||
case *dto.ClaudeRequest:
|
||
meta.MaxTokens = int(r.MaxTokens)
|
||
case *dto.ImageRequest:
|
||
// Pricing for image requests depends on ImagePriceRatio; safe to compute even when CountToken is disabled.
|
||
return r.GetTokenCountMeta()
|
||
default:
|
||
// Best-effort: leave CombineText empty to avoid large allocations.
|
||
}
|
||
return meta
|
||
}
|
||
|
||
func getChannel(c *gin.Context, info *relaycommon.RelayInfo, retryParam *service.RetryParam) (*model.Channel, *types.NewAPIError) {
|
||
if info.ChannelMeta == nil {
|
||
autoBan := c.GetBool("auto_ban")
|
||
autoBanInt := 1
|
||
if !autoBan {
|
||
autoBanInt = 0
|
||
}
|
||
return &model.Channel{
|
||
Id: c.GetInt("channel_id"),
|
||
Type: c.GetInt("channel_type"),
|
||
Name: c.GetString("channel_name"),
|
||
AutoBan: &autoBanInt,
|
||
}, nil
|
||
}
|
||
channel, selectGroup, err := service.CacheGetRandomSatisfiedChannel(retryParam)
|
||
|
||
info.PriceData.GroupRatioInfo = helper.HandleGroupRatio(c, info)
|
||
|
||
if err != nil {
|
||
return nil, types.NewError(fmt.Errorf("获取分组 %s 下模型 %s 的可用渠道失败（retry）: %s", selectGroup, info.OriginModelName, err.Error()), types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry())
|
||
}
|
||
if channel == nil {
|
||
return nil, types.NewError(fmt.Errorf("分组 %s 下模型 %s 的可用渠道不存在（retry）", selectGroup, info.OriginModelName), types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry())
|
||
}
|
||
|
||
newAPIError := middleware.SetupContextForSelectedChannel(c, channel, info.OriginModelName)
|
||
if newAPIError != nil {
|
||
return nil, newAPIError
|
||
}
|
||
return channel, nil
|
||
}
|
||
|
||
func shouldRetry(c *gin.Context, openaiErr *types.NewAPIError, retryTimes int) bool {
|
||
if openaiErr == nil {
|
||
return false
|
||
}
|
||
if service.ShouldSkipRetryAfterChannelAffinityFailure(c) {
|
||
return false
|
||
}
|
||
if types.IsChannelError(openaiErr) {
|
||
return true
|
||
}
|
||
if types.IsSkipRetryError(openaiErr) {
|
||
return false
|
||
}
|
||
if retryTimes <= 0 {
|
||
return false
|
||
}
|
||
if _, ok := c.Get("specific_channel_id"); ok {
|
||
return false
|
||
}
|
||
code := openaiErr.StatusCode
|
||
if code >= 200 && code < 300 {
|
||
return false
|
||
}
|
||
if code < 100 || code > 599 {
|
||
return true
|
||
}
|
||
return operation_setting.ShouldRetryByStatusCode(code)
|
||
}
|
||
|
||
func processChannelError(c *gin.Context, channelError types.ChannelError, err *types.NewAPIError) {
|
||
logger.LogError(c, fmt.Sprintf("channel error (channel #%d, status code: %d): %s", channelError.ChannelId, err.StatusCode, err.Error()))
|
||
// 不要使用context获取渠道信息，异步处理时可能会出现渠道信息不一致的情况
|
||
// do not use context to get channel info, there may be inconsistent channel info when processing asynchronously
|
||
if service.ShouldDisableChannel(channelError.ChannelType, err) && channelError.AutoBan {
|
||
gopool.Go(func() {
|
||
service.DisableChannel(channelError, err.ErrorWithStatusCode())
|
||
})
|
||
}
|
||
|
||
if constant.ErrorLogEnabled && types.IsRecordErrorLog(err) {
|
||
// 保存错误日志到mysql中
|
||
userId := c.GetInt("id")
|
||
tokenName := c.GetString("token_name")
|
||
modelName := c.GetString("original_model")
|
||
tokenId := c.GetInt("token_id")
|
||
userGroup := c.GetString("group")
|
||
channelId := c.GetInt("channel_id")
|
||
other := make(map[string]interface{})
|
||
if c.Request != nil && c.Request.URL != nil {
|
||
other["request_path"] = c.Request.URL.Path
|
||
}
|
||
other["error_type"] = err.GetErrorType()
|
||
other["error_code"] = err.GetErrorCode()
|
||
other["status_code"] = err.StatusCode
|
||
other["channel_id"] = channelId
|
||
other["channel_name"] = c.GetString("channel_name")
|
||
other["channel_type"] = c.GetInt("channel_type")
|
||
adminInfo := make(map[string]interface{})
|
||
adminInfo["use_channel"] = c.GetStringSlice("use_channel")
|
||
isMultiKey := common.GetContextKeyBool(c, constant.ContextKeyChannelIsMultiKey)
|
||
if isMultiKey {
|
||
adminInfo["is_multi_key"] = true
|
||
adminInfo["multi_key_index"] = common.GetContextKeyInt(c, constant.ContextKeyChannelMultiKeyIndex)
|
||
}
|
||
service.AppendChannelAffinityAdminInfo(c, adminInfo)
|
||
other["admin_info"] = adminInfo
|
||
startTime := common.GetContextKeyTime(c, constant.ContextKeyRequestStartTime)
|
||
if startTime.IsZero() {
|
||
startTime = time.Now()
|
||
}
|
||
useTimeSeconds := int(time.Since(startTime).Seconds())
|
||
model.RecordErrorLog(c, userId, channelId, modelName, tokenName, err.MaskSensitiveErrorWithStatusCode(), tokenId, useTimeSeconds, false, userGroup, other)
|
||
}
|
||
|
||
}
|
||
|
||
func RelayMidjourney(c *gin.Context) {
|
||
relayInfo, err := relaycommon.GenRelayInfo(c, types.RelayFormatMjProxy, nil, nil)
|
||
|
||
if err != nil {
|
||
c.JSON(http.StatusInternalServerError, gin.H{
|
||
"description": fmt.Sprintf("failed to generate relay info: %s", err.Error()),
|
||
"type": "upstream_error",
|
||
"code": 4,
|
||
})
|
||
return
|
||
}
|
||
|
||
var mjErr *dto.MidjourneyResponse
|
||
switch relayInfo.RelayMode {
|
||
case relayconstant.RelayModeMidjourneyNotify:
|
||
mjErr = relay.RelayMidjourneyNotify(c)
|
||
case relayconstant.RelayModeMidjourneyTaskFetch, relayconstant.RelayModeMidjourneyTaskFetchByCondition:
|
||
mjErr = relay.RelayMidjourneyTask(c, relayInfo.RelayMode)
|
||
case relayconstant.RelayModeMidjourneyTaskImageSeed:
|
||
mjErr = relay.RelayMidjourneyTaskImageSeed(c)
|
||
case relayconstant.RelayModeSwapFace:
|
||
mjErr = relay.RelaySwapFace(c, relayInfo)
|
||
default:
|
||
mjErr = relay.RelayMidjourneySubmit(c, relayInfo)
|
||
}
|
||
//err = relayMidjourneySubmit(c, relayMode)
|
||
log.Println(mjErr)
|
||
if mjErr != nil {
|
||
statusCode := http.StatusBadRequest
|
||
if mjErr.Code == 30 {
|
||
mjErr.Result = "当前分组负载已饱和，请稍后再试，或升级账户以提升服务质量。"
|
||
statusCode = http.StatusTooManyRequests
|
||
}
|
||
c.JSON(statusCode, gin.H{
|
||
"description": fmt.Sprintf("%s %s", mjErr.Description, mjErr.Result),
|
||
"type": "upstream_error",
|
||
"code": mjErr.Code,
|
||
})
|
||
channelId := c.GetInt("channel_id")
|
||
logger.LogError(c, fmt.Sprintf("relay error (channel #%d, status code %d): %s", channelId, statusCode, fmt.Sprintf("%s %s", mjErr.Description, mjErr.Result)))
|
||
}
|
||
}
|
||
|
||
func RelayNotImplemented(c *gin.Context) {
|
||
err := types.OpenAIError{
|
||
Message: "API not implemented",
|
||
Type: "new_api_error",
|
||
Param: "",
|
||
Code: "api_not_implemented",
|
||
}
|
||
c.JSON(http.StatusNotImplemented, gin.H{
|
||
"error": err,
|
||
})
|
||
}
|
||
|
||
func RelayNotFound(c *gin.Context) {
|
||
err := types.OpenAIError{
|
||
Message: fmt.Sprintf("Invalid URL (%s %s)", c.Request.Method, c.Request.URL.Path),
|
||
Type: "invalid_request_error",
|
||
Param: "",
|
||
Code: "",
|
||
}
|
||
c.JSON(http.StatusNotFound, gin.H{
|
||
"error": err,
|
||
})
|
||
}
|
||
|
||
func RelayTask(c *gin.Context) {
|
||
channelId := c.GetInt("channel_id")
|
||
c.Set("use_channel", []string{fmt.Sprintf("%d", channelId)})
|
||
relayInfo, err := relaycommon.GenRelayInfo(c, types.RelayFormatTask, nil, nil)
|
||
if err != nil {
|
||
c.JSON(http.StatusInternalServerError, &dto.TaskError{
|
||
Code: "gen_relay_info_failed",
|
||
Message: err.Error(),
|
||
StatusCode: http.StatusInternalServerError,
|
||
})
|
||
return
|
||
}
|
||
|
||
// Fetch 操作是纯 DB 查询（或 task 自带 channelId 的上游查询），不依赖上下文 channel，无需重试
|
||
// TODO: 在video-route层面优化，避免无谓的 channel 选择和上下文设置，也没必要吧代码放到这里来写这么多屎山
|
||
switch relayInfo.RelayMode {
|
||
case relayconstant.RelayModeSunoFetch, relayconstant.RelayModeSunoFetchByID, relayconstant.RelayModeVideoFetchByID:
|
||
if taskErr := relay.RelayTaskFetch(c, relayInfo.RelayMode); taskErr != nil {
|
||
respondTaskError(c, taskErr)
|
||
}
|
||
return
|
||
}
|
||
|
||
// ── Submit 路径 ─────────────────────────────────────────────────
|
||
|
||
// 1. 解析原始任务（remix / continuation），一次性，可能锁定渠道并禁止重试
|
||
if taskErr := relay.ResolveOriginTask(c, relayInfo); taskErr != nil {
|
||
respondTaskError(c, taskErr)
|
||
return
|
||
}
|
||
|
||
// 2. defer Refund（全部失败时回滚预扣费）
|
||
var result *relay.TaskSubmitResult
|
||
var taskErr *dto.TaskError
|
||
defer func() {
|
||
if taskErr != nil && relayInfo.Billing != nil {
|
||
relayInfo.Billing.Refund(c)
|
||
}
|
||
}()
|
||
|
||
// 3. 执行 + 重试（RelayTaskSubmit 内部在首次调用时自动预扣费）
|
||
taskErr = taskSubmitWithRetry(c, relayInfo, channelId, common.RetryTimes, func() *dto.TaskError {
|
||
var te *dto.TaskError
|
||
result, te = relay.RelayTaskSubmit(c, relayInfo)
|
||
return te
|
||
})
|
||
|
||
// 4. 成功：结算 + 日志 + 插入任务
|
||
if taskErr == nil {
|
||
if settleErr := service.SettleBilling(c, relayInfo, result.Quota); settleErr != nil {
|
||
common.SysError("settle task billing error: " + settleErr.Error())
|
||
}
|
||
service.LogTaskConsumption(c, relayInfo, result.ModelName)
|
||
|
||
task := model.InitTask(result.Platform, relayInfo)
|
||
task.PrivateData.UpstreamTaskID = result.UpstreamTaskID
|
||
task.PrivateData.BillingSource = relayInfo.BillingSource
|
||
task.PrivateData.SubscriptionId = relayInfo.SubscriptionId
|
||
task.PrivateData.TokenId = relayInfo.TokenId
|
||
task.PrivateData.BillingContext = &model.TaskBillingContext{
|
||
ModelPrice: relayInfo.PriceData.ModelPrice,
|
||
GroupRatio: relayInfo.PriceData.GroupRatioInfo.GroupRatio,
|
||
ModelRatio: relayInfo.PriceData.ModelRatio,
|
||
OtherRatios: relayInfo.PriceData.OtherRatios,
|
||
ModelName: result.ModelName,
|
||
}
|
||
task.Quota = result.Quota
|
||
task.Data = result.TaskData
|
||
task.Action = relayInfo.Action
|
||
if insertErr := task.Insert(); insertErr != nil {
|
||
//taskErr = service.TaskErrorWrapper(insertErr, "insert_task_failed", http.StatusInternalServerError)
|
||
common.SysError("insert task error: " + insertErr.Error())
|
||
}
|
||
}
|
||
|
||
if taskErr != nil {
|
||
respondTaskError(c, taskErr)
|
||
}
|
||
}
|
||
|
||
// respondTaskError 统一输出 Task 错误响应（含 429 限流提示改写）
|
||
func respondTaskError(c *gin.Context, taskErr *dto.TaskError) {
|
||
if taskErr.StatusCode == http.StatusTooManyRequests {
|
||
taskErr.Message = "当前分组上游负载已饱和，请稍后再试"
|
||
}
|
||
c.JSON(taskErr.StatusCode, taskErr)
|
||
}
|
||
|
||
// taskSubmitWithRetry 执行首次尝试并在失败时切换渠道重试，返回最终的 taskErr。
|
||
// attempt 闭包负责实际的上游请求，不涉及计费。
|
||
func taskSubmitWithRetry(c *gin.Context, relayInfo *relaycommon.RelayInfo,
|
||
channelId int, retryTimes int, attempt func() *dto.TaskError) *dto.TaskError {
|
||
|
||
taskErr := attempt()
|
||
if taskErr == nil {
|
||
return nil
|
||
}
|
||
if !taskErr.LocalError {
|
||
processChannelError(c,
|
||
*types.NewChannelError(channelId, c.GetInt("channel_type"), c.GetString("channel_name"), common.GetContextKeyBool(c, constant.ContextKeyChannelIsMultiKey),
|
||
common.GetContextKeyString(c, constant.ContextKeyChannelKey), common.GetContextKeyBool(c, constant.ContextKeyChannelAutoBan)),
|
||
types.NewOpenAIError(taskErr.Error, types.ErrorCodeBadResponseStatusCode, taskErr.StatusCode))
|
||
}
|
||
|
||
retryParam := &service.RetryParam{
|
||
Ctx: c,
|
||
TokenGroup: relayInfo.TokenGroup,
|
||
ModelName: relayInfo.OriginModelName,
|
||
Retry: common.GetPointer(0),
|
||
}
|
||
for ; shouldRetryTaskRelay(c, channelId, taskErr, retryTimes) && retryParam.GetRetry() < retryTimes; retryParam.IncreaseRetry() {
|
||
channel, newAPIError := getChannel(c, relayInfo, retryParam)
|
||
if newAPIError != nil {
|
||
logger.LogError(c, fmt.Sprintf("CacheGetRandomSatisfiedChannel failed: %s", newAPIError.Error()))
|
||
taskErr = service.TaskErrorWrapperLocal(newAPIError.Err, "get_channel_failed", http.StatusInternalServerError)
|
||
break
|
||
}
|
||
channelId = channel.Id
|
||
useChannel := c.GetStringSlice("use_channel")
|
||
useChannel = append(useChannel, fmt.Sprintf("%d", channelId))
|
||
c.Set("use_channel", useChannel)
|
||
logger.LogInfo(c, fmt.Sprintf("using channel #%d to retry (remain times %d)", channel.Id, retryParam.GetRetry()))
|
||
middleware.SetupContextForSelectedChannel(c, channel, c.GetString("original_model"))
|
||
|
||
bodyStorage, err := common.GetBodyStorage(c)
|
||
if err != nil {
|
||
if common.IsRequestBodyTooLargeError(err) || errors.Is(err, common.ErrRequestBodyTooLarge) {
|
||
taskErr = service.TaskErrorWrapperLocal(err, "read_request_body_failed", http.StatusRequestEntityTooLarge)
|
||
} else {
|
||
taskErr = service.TaskErrorWrapperLocal(err, "read_request_body_failed", http.StatusBadRequest)
|
||
}
|
||
break
|
||
}
|
||
c.Request.Body = io.NopCloser(bodyStorage)
|
||
taskErr = attempt()
|
||
if taskErr != nil && !taskErr.LocalError {
|
||
processChannelError(c,
|
||
*types.NewChannelError(channel.Id, channel.Type, channel.Name, channel.ChannelInfo.IsMultiKey,
|
||
common.GetContextKeyString(c, constant.ContextKeyChannelKey), channel.GetAutoBan()),
|
||
types.NewOpenAIError(taskErr.Error, types.ErrorCodeBadResponseStatusCode, taskErr.StatusCode))
|
||
}
|
||
}
|
||
|
||
useChannel := c.GetStringSlice("use_channel")
|
||
if len(useChannel) > 1 {
|
||
retryLogStr := fmt.Sprintf("重试：%s", strings.Trim(strings.Join(strings.Fields(fmt.Sprint(useChannel)), "->"), "[]"))
|
||
logger.LogInfo(c, retryLogStr)
|
||
}
|
||
return taskErr
|
||
}
|
||
|
||
func shouldRetryTaskRelay(c *gin.Context, channelId int, taskErr *dto.TaskError, retryTimes int) bool {
|
||
if taskErr == nil {
|
||
return false
|
||
}
|
||
if service.ShouldSkipRetryAfterChannelAffinityFailure(c) {
|
||
return false
|
||
}
|
||
if retryTimes <= 0 {
|
||
return false
|
||
}
|
||
if _, ok := c.Get("specific_channel_id"); ok {
|
||
return false
|
||
}
|
||
if taskErr.StatusCode == http.StatusTooManyRequests {
|
||
return true
|
||
}
|
||
if taskErr.StatusCode == 307 {
|
||
return true
|
||
}
|
||
if taskErr.StatusCode/100 == 5 {
|
||
// 超时不重试
|
||
if taskErr.StatusCode == 504 || taskErr.StatusCode == 524 {
|
||
return false
|
||
}
|
||
return true
|
||
}
|
||
if taskErr.StatusCode == http.StatusBadRequest {
|
||
return false
|
||
}
|
||
if taskErr.StatusCode == 408 {
|
||
// azure处理超时不重试
|
||
return false
|
||
}
|
||
if taskErr.LocalError {
|
||
return false
|
||
}
|
||
if taskErr.StatusCode/100 == 2 {
|
||
return false
|
||
}
|
||
return true
|
||
}
|