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.
 
 
 

659 lines
22 KiB

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