You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

640 lines
20 KiB

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