Não pode escolher mais do que 25 tópicos Os tópicos devem começar com uma letra ou um número, podem incluir traços ('-') e podem ter até 35 caracteres.
 
 
 

834 linhas
28 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/setting/ratio_setting"
  26. "github.com/QuantumNous/new-api/types"
  27. "github.com/bytedance/gopkg/util/gopool"
  28. "github.com/gin-gonic/gin"
  29. "github.com/gorilla/websocket"
  30. )
  31. func relayHandler(c *gin.Context, info *relaycommon.RelayInfo) *types.NewAPIError {
  32. var err *types.NewAPIError
  33. switch info.RelayMode {
  34. case relayconstant.RelayModeImagesGenerations, relayconstant.RelayModeImagesEdits:
  35. err = relay.ImageHelper(c, info)
  36. case relayconstant.RelayModeAudioSpeech:
  37. fallthrough
  38. case relayconstant.RelayModeAudioTranslation:
  39. fallthrough
  40. case relayconstant.RelayModeAudioTranscription:
  41. err = relay.AudioHelper(c, info)
  42. case relayconstant.RelayModeRerank:
  43. err = relay.RerankHelper(c, info)
  44. case relayconstant.RelayModeEmbeddings:
  45. err = relay.EmbeddingHelper(c, info)
  46. case relayconstant.RelayModeResponses, relayconstant.RelayModeResponsesCompact:
  47. err = relay.ResponsesHelper(c, info)
  48. default:
  49. err = relay.TextHelper(c, info)
  50. }
  51. return err
  52. }
  53. func geminiRelayHandler(c *gin.Context, info *relaycommon.RelayInfo) *types.NewAPIError {
  54. var err *types.NewAPIError
  55. if strings.Contains(c.Request.URL.Path, "embed") {
  56. err = relay.GeminiEmbeddingHandler(c, info)
  57. } else {
  58. err = relay.GeminiHelper(c, info)
  59. }
  60. return err
  61. }
  62. func Relay(c *gin.Context, relayFormat types.RelayFormat) {
  63. requestId := c.GetString(common.RequestIdKey)
  64. //group := common.GetContextKeyString(c, constant.ContextKeyUsingGroup)
  65. //originalModel := common.GetContextKeyString(c, constant.ContextKeyOriginalModel)
  66. var (
  67. newAPIError *types.NewAPIError
  68. ws *websocket.Conn
  69. )
  70. if relayFormat == types.RelayFormatOpenAIRealtime {
  71. var err error
  72. ws, err = upgrader.Upgrade(c.Writer, c.Request, nil)
  73. if err != nil {
  74. helper.WssError(c, ws, types.NewError(err, types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry()).ToOpenAIError())
  75. return
  76. }
  77. defer ws.Close()
  78. }
  79. defer func() {
  80. if newAPIError != nil {
  81. logger.LogError(c, fmt.Sprintf("relay error: %s", newAPIError.Error()))
  82. // Prometheus: 错误计数
  83. if metrics.IsEnabled() {
  84. channelId := strconv.Itoa(c.GetInt("channel_id"))
  85. modelName := c.GetString("original_model")
  86. errorType, errorCode := metrics.ClassifyError(newAPIError.StatusCode, newAPIError.Error())
  87. metrics.GetMetrics().RequestErrorsTotal.WithLabelValues(channelId, modelName, errorType, errorCode).Inc()
  88. }
  89. newAPIError.SetMessage(common.MessageWithRequestId(newAPIError.Error(), requestId))
  90. switch relayFormat {
  91. case types.RelayFormatOpenAIRealtime:
  92. helper.WssError(c, ws, newAPIError.ToOpenAIError())
  93. case types.RelayFormatClaude:
  94. c.JSON(newAPIError.StatusCode, gin.H{
  95. "type": "error",
  96. "error": newAPIError.ToClaudeError(),
  97. })
  98. default:
  99. c.JSON(newAPIError.StatusCode, gin.H{
  100. "error": newAPIError.ToOpenAIError(),
  101. })
  102. }
  103. }
  104. }()
  105. request, err := helper.GetAndValidateRequest(c, relayFormat)
  106. if err != nil {
  107. // Map "request body too large" to 413 so clients can handle it correctly
  108. if common.IsRequestBodyTooLargeError(err) || errors.Is(err, common.ErrRequestBodyTooLarge) {
  109. newAPIError = types.NewErrorWithStatusCode(err, types.ErrorCodeReadRequestBodyFailed, http.StatusRequestEntityTooLarge, types.ErrOptionWithSkipRetry())
  110. } else {
  111. newAPIError = types.NewError(err, types.ErrorCodeInvalidRequest)
  112. }
  113. return
  114. }
  115. relayInfo, err := relaycommon.GenRelayInfo(c, relayFormat, request, ws)
  116. if err != nil {
  117. newAPIError = types.NewError(err, types.ErrorCodeGenRelayInfoFailed)
  118. return
  119. }
  120. needSensitiveCheck := setting.ShouldCheckPromptSensitive()
  121. needCountToken := constant.CountToken
  122. // Avoid building huge CombineText (strings.Join) when token counting and sensitive check are both disabled.
  123. var meta *types.TokenCountMeta
  124. if needSensitiveCheck || needCountToken {
  125. meta = request.GetTokenCountMeta()
  126. } else {
  127. meta = fastTokenCountMetaForPricing(request)
  128. }
  129. if needSensitiveCheck && meta != nil {
  130. contains, words := service.CheckSensitiveText(meta.CombineText)
  131. if contains {
  132. logger.LogWarn(c, fmt.Sprintf("user sensitive words detected: %s", strings.Join(words, ", ")))
  133. newAPIError = types.NewError(err, types.ErrorCodeSensitiveWordsDetected)
  134. return
  135. }
  136. }
  137. tokens, err := service.EstimateRequestToken(c, meta, relayInfo)
  138. if err != nil {
  139. newAPIError = types.NewError(err, types.ErrorCodeCountTokenFailed)
  140. return
  141. }
  142. relayInfo.SetEstimatePromptTokens(tokens)
  143. priceData, err := helper.ModelPriceHelper(c, relayInfo, tokens, meta)
  144. if err != nil {
  145. newAPIError = types.NewError(err, types.ErrorCodeModelPriceError)
  146. return
  147. }
  148. // common.SetContextKey(c, constant.ContextKeyTokenCountMeta, meta)
  149. if priceData.FreeModel {
  150. logger.LogInfo(c, fmt.Sprintf("模型 %s 免费,跳过预扣费", relayInfo.OriginModelName))
  151. } else {
  152. common.SysLog(fmt.Sprintf("[RegionSync] relay PreConsumeBilling: userId=%d, QuotaToPreConsume=%d, FreeModel=%v, UsePrice=%v, ModelRatio=%.4f, GroupRatio=%.4f, ModelPrice=%.4f, model=%s",
  153. relayInfo.UserId, priceData.QuotaToPreConsume, priceData.FreeModel, priceData.UsePrice, priceData.ModelRatio, priceData.GroupRatioInfo.GroupRatio, priceData.ModelPrice, relayInfo.OriginModelName))
  154. newAPIError = service.PreConsumeBilling(c, priceData.QuotaToPreConsume, relayInfo)
  155. if newAPIError != nil {
  156. return
  157. }
  158. }
  159. defer func() {
  160. // Only return quota if downstream failed and quota was actually pre-consumed
  161. if newAPIError != nil {
  162. newAPIError = service.NormalizeViolationFeeError(newAPIError)
  163. if relayInfo.Billing != nil {
  164. relayInfo.Billing.Refund(c)
  165. }
  166. service.ChargeViolationFeeIfNeeded(c, relayInfo, newAPIError)
  167. }
  168. }()
  169. retryParam := &service.RetryParam{
  170. Ctx: c,
  171. TokenGroup: relayInfo.TokenGroup,
  172. ModelName: relayInfo.OriginModelName,
  173. Retry: common.GetPointer(0),
  174. RequireMatrixUsageBilling: relayInfo.RequireMatrixUsageBilling,
  175. }
  176. for ; retryParam.GetRetry() <= common.RetryTimes; retryParam.IncreaseRetry() {
  177. channel, _, channelErr := getChannel(c, relayInfo, retryParam)
  178. if channelErr != nil {
  179. logger.LogError(c, channelErr.Error())
  180. newAPIError = channelErr
  181. break
  182. }
  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 concreteTaskVideoBindingGroup(c *gin.Context, group string) string {
  263. group = strings.TrimSpace(group)
  264. if group == "" {
  265. group = strings.TrimSpace(common.GetContextKeyString(c, constant.ContextKeyTokenGroup))
  266. }
  267. if group == "auto" {
  268. autoGroup := strings.TrimSpace(common.GetContextKeyString(c, constant.ContextKeyAutoGroup))
  269. if autoGroup == "" || autoGroup == "auto" {
  270. return ""
  271. }
  272. return autoGroup
  273. }
  274. return group
  275. }
  276. func persistTaskVideoBindingIfNeeded(userId int, group string, channel *model.Channel) error {
  277. group = strings.TrimSpace(group)
  278. if group == "" || group == "auto" || channel == nil {
  279. return nil
  280. }
  281. family, ok := service.VideoAssetFamilyForChannelType(channel.Type)
  282. if !ok {
  283. return nil
  284. }
  285. if !service.IsUsableVideoAssetChannelForFamily(channel, group, "", family) {
  286. return nil
  287. }
  288. return service.BindVideoAssetChannel(userId, group, channel, family)
  289. }
  290. func getChannel(c *gin.Context, info *relaycommon.RelayInfo, retryParam *service.RetryParam) (*model.Channel, string, *types.NewAPIError) {
  291. if info.ChannelMeta == nil {
  292. autoBan := c.GetBool("auto_ban")
  293. autoBanInt := 1
  294. if !autoBan {
  295. autoBanInt = 0
  296. }
  297. return &model.Channel{
  298. Id: c.GetInt("channel_id"),
  299. Type: c.GetInt("channel_type"),
  300. Name: c.GetString("channel_name"),
  301. AutoBan: &autoBanInt,
  302. }, concreteTaskVideoBindingGroup(c, info.TokenGroup), nil
  303. }
  304. channel, selectGroup, err := service.CacheGetRandomSatisfiedChannel(retryParam)
  305. info.PriceData.GroupRatioInfo = helper.HandleGroupRatio(c, info)
  306. if err != nil {
  307. if apiErr, ok := err.(*types.NewAPIError); ok {
  308. return nil, selectGroup, apiErr
  309. }
  310. return nil, selectGroup, types.NewError(fmt.Errorf("获取分组 %s 下模型 %s 的可用渠道失败(retry): %s", selectGroup, info.OriginModelName, err.Error()), types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry())
  311. }
  312. if channel == nil {
  313. return nil, selectGroup, types.NewError(fmt.Errorf("分组 %s 下模型 %s 的可用渠道不存在(retry)", selectGroup, info.OriginModelName), types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry())
  314. }
  315. newAPIError := middleware.SetupContextForSelectedChannel(c, channel, info.OriginModelName)
  316. if newAPIError != nil {
  317. return nil, selectGroup, newAPIError
  318. }
  319. return channel, selectGroup, nil
  320. }
  321. func requiredTaskChannelTypeForRequest(c *gin.Context) int {
  322. return 0
  323. }
  324. func allowedTaskChannelTypesForRequest(c *gin.Context) []int {
  325. if strings.HasPrefix(c.Request.URL.Path, "/api/v3/contents/generations/tasks") {
  326. // Keep in sync with allowedChannelTypesForRequest in middleware/distributor.go
  327. return service.VideoAssetChannelTypesForFamily(service.VideoAssetFamilySeedance)
  328. }
  329. if isKlingAipingNativePath(c.Request.URL.Path) {
  330. return service.VideoAssetChannelTypesForFamily(service.VideoAssetFamilyKling)
  331. }
  332. return nil
  333. }
  334. func shouldRetry(c *gin.Context, openaiErr *types.NewAPIError, retryTimes int) bool {
  335. if openaiErr == nil {
  336. return false
  337. }
  338. if service.ShouldSkipRetryAfterChannelAffinityFailure(c) {
  339. return false
  340. }
  341. if types.IsChannelError(openaiErr) {
  342. return true
  343. }
  344. if types.IsSkipRetryError(openaiErr) {
  345. return false
  346. }
  347. if retryTimes <= 0 {
  348. return false
  349. }
  350. code := openaiErr.StatusCode
  351. if code >= 200 && code < 300 {
  352. return false
  353. }
  354. if code < 100 || code > 599 {
  355. return true
  356. }
  357. return operation_setting.ShouldRetryByStatusCode(code)
  358. }
  359. func processChannelError(c *gin.Context, channelError types.ChannelError, err *types.NewAPIError) {
  360. logger.LogError(c, fmt.Sprintf("channel error (channel #%d, status code: %d): %s", channelError.ChannelId, err.StatusCode, err.Error()))
  361. // 不要使用context获取渠道信息,异步处理时可能会出现渠道信息不一致的情况
  362. // do not use context to get channel info, there may be inconsistent channel info when processing asynchronously
  363. if service.ShouldDisableChannel(channelError.ChannelType, err) && channelError.AutoBan {
  364. gopool.Go(func() {
  365. service.DisableChannel(channelError, err.ErrorWithStatusCode())
  366. })
  367. }
  368. if constant.ErrorLogEnabled && types.IsRecordErrorLog(err) {
  369. // 保存错误日志到mysql中
  370. userId := c.GetInt("id")
  371. tokenName := c.GetString("token_name")
  372. modelName := c.GetString("original_model")
  373. tokenId := c.GetInt("token_id")
  374. userGroup := c.GetString("group")
  375. channelId := c.GetInt("channel_id")
  376. other := make(map[string]interface{})
  377. if c.Request != nil && c.Request.URL != nil {
  378. other["request_path"] = c.Request.URL.Path
  379. }
  380. other["error_type"] = err.GetErrorType()
  381. other["error_code"] = err.GetErrorCode()
  382. other["status_code"] = err.StatusCode
  383. other["channel_id"] = channelId
  384. other["channel_name"] = c.GetString("channel_name")
  385. other["channel_type"] = c.GetInt("channel_type")
  386. if err.UpstreamRequestId != "" {
  387. other["upstream_request_id"] = err.UpstreamRequestId
  388. }
  389. if err.UpstreamBody != "" {
  390. other["upstream_body"] = err.UpstreamBody
  391. }
  392. adminInfo := make(map[string]interface{})
  393. adminInfo["use_channel"] = c.GetStringSlice("use_channel")
  394. isMultiKey := common.GetContextKeyBool(c, constant.ContextKeyChannelIsMultiKey)
  395. if isMultiKey {
  396. adminInfo["is_multi_key"] = true
  397. adminInfo["multi_key_index"] = common.GetContextKeyInt(c, constant.ContextKeyChannelMultiKeyIndex)
  398. }
  399. service.AppendChannelAffinityAdminInfo(c, adminInfo)
  400. other["admin_info"] = adminInfo
  401. startTime := common.GetContextKeyTime(c, constant.ContextKeyRequestStartTime)
  402. if startTime.IsZero() {
  403. startTime = time.Now()
  404. }
  405. useTimeSeconds := int(time.Since(startTime).Seconds())
  406. model.RecordErrorLog(c, userId, channelId, modelName, tokenName, err.MaskSensitiveErrorWithStatusCode(), tokenId, useTimeSeconds, false, userGroup, other)
  407. }
  408. }
  409. func RelayMidjourney(c *gin.Context) {
  410. relayInfo, err := relaycommon.GenRelayInfo(c, types.RelayFormatMjProxy, nil, nil)
  411. if err != nil {
  412. c.JSON(http.StatusInternalServerError, gin.H{
  413. "description": fmt.Sprintf("failed to generate relay info: %s", err.Error()),
  414. "type": "upstream_error",
  415. "code": 4,
  416. })
  417. return
  418. }
  419. var mjErr *dto.MidjourneyResponse
  420. switch relayInfo.RelayMode {
  421. case relayconstant.RelayModeMidjourneyNotify:
  422. mjErr = relay.RelayMidjourneyNotify(c)
  423. case relayconstant.RelayModeMidjourneyTaskFetch, relayconstant.RelayModeMidjourneyTaskFetchByCondition:
  424. mjErr = relay.RelayMidjourneyTask(c, relayInfo.RelayMode)
  425. case relayconstant.RelayModeMidjourneyTaskImageSeed:
  426. mjErr = relay.RelayMidjourneyTaskImageSeed(c)
  427. case relayconstant.RelayModeSwapFace:
  428. mjErr = relay.RelaySwapFace(c, relayInfo)
  429. default:
  430. mjErr = relay.RelayMidjourneySubmit(c, relayInfo)
  431. }
  432. //err = relayMidjourneySubmit(c, relayMode)
  433. log.Println(mjErr)
  434. if mjErr != nil {
  435. statusCode := http.StatusBadRequest
  436. if mjErr.Code == 30 {
  437. mjErr.Result = "当前分组负载已饱和,请稍后再试,或升级账户以提升服务质量。"
  438. statusCode = http.StatusTooManyRequests
  439. }
  440. c.JSON(statusCode, gin.H{
  441. "description": fmt.Sprintf("%s %s", mjErr.Description, mjErr.Result),
  442. "type": "upstream_error",
  443. "code": mjErr.Code,
  444. })
  445. channelId := c.GetInt("channel_id")
  446. logger.LogError(c, fmt.Sprintf("relay error (channel #%d, status code %d): %s", channelId, statusCode, fmt.Sprintf("%s %s", mjErr.Description, mjErr.Result)))
  447. }
  448. }
  449. func RelayNotImplemented(c *gin.Context) {
  450. err := types.OpenAIError{
  451. Message: "API not implemented",
  452. Type: "new_api_error",
  453. Param: "",
  454. Code: "api_not_implemented",
  455. }
  456. c.JSON(http.StatusNotImplemented, gin.H{
  457. "error": err,
  458. })
  459. }
  460. func RelayNotFound(c *gin.Context) {
  461. err := types.OpenAIError{
  462. Message: fmt.Sprintf("Invalid URL (%s %s)", c.Request.Method, c.Request.URL.Path),
  463. Type: "invalid_request_error",
  464. Param: "",
  465. Code: "",
  466. }
  467. c.JSON(http.StatusNotFound, gin.H{
  468. "error": err,
  469. })
  470. }
  471. func RelayTaskFetch(c *gin.Context) {
  472. relayInfo, err := relaycommon.GenRelayInfo(c, types.RelayFormatTask, nil, nil)
  473. if err != nil {
  474. c.JSON(http.StatusInternalServerError, &dto.TaskError{
  475. Code: "gen_relay_info_failed",
  476. Message: err.Error(),
  477. StatusCode: http.StatusInternalServerError,
  478. })
  479. return
  480. }
  481. if taskErr := relay.RelayTaskFetch(c, relayInfo.RelayMode); taskErr != nil {
  482. respondTaskError(c, taskErr)
  483. }
  484. }
  485. func preloadTaskPricingConfig(c *gin.Context, info *relaycommon.RelayInfo) *dto.TaskError {
  486. modelName := strings.TrimSpace(info.OriginModelName)
  487. action := strings.TrimSpace(info.Action)
  488. contentType := c.Request.Header.Get("Content-Type")
  489. storage, err := common.GetBodyStorage(c)
  490. if err != nil {
  491. status := http.StatusBadRequest
  492. if common.IsRequestBodyTooLargeError(err) || errors.Is(err, common.ErrRequestBodyTooLarge) {
  493. status = http.StatusRequestEntityTooLarge
  494. }
  495. return service.TaskErrorWrapperLocal(err, "read_request_body_failed", status)
  496. }
  497. defer func() {
  498. _, _ = storage.Seek(0, io.SeekStart)
  499. c.Request.Body = io.NopCloser(storage)
  500. }()
  501. switch {
  502. case strings.HasPrefix(contentType, "application/json"):
  503. body, err := storage.Bytes()
  504. if err != nil {
  505. return service.TaskErrorWrapperLocal(err, "read_request_body_failed", http.StatusBadRequest)
  506. }
  507. if strings.TrimSpace(string(body)) != "" {
  508. var payload map[string]any
  509. if err := common.Unmarshal(body, &payload); err != nil {
  510. return service.TaskErrorWrapperLocal(err, "invalid_json", http.StatusBadRequest)
  511. }
  512. if raw, ok := payload["model_name"].(string); ok && strings.TrimSpace(raw) != "" {
  513. modelName = strings.TrimSpace(raw)
  514. } else if raw, ok := payload["model"].(string); ok && strings.TrimSpace(raw) != "" {
  515. modelName = strings.TrimSpace(raw)
  516. }
  517. if raw, ok := payload["action"].(string); ok && strings.TrimSpace(raw) != "" {
  518. action = strings.TrimSpace(raw)
  519. }
  520. }
  521. case strings.Contains(contentType, gin.MIMEMultipartPOSTForm):
  522. form, err := common.ParseMultipartFormReusable(c)
  523. if err != nil {
  524. return service.TaskErrorWrapperLocal(err, "invalid_multipart_form", http.StatusBadRequest)
  525. }
  526. if vals := form.Value["model"]; len(vals) > 0 && strings.TrimSpace(vals[0]) != "" {
  527. modelName = strings.TrimSpace(vals[0])
  528. }
  529. if vals := form.Value["action"]; len(vals) > 0 && strings.TrimSpace(vals[0]) != "" {
  530. action = strings.TrimSpace(vals[0])
  531. }
  532. }
  533. if modelName == "" && action != "" {
  534. platform := constant.TaskPlatform(c.GetString("platform"))
  535. modelName = service.CoverTaskActionToModelName(platform, action)
  536. }
  537. if modelName != "" {
  538. info.OriginModelName = modelName
  539. }
  540. if action != "" {
  541. info.Action = action
  542. }
  543. info.PricingConfigSnapshotLoaded = true
  544. info.PricingConfigSnapshot = ratio_setting.GetPricingConfig(info.OriginModelName)
  545. info.RequireMatrixUsageBilling = info.PricingConfigSnapshot != nil &&
  546. info.PricingConfigSnapshot.BillingUnit == types.BillingUnitPer1MTokens
  547. return nil
  548. }
  549. func buildTaskBillingContext(info *relaycommon.RelayInfo) *model.TaskBillingContext {
  550. bc := &model.TaskBillingContext{
  551. ModelPrice: info.PriceData.ModelPrice,
  552. GroupRatio: info.PriceData.GroupRatioInfo.GroupRatio,
  553. ModelRatio: info.PriceData.ModelRatio,
  554. OtherRatios: info.PriceData.OtherRatios,
  555. OriginModelName: info.OriginModelName,
  556. PerCallBilling: common.StringsContains(constant.TaskPricePatches, info.OriginModelName),
  557. }
  558. if decision := info.PricingDecisionFrozen; decision != nil && decision.BillingMode == types.BillingModeMatrix {
  559. bc.BillingMode = decision.BillingMode
  560. bc.PricingSnapshot = types.CloneMapAny(decision.Snapshot)
  561. bc.BillingUnit = decision.BillingUnit
  562. bc.TokenUnitPriceUSD = decision.TokenUnitPriceUSD
  563. bc.PerCallBilling = decision.PerCallBilling
  564. if decision.BillingUnit == types.BillingUnitPer1MTokens {
  565. bc.ModelPrice = decision.TokenUnitPriceUSD
  566. } else {
  567. bc.ModelPrice = decision.PriceUSD
  568. }
  569. bc.GroupRatio = decision.GroupRatioInfo.GroupRatio
  570. bc.OtherRatios = types.CloneRatios(decision.OtherRatios)
  571. }
  572. return bc
  573. }
  574. func RelayTask(c *gin.Context) {
  575. relayInfo, err := relaycommon.GenRelayInfo(c, types.RelayFormatTask, nil, nil)
  576. if err != nil {
  577. c.JSON(http.StatusInternalServerError, &dto.TaskError{
  578. Code: "gen_relay_info_failed",
  579. Message: err.Error(),
  580. StatusCode: http.StatusInternalServerError,
  581. })
  582. return
  583. }
  584. relayTaskWithInfo(c, relayInfo)
  585. }
  586. func relayTaskWithInfo(c *gin.Context, relayInfo *relaycommon.RelayInfo) {
  587. if taskErr := relay.ResolveOriginTask(c, relayInfo); taskErr != nil {
  588. respondTaskError(c, taskErr)
  589. return
  590. }
  591. if taskErr := preloadTaskPricingConfig(c, relayInfo); taskErr != nil {
  592. respondTaskError(c, taskErr)
  593. return
  594. }
  595. var result *relay.TaskSubmitResult
  596. var taskErr *dto.TaskError
  597. defer func() {
  598. if taskErr != nil && relayInfo.Billing != nil {
  599. relayInfo.Billing.Refund(c)
  600. }
  601. }()
  602. retryParam := &service.RetryParam{
  603. Ctx: c,
  604. TokenGroup: relayInfo.TokenGroup,
  605. ModelName: relayInfo.OriginModelName,
  606. Retry: common.GetPointer(0),
  607. RequiredChannelType: requiredTaskChannelTypeForRequest(c),
  608. AllowedChannelTypes: allowedTaskChannelTypesForRequest(c),
  609. }
  610. for ; retryParam.GetRetry() <= common.RetryTimes; retryParam.IncreaseRetry() {
  611. var channel *model.Channel
  612. var selectedGroup string
  613. if lockedCh, ok := relayInfo.LockedChannel.(*model.Channel); ok && lockedCh != nil {
  614. channel = lockedCh
  615. selectedGroup = concreteTaskVideoBindingGroup(c, relayInfo.TokenGroup)
  616. if setupErr := middleware.SetupContextForSelectedChannel(c, channel, relayInfo.OriginModelName); setupErr != nil {
  617. taskErr = service.TaskErrorWrapperLocal(setupErr.Err, "setup_locked_channel_failed", http.StatusInternalServerError)
  618. break
  619. }
  620. } else {
  621. var channelErr *types.NewAPIError
  622. channel, selectedGroup, channelErr = getChannel(c, relayInfo, retryParam)
  623. if channelErr != nil {
  624. logger.LogError(c, channelErr.Error())
  625. taskErr = service.TaskErrorWrapperLocal(channelErr.Err, "get_channel_failed", http.StatusInternalServerError)
  626. break
  627. }
  628. }
  629. if bindErr := persistTaskVideoBindingIfNeeded(c.GetInt("id"), selectedGroup, channel); bindErr != nil {
  630. taskErr = service.TaskErrorWrapperLocal(bindErr, "bind_task_video_channel_failed", http.StatusServiceUnavailable)
  631. break
  632. }
  633. addUsedChannel(c, channel.Id)
  634. bodyStorage, bodyErr := common.GetBodyStorage(c)
  635. if bodyErr != nil {
  636. if common.IsRequestBodyTooLargeError(bodyErr) || errors.Is(bodyErr, common.ErrRequestBodyTooLarge) {
  637. taskErr = service.TaskErrorWrapperLocal(bodyErr, "read_request_body_failed", http.StatusRequestEntityTooLarge)
  638. } else {
  639. taskErr = service.TaskErrorWrapperLocal(bodyErr, "read_request_body_failed", http.StatusBadRequest)
  640. }
  641. break
  642. }
  643. c.Request.Body = io.NopCloser(bodyStorage)
  644. result, taskErr = relay.RelayTaskSubmit(c, relayInfo)
  645. if taskErr == nil {
  646. break
  647. }
  648. if !taskErr.LocalError {
  649. processChannelError(c,
  650. *types.NewChannelError(channel.Id, channel.Type, channel.Name, channel.ChannelInfo.IsMultiKey,
  651. common.GetContextKeyString(c, constant.ContextKeyChannelKey), channel.GetAutoBan()),
  652. types.NewOpenAIError(taskErr.Error, types.ErrorCodeBadResponseStatusCode, taskErr.StatusCode))
  653. }
  654. if !shouldRetryTaskRelay(c, channel.Id, taskErr, common.RetryTimes-retryParam.GetRetry()) {
  655. break
  656. }
  657. }
  658. useChannel := c.GetStringSlice("use_channel")
  659. if len(useChannel) > 1 {
  660. retryLogStr := fmt.Sprintf("重试:%s", strings.Trim(strings.Join(strings.Fields(fmt.Sprint(useChannel)), "->"), "[]"))
  661. logger.LogInfo(c, retryLogStr)
  662. }
  663. // ── 成功:结算 + 日志 + 插入任务 ──
  664. if taskErr == nil {
  665. if settleErr := service.SettleBilling(c, relayInfo, result.Quota); settleErr != nil {
  666. common.SysError("settle task billing error: " + settleErr.Error())
  667. }
  668. service.LogTaskConsumption(c, relayInfo)
  669. task := model.InitTask(result.Platform, relayInfo)
  670. task.PrivateData.UpstreamTaskID = result.UpstreamTaskID
  671. task.PrivateData.BillingSource = relayInfo.BillingSource
  672. task.PrivateData.SubscriptionId = relayInfo.SubscriptionId
  673. task.PrivateData.TokenId = relayInfo.TokenId
  674. task.PrivateData.BillingContext = buildTaskBillingContext(relayInfo)
  675. task.PrivateData.UpstreamRequest = relay.BuildUpstreamRequestSnapshotForTask(result.UpstreamReqJSON)
  676. task.Quota = result.Quota
  677. task.Data = result.TaskData
  678. task.Action = relayInfo.Action
  679. if insertErr := task.Insert(); insertErr != nil {
  680. common.SysError("insert task error: " + insertErr.Error())
  681. }
  682. }
  683. if taskErr != nil {
  684. respondTaskError(c, taskErr)
  685. }
  686. }
  687. // respondTaskError 统一输出 Task 错误响应(含 429 限流提示改写)
  688. func respondTaskError(c *gin.Context, taskErr *dto.TaskError) {
  689. if taskErr.StatusCode == http.StatusTooManyRequests {
  690. taskErr.Message = "当前分组上游负载已饱和,请稍后再试"
  691. }
  692. if isKlingAipingNativePath(c.Request.URL.Path) {
  693. normalizeKlingAipingTaskError(taskErr)
  694. }
  695. c.JSON(taskErr.StatusCode, taskErr)
  696. }
  697. func isKlingAipingNativePath(path string) bool {
  698. return pathMatchesAnyKlingAipingNativePrefix(path,
  699. "/v1/videos/text2video",
  700. "/v1/videos/image2video",
  701. "/v1/videos/motion-control",
  702. "/v1/videos/omni-video",
  703. "/v1/videos/multi-image2video",
  704. "/v1/videos/video-extend",
  705. "/v1/general/advanced-custom-elements",
  706. "/v1/general/custom-voices",
  707. )
  708. }
  709. func pathMatchesAnyKlingAipingNativePrefix(path string, prefixes ...string) bool {
  710. for _, prefix := range prefixes {
  711. if path == prefix || strings.HasPrefix(path, prefix+"/") {
  712. return true
  713. }
  714. }
  715. return false
  716. }
  717. func shouldRetryTaskRelay(c *gin.Context, channelId int, taskErr *dto.TaskError, retryTimes int) bool {
  718. if taskErr == nil {
  719. return false
  720. }
  721. if service.ShouldSkipRetryAfterChannelAffinityFailure(c) {
  722. return false
  723. }
  724. if retryTimes <= 0 {
  725. return false
  726. }
  727. if taskErr.StatusCode == http.StatusTooManyRequests {
  728. return true
  729. }
  730. if taskErr.StatusCode == 307 {
  731. return true
  732. }
  733. if taskErr.StatusCode/100 == 5 {
  734. // 超时不重试
  735. if operation_setting.IsAlwaysSkipRetryStatusCode(taskErr.StatusCode) {
  736. return false
  737. }
  738. return true
  739. }
  740. if taskErr.StatusCode == http.StatusBadRequest {
  741. return false
  742. }
  743. if taskErr.StatusCode == 408 {
  744. // azure处理超时不重试
  745. return false
  746. }
  747. if taskErr.LocalError {
  748. return false
  749. }
  750. if taskErr.StatusCode/100 == 2 {
  751. return false
  752. }
  753. return true
  754. }