relay.go 21 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644
  1. package controller
  2. import (
  3. "errors"
  4. "fmt"
  5. "io"
  6. "log"
  7. "net/http"
  8. "strings"
  9. "time"
  10. "github.com/QuantumNous/new-api/common"
  11. "github.com/QuantumNous/new-api/constant"
  12. "github.com/QuantumNous/new-api/dto"
  13. "github.com/QuantumNous/new-api/logger"
  14. "github.com/QuantumNous/new-api/middleware"
  15. "github.com/QuantumNous/new-api/model"
  16. "github.com/QuantumNous/new-api/relay"
  17. relaycommon "github.com/QuantumNous/new-api/relay/common"
  18. relayconstant "github.com/QuantumNous/new-api/relay/constant"
  19. "github.com/QuantumNous/new-api/relay/helper"
  20. "github.com/QuantumNous/new-api/service"
  21. "github.com/QuantumNous/new-api/setting"
  22. "github.com/QuantumNous/new-api/setting/operation_setting"
  23. "github.com/QuantumNous/new-api/types"
  24. "github.com/bytedance/gopkg/util/gopool"
  25. "github.com/samber/lo"
  26. "github.com/gin-gonic/gin"
  27. "github.com/gorilla/websocket"
  28. )
  29. func relayHandler(c *gin.Context, info *relaycommon.RelayInfo) *types.NewAPIError {
  30. var err *types.NewAPIError
  31. switch info.RelayMode {
  32. case relayconstant.RelayModeImagesGenerations, relayconstant.RelayModeImagesEdits:
  33. err = relay.ImageHelper(c, info)
  34. case relayconstant.RelayModeAudioSpeech:
  35. fallthrough
  36. case relayconstant.RelayModeAudioTranslation:
  37. fallthrough
  38. case relayconstant.RelayModeAudioTranscription:
  39. err = relay.AudioHelper(c, info)
  40. case relayconstant.RelayModeRerank:
  41. err = relay.RerankHelper(c, info)
  42. case relayconstant.RelayModeEmbeddings:
  43. err = relay.EmbeddingHelper(c, info)
  44. case relayconstant.RelayModeResponses, relayconstant.RelayModeResponsesCompact:
  45. err = relay.ResponsesHelper(c, info)
  46. default:
  47. err = relay.TextHelper(c, info)
  48. }
  49. return err
  50. }
  51. func geminiRelayHandler(c *gin.Context, info *relaycommon.RelayInfo) *types.NewAPIError {
  52. var err *types.NewAPIError
  53. if strings.Contains(c.Request.URL.Path, "embed") {
  54. err = relay.GeminiEmbeddingHandler(c, info)
  55. } else {
  56. err = relay.GeminiHelper(c, info)
  57. }
  58. return err
  59. }
  60. func Relay(c *gin.Context, relayFormat types.RelayFormat) {
  61. requestId := c.GetString(common.RequestIdKey)
  62. //group := common.GetContextKeyString(c, constant.ContextKeyUsingGroup)
  63. //originalModel := common.GetContextKeyString(c, constant.ContextKeyOriginalModel)
  64. var (
  65. newAPIError *types.NewAPIError
  66. ws *websocket.Conn
  67. )
  68. if relayFormat == types.RelayFormatOpenAIRealtime {
  69. var err error
  70. ws, err = upgrader.Upgrade(c.Writer, c.Request, nil)
  71. if err != nil {
  72. helper.WssError(c, ws, types.NewError(err, types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry()).ToOpenAIError())
  73. return
  74. }
  75. defer ws.Close()
  76. }
  77. defer func() {
  78. if newAPIError != nil {
  79. logger.LogError(c, fmt.Sprintf("relay error: %s", newAPIError.Error()))
  80. newAPIError.SetMessage(common.MessageWithRequestId(newAPIError.Error(), requestId))
  81. switch relayFormat {
  82. case types.RelayFormatOpenAIRealtime:
  83. helper.WssError(c, ws, newAPIError.ToOpenAIError())
  84. case types.RelayFormatClaude:
  85. c.JSON(newAPIError.StatusCode, gin.H{
  86. "type": "error",
  87. "error": newAPIError.ToClaudeError(),
  88. })
  89. default:
  90. c.JSON(newAPIError.StatusCode, gin.H{
  91. "error": newAPIError.ToOpenAIError(),
  92. })
  93. }
  94. }
  95. }()
  96. request, err := helper.GetAndValidateRequest(c, relayFormat)
  97. if err != nil {
  98. // Map "request body too large" to 413 so clients can handle it correctly
  99. if common.IsRequestBodyTooLargeError(err) || errors.Is(err, common.ErrRequestBodyTooLarge) {
  100. newAPIError = types.NewErrorWithStatusCode(err, types.ErrorCodeReadRequestBodyFailed, http.StatusRequestEntityTooLarge, types.ErrOptionWithSkipRetry())
  101. } else {
  102. newAPIError = types.NewError(err, types.ErrorCodeInvalidRequest)
  103. }
  104. return
  105. }
  106. relayInfo, err := relaycommon.GenRelayInfo(c, relayFormat, request, ws)
  107. if err != nil {
  108. newAPIError = types.NewError(err, types.ErrorCodeGenRelayInfoFailed)
  109. return
  110. }
  111. needSensitiveCheck := setting.ShouldCheckPromptSensitive()
  112. needCountToken := constant.CountToken
  113. // Avoid building huge CombineText (strings.Join) when token counting and sensitive check are both disabled.
  114. var meta *types.TokenCountMeta
  115. if needSensitiveCheck || needCountToken {
  116. meta = request.GetTokenCountMeta()
  117. } else {
  118. meta = fastTokenCountMetaForPricing(request)
  119. }
  120. if needSensitiveCheck && meta != nil {
  121. contains, words := service.CheckSensitiveText(meta.CombineText)
  122. if contains {
  123. logger.LogWarn(c, fmt.Sprintf("user sensitive words detected: %s", strings.Join(words, ", ")))
  124. newAPIError = types.NewError(err, types.ErrorCodeSensitiveWordsDetected)
  125. return
  126. }
  127. }
  128. tokens, err := service.EstimateRequestToken(c, meta, relayInfo)
  129. if err != nil {
  130. newAPIError = types.NewError(err, types.ErrorCodeCountTokenFailed)
  131. return
  132. }
  133. relayInfo.SetEstimatePromptTokens(tokens)
  134. priceData, err := helper.ModelPriceHelper(c, relayInfo, tokens, meta)
  135. if err != nil {
  136. newAPIError = types.NewError(err, types.ErrorCodeModelPriceError)
  137. return
  138. }
  139. // common.SetContextKey(c, constant.ContextKeyTokenCountMeta, meta)
  140. if priceData.FreeModel {
  141. logger.LogInfo(c, fmt.Sprintf("模型 %s 免费,跳过预扣费", relayInfo.OriginModelName))
  142. } else {
  143. newAPIError = service.PreConsumeBilling(c, priceData.QuotaToPreConsume, relayInfo)
  144. if newAPIError != nil {
  145. return
  146. }
  147. }
  148. defer func() {
  149. // Only return quota if downstream failed and quota was actually pre-consumed
  150. if newAPIError != nil {
  151. newAPIError = service.NormalizeViolationFeeError(newAPIError)
  152. if relayInfo.Billing != nil {
  153. relayInfo.Billing.Refund(c)
  154. }
  155. service.ChargeViolationFeeIfNeeded(c, relayInfo, newAPIError)
  156. }
  157. }()
  158. retryParam := &service.RetryParam{
  159. Ctx: c,
  160. TokenGroup: relayInfo.TokenGroup,
  161. ModelName: relayInfo.OriginModelName,
  162. Retry: common.GetPointer(0),
  163. }
  164. relayInfo.RetryIndex = 0
  165. relayInfo.LastError = nil
  166. for ; retryParam.GetRetry() <= common.RetryTimes; retryParam.IncreaseRetry() {
  167. relayInfo.RetryIndex = retryParam.GetRetry()
  168. channel, channelErr := getChannel(c, relayInfo, retryParam)
  169. if channelErr != nil {
  170. logger.LogError(c, channelErr.Error())
  171. newAPIError = channelErr
  172. break
  173. }
  174. addUsedChannel(c, channel.Id)
  175. bodyStorage, bodyErr := common.GetBodyStorage(c)
  176. if bodyErr != nil {
  177. // Ensure consistent 413 for oversized bodies even when error occurs later (e.g., retry path)
  178. if common.IsRequestBodyTooLargeError(bodyErr) || errors.Is(bodyErr, common.ErrRequestBodyTooLarge) {
  179. newAPIError = types.NewErrorWithStatusCode(bodyErr, types.ErrorCodeReadRequestBodyFailed, http.StatusRequestEntityTooLarge, types.ErrOptionWithSkipRetry())
  180. } else {
  181. newAPIError = types.NewErrorWithStatusCode(bodyErr, types.ErrorCodeReadRequestBodyFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry())
  182. }
  183. break
  184. }
  185. c.Request.Body = io.NopCloser(bodyStorage)
  186. switch relayFormat {
  187. case types.RelayFormatOpenAIRealtime:
  188. newAPIError = relay.WssHelper(c, relayInfo)
  189. case types.RelayFormatClaude:
  190. newAPIError = relay.ClaudeHelper(c, relayInfo)
  191. case types.RelayFormatGemini:
  192. newAPIError = geminiRelayHandler(c, relayInfo)
  193. default:
  194. newAPIError = relayHandler(c, relayInfo)
  195. }
  196. if newAPIError == nil {
  197. relayInfo.LastError = nil
  198. return
  199. }
  200. newAPIError = service.NormalizeViolationFeeError(newAPIError)
  201. relayInfo.LastError = newAPIError
  202. processChannelError(c, *types.NewChannelError(channel.Id, channel.Type, channel.Name, channel.ChannelInfo.IsMultiKey, common.GetContextKeyString(c, constant.ContextKeyChannelKey), channel.GetAutoBan()), newAPIError)
  203. if !shouldRetry(c, newAPIError, common.RetryTimes-retryParam.GetRetry()) {
  204. break
  205. }
  206. }
  207. useChannel := c.GetStringSlice("use_channel")
  208. if len(useChannel) > 1 {
  209. retryLogStr := fmt.Sprintf("重试:%s", strings.Trim(strings.Join(strings.Fields(fmt.Sprint(useChannel)), "->"), "[]"))
  210. logger.LogInfo(c, retryLogStr)
  211. }
  212. }
  213. var upgrader = websocket.Upgrader{
  214. Subprotocols: []string{"realtime"}, // WS 握手支持的协议,如果有使用 Sec-WebSocket-Protocol,则必须在此声明对应的 Protocol TODO add other protocol
  215. CheckOrigin: func(r *http.Request) bool {
  216. return true // 允许跨域
  217. },
  218. }
  219. func addUsedChannel(c *gin.Context, channelId int) {
  220. useChannel := c.GetStringSlice("use_channel")
  221. useChannel = append(useChannel, fmt.Sprintf("%d", channelId))
  222. c.Set("use_channel", useChannel)
  223. }
  224. func fastTokenCountMetaForPricing(request dto.Request) *types.TokenCountMeta {
  225. if request == nil {
  226. return &types.TokenCountMeta{}
  227. }
  228. meta := &types.TokenCountMeta{
  229. TokenType: types.TokenTypeTokenizer,
  230. }
  231. switch r := request.(type) {
  232. case *dto.GeneralOpenAIRequest:
  233. maxCompletionTokens := lo.FromPtrOr(r.MaxCompletionTokens, uint(0))
  234. maxTokens := lo.FromPtrOr(r.MaxTokens, uint(0))
  235. if maxCompletionTokens > maxTokens {
  236. meta.MaxTokens = int(maxCompletionTokens)
  237. } else {
  238. meta.MaxTokens = int(maxTokens)
  239. }
  240. case *dto.OpenAIResponsesRequest:
  241. meta.MaxTokens = int(lo.FromPtrOr(r.MaxOutputTokens, uint(0)))
  242. case *dto.ClaudeRequest:
  243. meta.MaxTokens = int(lo.FromPtr(r.MaxTokens))
  244. case *dto.ImageRequest:
  245. // Pricing for image requests depends on ImagePriceRatio; safe to compute even when CountToken is disabled.
  246. return r.GetTokenCountMeta()
  247. default:
  248. // Best-effort: leave CombineText empty to avoid large allocations.
  249. }
  250. return meta
  251. }
  252. func getChannel(c *gin.Context, info *relaycommon.RelayInfo, retryParam *service.RetryParam) (*model.Channel, *types.NewAPIError) {
  253. if info.ChannelMeta == nil {
  254. autoBan := c.GetBool("auto_ban")
  255. autoBanInt := 1
  256. if !autoBan {
  257. autoBanInt = 0
  258. }
  259. return &model.Channel{
  260. Id: c.GetInt("channel_id"),
  261. Type: c.GetInt("channel_type"),
  262. Name: c.GetString("channel_name"),
  263. AutoBan: &autoBanInt,
  264. }, nil
  265. }
  266. channel, selectGroup, err := service.CacheGetRandomSatisfiedChannel(retryParam)
  267. info.PriceData.GroupRatioInfo = helper.HandleGroupRatio(c, info)
  268. if err != nil {
  269. return nil, types.NewError(fmt.Errorf("获取分组 %s 下模型 %s 的可用渠道失败(retry): %s", selectGroup, info.OriginModelName, err.Error()), types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry())
  270. }
  271. if channel == nil {
  272. return nil, types.NewError(fmt.Errorf("分组 %s 下模型 %s 的可用渠道不存在(retry)", selectGroup, info.OriginModelName), types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry())
  273. }
  274. newAPIError := middleware.SetupContextForSelectedChannel(c, channel, info.OriginModelName)
  275. if newAPIError != nil {
  276. return nil, newAPIError
  277. }
  278. return channel, nil
  279. }
  280. func shouldRetry(c *gin.Context, openaiErr *types.NewAPIError, retryTimes int) bool {
  281. if openaiErr == nil {
  282. return false
  283. }
  284. if service.ShouldSkipRetryAfterChannelAffinityFailure(c) {
  285. return false
  286. }
  287. if types.IsChannelError(openaiErr) {
  288. return true
  289. }
  290. if types.IsSkipRetryError(openaiErr) {
  291. return false
  292. }
  293. if retryTimes <= 0 {
  294. return false
  295. }
  296. if _, ok := c.Get("specific_channel_id"); ok {
  297. return false
  298. }
  299. code := openaiErr.StatusCode
  300. if code >= 200 && code < 300 {
  301. return false
  302. }
  303. if code < 100 || code > 599 {
  304. return true
  305. }
  306. return operation_setting.ShouldRetryByStatusCode(code)
  307. }
  308. func processChannelError(c *gin.Context, channelError types.ChannelError, err *types.NewAPIError) {
  309. logger.LogError(c, fmt.Sprintf("channel error (channel #%d, status code: %d): %s", channelError.ChannelId, err.StatusCode, err.Error()))
  310. // 不要使用context获取渠道信息,异步处理时可能会出现渠道信息不一致的情况
  311. // do not use context to get channel info, there may be inconsistent channel info when processing asynchronously
  312. if service.ShouldDisableChannel(channelError.ChannelType, err) && channelError.AutoBan {
  313. gopool.Go(func() {
  314. service.DisableChannel(channelError, err.ErrorWithStatusCode())
  315. })
  316. }
  317. if constant.ErrorLogEnabled && types.IsRecordErrorLog(err) {
  318. // 保存错误日志到mysql中
  319. userId := c.GetInt("id")
  320. tokenName := c.GetString("token_name")
  321. modelName := c.GetString("original_model")
  322. tokenId := c.GetInt("token_id")
  323. userGroup := c.GetString("group")
  324. channelId := c.GetInt("channel_id")
  325. other := make(map[string]interface{})
  326. if c.Request != nil && c.Request.URL != nil {
  327. other["request_path"] = c.Request.URL.Path
  328. }
  329. other["error_type"] = err.GetErrorType()
  330. other["error_code"] = err.GetErrorCode()
  331. other["status_code"] = err.StatusCode
  332. other["channel_id"] = channelId
  333. other["channel_name"] = c.GetString("channel_name")
  334. other["channel_type"] = c.GetInt("channel_type")
  335. adminInfo := make(map[string]interface{})
  336. adminInfo["use_channel"] = c.GetStringSlice("use_channel")
  337. isMultiKey := common.GetContextKeyBool(c, constant.ContextKeyChannelIsMultiKey)
  338. if isMultiKey {
  339. adminInfo["is_multi_key"] = true
  340. adminInfo["multi_key_index"] = common.GetContextKeyInt(c, constant.ContextKeyChannelMultiKeyIndex)
  341. }
  342. service.AppendChannelAffinityAdminInfo(c, adminInfo)
  343. other["admin_info"] = adminInfo
  344. startTime := common.GetContextKeyTime(c, constant.ContextKeyRequestStartTime)
  345. if startTime.IsZero() {
  346. startTime = time.Now()
  347. }
  348. useTimeSeconds := int(time.Since(startTime).Seconds())
  349. model.RecordErrorLog(c, userId, channelId, modelName, tokenName, err.MaskSensitiveErrorWithStatusCode(), tokenId, useTimeSeconds, false, userGroup, other)
  350. }
  351. }
  352. func RelayMidjourney(c *gin.Context) {
  353. relayInfo, err := relaycommon.GenRelayInfo(c, types.RelayFormatMjProxy, nil, nil)
  354. if err != nil {
  355. c.JSON(http.StatusInternalServerError, gin.H{
  356. "description": fmt.Sprintf("failed to generate relay info: %s", err.Error()),
  357. "type": "upstream_error",
  358. "code": 4,
  359. })
  360. return
  361. }
  362. var mjErr *dto.MidjourneyResponse
  363. switch relayInfo.RelayMode {
  364. case relayconstant.RelayModeMidjourneyNotify:
  365. mjErr = relay.RelayMidjourneyNotify(c)
  366. case relayconstant.RelayModeMidjourneyTaskFetch, relayconstant.RelayModeMidjourneyTaskFetchByCondition:
  367. mjErr = relay.RelayMidjourneyTask(c, relayInfo.RelayMode)
  368. case relayconstant.RelayModeMidjourneyTaskImageSeed:
  369. mjErr = relay.RelayMidjourneyTaskImageSeed(c)
  370. case relayconstant.RelayModeSwapFace:
  371. mjErr = relay.RelaySwapFace(c, relayInfo)
  372. default:
  373. mjErr = relay.RelayMidjourneySubmit(c, relayInfo)
  374. }
  375. //err = relayMidjourneySubmit(c, relayMode)
  376. log.Println(mjErr)
  377. if mjErr != nil {
  378. statusCode := http.StatusBadRequest
  379. if mjErr.Code == 30 {
  380. mjErr.Result = "当前分组负载已饱和,请稍后再试,或升级账户以提升服务质量。"
  381. statusCode = http.StatusTooManyRequests
  382. }
  383. c.JSON(statusCode, gin.H{
  384. "description": fmt.Sprintf("%s %s", mjErr.Description, mjErr.Result),
  385. "type": "upstream_error",
  386. "code": mjErr.Code,
  387. })
  388. channelId := c.GetInt("channel_id")
  389. logger.LogError(c, fmt.Sprintf("relay error (channel #%d, status code %d): %s", channelId, statusCode, fmt.Sprintf("%s %s", mjErr.Description, mjErr.Result)))
  390. }
  391. }
  392. func RelayNotImplemented(c *gin.Context) {
  393. err := types.OpenAIError{
  394. Message: "API not implemented",
  395. Type: "new_api_error",
  396. Param: "",
  397. Code: "api_not_implemented",
  398. }
  399. c.JSON(http.StatusNotImplemented, gin.H{
  400. "error": err,
  401. })
  402. }
  403. func RelayNotFound(c *gin.Context) {
  404. err := types.OpenAIError{
  405. Message: fmt.Sprintf("Invalid URL (%s %s)", c.Request.Method, c.Request.URL.Path),
  406. Type: "invalid_request_error",
  407. Param: "",
  408. Code: "",
  409. }
  410. c.JSON(http.StatusNotFound, gin.H{
  411. "error": err,
  412. })
  413. }
  414. func RelayTaskFetch(c *gin.Context) {
  415. relayInfo, err := relaycommon.GenRelayInfo(c, types.RelayFormatTask, nil, nil)
  416. if err != nil {
  417. c.JSON(http.StatusInternalServerError, &dto.TaskError{
  418. Code: "gen_relay_info_failed",
  419. Message: err.Error(),
  420. StatusCode: http.StatusInternalServerError,
  421. })
  422. return
  423. }
  424. if taskErr := relay.RelayTaskFetch(c, relayInfo.RelayMode); taskErr != nil {
  425. respondTaskError(c, taskErr)
  426. }
  427. }
  428. func RelayTask(c *gin.Context) {
  429. relayInfo, err := relaycommon.GenRelayInfo(c, types.RelayFormatTask, nil, nil)
  430. if err != nil {
  431. c.JSON(http.StatusInternalServerError, &dto.TaskError{
  432. Code: "gen_relay_info_failed",
  433. Message: err.Error(),
  434. StatusCode: http.StatusInternalServerError,
  435. })
  436. return
  437. }
  438. if taskErr := relay.ResolveOriginTask(c, relayInfo); taskErr != nil {
  439. respondTaskError(c, taskErr)
  440. return
  441. }
  442. var result *relay.TaskSubmitResult
  443. var taskErr *dto.TaskError
  444. defer func() {
  445. if taskErr != nil && relayInfo.Billing != nil {
  446. relayInfo.Billing.Refund(c)
  447. }
  448. }()
  449. retryParam := &service.RetryParam{
  450. Ctx: c,
  451. TokenGroup: relayInfo.TokenGroup,
  452. ModelName: relayInfo.OriginModelName,
  453. Retry: common.GetPointer(0),
  454. }
  455. for ; retryParam.GetRetry() <= common.RetryTimes; retryParam.IncreaseRetry() {
  456. var channel *model.Channel
  457. if lockedCh, ok := relayInfo.LockedChannel.(*model.Channel); ok && lockedCh != nil {
  458. channel = lockedCh
  459. if retryParam.GetRetry() > 0 {
  460. if setupErr := middleware.SetupContextForSelectedChannel(c, channel, relayInfo.OriginModelName); setupErr != nil {
  461. taskErr = service.TaskErrorWrapperLocal(setupErr.Err, "setup_locked_channel_failed", http.StatusInternalServerError)
  462. break
  463. }
  464. }
  465. } else {
  466. var channelErr *types.NewAPIError
  467. channel, channelErr = getChannel(c, relayInfo, retryParam)
  468. if channelErr != nil {
  469. logger.LogError(c, channelErr.Error())
  470. taskErr = service.TaskErrorWrapperLocal(channelErr.Err, "get_channel_failed", http.StatusInternalServerError)
  471. break
  472. }
  473. }
  474. addUsedChannel(c, channel.Id)
  475. bodyStorage, bodyErr := common.GetBodyStorage(c)
  476. if bodyErr != nil {
  477. if common.IsRequestBodyTooLargeError(bodyErr) || errors.Is(bodyErr, common.ErrRequestBodyTooLarge) {
  478. taskErr = service.TaskErrorWrapperLocal(bodyErr, "read_request_body_failed", http.StatusRequestEntityTooLarge)
  479. } else {
  480. taskErr = service.TaskErrorWrapperLocal(bodyErr, "read_request_body_failed", http.StatusBadRequest)
  481. }
  482. break
  483. }
  484. c.Request.Body = io.NopCloser(bodyStorage)
  485. result, taskErr = relay.RelayTaskSubmit(c, relayInfo)
  486. if taskErr == nil {
  487. break
  488. }
  489. if !taskErr.LocalError {
  490. processChannelError(c,
  491. *types.NewChannelError(channel.Id, channel.Type, channel.Name, channel.ChannelInfo.IsMultiKey,
  492. common.GetContextKeyString(c, constant.ContextKeyChannelKey), channel.GetAutoBan()),
  493. types.NewOpenAIError(taskErr.Error, types.ErrorCodeBadResponseStatusCode, taskErr.StatusCode))
  494. }
  495. if !shouldRetryTaskRelay(c, channel.Id, taskErr, common.RetryTimes-retryParam.GetRetry()) {
  496. break
  497. }
  498. }
  499. useChannel := c.GetStringSlice("use_channel")
  500. if len(useChannel) > 1 {
  501. retryLogStr := fmt.Sprintf("重试:%s", strings.Trim(strings.Join(strings.Fields(fmt.Sprint(useChannel)), "->"), "[]"))
  502. logger.LogInfo(c, retryLogStr)
  503. }
  504. // ── 成功:结算 + 日志 + 插入任务 ──
  505. if taskErr == nil {
  506. if settleErr := service.SettleBilling(c, relayInfo, result.Quota); settleErr != nil {
  507. common.SysError("settle task billing error: " + settleErr.Error())
  508. }
  509. service.LogTaskConsumption(c, relayInfo)
  510. task := model.InitTask(result.Platform, relayInfo)
  511. task.PrivateData.UpstreamTaskID = result.UpstreamTaskID
  512. task.PrivateData.BillingSource = relayInfo.BillingSource
  513. task.PrivateData.SubscriptionId = relayInfo.SubscriptionId
  514. task.PrivateData.TokenId = relayInfo.TokenId
  515. task.PrivateData.BillingContext = &model.TaskBillingContext{
  516. ModelPrice: relayInfo.PriceData.ModelPrice,
  517. GroupRatio: relayInfo.PriceData.GroupRatioInfo.GroupRatio,
  518. ModelRatio: relayInfo.PriceData.ModelRatio,
  519. OtherRatios: relayInfo.PriceData.OtherRatios,
  520. OriginModelName: relayInfo.OriginModelName,
  521. PerCallBilling: common.StringsContains(constant.TaskPricePatches, relayInfo.OriginModelName),
  522. }
  523. task.Quota = result.Quota
  524. task.Data = result.TaskData
  525. task.Action = relayInfo.Action
  526. if insertErr := task.Insert(); insertErr != nil {
  527. common.SysError("insert task error: " + insertErr.Error())
  528. }
  529. }
  530. if taskErr != nil {
  531. respondTaskError(c, taskErr)
  532. }
  533. }
  534. // respondTaskError 统一输出 Task 错误响应(含 429 限流提示改写)
  535. func respondTaskError(c *gin.Context, taskErr *dto.TaskError) {
  536. if taskErr.StatusCode == http.StatusTooManyRequests {
  537. taskErr.Message = "当前分组上游负载已饱和,请稍后再试"
  538. }
  539. c.JSON(taskErr.StatusCode, taskErr)
  540. }
  541. func shouldRetryTaskRelay(c *gin.Context, channelId int, taskErr *dto.TaskError, retryTimes int) bool {
  542. if taskErr == nil {
  543. return false
  544. }
  545. if service.ShouldSkipRetryAfterChannelAffinityFailure(c) {
  546. return false
  547. }
  548. if retryTimes <= 0 {
  549. return false
  550. }
  551. if _, ok := c.Get("specific_channel_id"); ok {
  552. return false
  553. }
  554. if taskErr.StatusCode == http.StatusTooManyRequests {
  555. return true
  556. }
  557. if taskErr.StatusCode == 307 {
  558. return true
  559. }
  560. if taskErr.StatusCode/100 == 5 {
  561. // 超时不重试
  562. if operation_setting.IsAlwaysSkipRetryStatusCode(taskErr.StatusCode) {
  563. return false
  564. }
  565. return true
  566. }
  567. if taskErr.StatusCode == http.StatusBadRequest {
  568. return false
  569. }
  570. if taskErr.StatusCode == 408 {
  571. // azure处理超时不重试
  572. return false
  573. }
  574. if taskErr.LocalError {
  575. return false
  576. }
  577. if taskErr.StatusCode/100 == 2 {
  578. return false
  579. }
  580. return true
  581. }