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.
 
 
 

469 regels
19 KiB

  1. package middleware
  2. import (
  3. "errors"
  4. "fmt"
  5. "net/http"
  6. "slices"
  7. "strconv"
  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/i18n"
  14. "github.com/QuantumNous/new-api/model"
  15. relayconstant "github.com/QuantumNous/new-api/relay/constant"
  16. "github.com/QuantumNous/new-api/service"
  17. "github.com/QuantumNous/new-api/setting/ratio_setting"
  18. "github.com/QuantumNous/new-api/types"
  19. "github.com/gin-gonic/gin"
  20. )
  21. type ModelRequest struct {
  22. Model string `json:"model"`
  23. Group string `json:"group,omitempty"`
  24. ChannelId int `json:"channel_id,omitempty"`
  25. }
  26. func Distribute() func(c *gin.Context) {
  27. return func(c *gin.Context) {
  28. var channel *model.Channel
  29. modelRequest, shouldSelectChannel, err := getModelRequest(c)
  30. if err != nil {
  31. abortWithOpenAiMessage(c, http.StatusBadRequest, i18n.T(c, i18n.MsgDistributorInvalidRequest, map[string]any{"Error": err.Error()}))
  32. return
  33. }
  34. channelId, ok := common.GetContextKey(c, constant.ContextKeyTokenSpecificChannelId)
  35. if ok {
  36. id, err := strconv.Atoi(channelId.(string))
  37. if err != nil {
  38. abortWithOpenAiMessage(c, http.StatusBadRequest, i18n.T(c, i18n.MsgDistributorInvalidChannelId))
  39. return
  40. }
  41. channel, err = model.GetChannelById(id, true)
  42. if err != nil {
  43. abortWithOpenAiMessage(c, http.StatusBadRequest, i18n.T(c, i18n.MsgDistributorInvalidChannelId))
  44. return
  45. }
  46. if channel.Status != common.ChannelStatusEnabled {
  47. abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgDistributorChannelDisabled))
  48. return
  49. }
  50. // 验证指定渠道是否支持请求的模型
  51. if shouldSelectChannel && modelRequest.Model != "" {
  52. usingGroup := common.GetContextKeyString(c, constant.ContextKeyUsingGroup)
  53. if !model.IsChannelEnabledForGroupModel(usingGroup, modelRequest.Model, channel.Id) {
  54. abortWithOpenAiMessage(c, http.StatusServiceUnavailable, i18n.T(c, i18n.MsgDistributorNoAvailableChannel, map[string]any{"Group": usingGroup, "Model": modelRequest.Model}), types.ErrorCodeModelNotFound)
  55. return
  56. }
  57. }
  58. } else {
  59. // Normal channel selection logic
  60. // check token model mapping
  61. modelLimitEnable := common.GetContextKeyBool(c, constant.ContextKeyTokenModelLimitEnabled)
  62. if modelLimitEnable {
  63. s, ok := common.GetContextKey(c, constant.ContextKeyTokenModelLimit)
  64. if !ok {
  65. // token model limit is empty, all models are not allowed
  66. abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgDistributorTokenNoModelAccess))
  67. return
  68. }
  69. var tokenModelLimit map[string]bool
  70. tokenModelLimit, ok = s.(map[string]bool)
  71. if !ok {
  72. tokenModelLimit = map[string]bool{}
  73. }
  74. matchName := ratio_setting.FormatMatchingModelName(modelRequest.Model) // match gpts & thinking-*
  75. if _, ok := tokenModelLimit[matchName]; !ok {
  76. abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgDistributorTokenModelForbidden, map[string]any{"Model": modelRequest.Model}))
  77. return
  78. }
  79. }
  80. if shouldSelectChannel {
  81. if modelRequest.Model == "" {
  82. abortWithOpenAiMessage(c, http.StatusBadRequest, i18n.T(c, i18n.MsgDistributorModelNameRequired))
  83. return
  84. }
  85. var selectGroup string
  86. usingGroup := common.GetContextKeyString(c, constant.ContextKeyUsingGroup)
  87. // check path is /pg/chat/completions
  88. if strings.HasPrefix(c.Request.URL.Path, "/pg/chat/completions") {
  89. playgroundRequest := &dto.PlayGroundRequest{}
  90. err = common.UnmarshalBodyReusable(c, playgroundRequest)
  91. if err != nil {
  92. abortWithOpenAiMessage(c, http.StatusBadRequest, i18n.T(c, i18n.MsgDistributorInvalidPlayground, map[string]any{"Error": err.Error()}))
  93. return
  94. }
  95. if playgroundRequest.Group != "" {
  96. if !service.GroupInUserUsableGroups(usingGroup, playgroundRequest.Group) && playgroundRequest.Group != usingGroup {
  97. abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgDistributorGroupAccessDenied))
  98. return
  99. }
  100. usingGroup = playgroundRequest.Group
  101. common.SetContextKey(c, constant.ContextKeyUsingGroup, usingGroup)
  102. }
  103. }
  104. // 默认通道检查(管理员配置的模型默认通道)
  105. if channel == nil {
  106. if defaultChannelId, ok := model.GetDefaultChannelId(modelRequest.Model); ok {
  107. defaultCh, err := model.CacheGetChannel(defaultChannelId)
  108. if err != nil || defaultCh == nil {
  109. common.SysLog(fmt.Sprintf("[Distribute] model=%s default_channel=%d not found in cache, fallback", modelRequest.Model, defaultChannelId))
  110. } else if defaultCh.Status != common.ChannelStatusEnabled {
  111. common.SysLog(fmt.Sprintf("[Distribute] model=%s default_channel=%d disabled(status=%d), fallback", modelRequest.Model, defaultChannelId, defaultCh.Status))
  112. } else if usingGroup == "auto" {
  113. userGroup := common.GetContextKeyString(c, constant.ContextKeyUserGroup)
  114. autoGroups := service.GetUserAutoGroup(userGroup)
  115. for _, g := range autoGroups {
  116. if model.IsChannelEnabledForGroupModel(g, modelRequest.Model, defaultCh.Id) {
  117. channel = defaultCh
  118. selectGroup = g
  119. common.SetContextKey(c, constant.ContextKeyAutoGroup, g)
  120. common.SysLog(fmt.Sprintf("[Distribute] model=%s using default_channel=%d (auto group=%s)", modelRequest.Model, defaultChannelId, g))
  121. break
  122. }
  123. }
  124. if channel == nil {
  125. common.SysLog(fmt.Sprintf("[Distribute] model=%s default_channel=%d not enabled for any auto group, fallback", modelRequest.Model, defaultChannelId))
  126. }
  127. } else if model.IsChannelEnabledForGroupModel(usingGroup, modelRequest.Model, defaultCh.Id) {
  128. channel = defaultCh
  129. selectGroup = usingGroup
  130. common.SysLog(fmt.Sprintf("[Distribute] model=%s using default_channel=%d (group=%s)", modelRequest.Model, defaultChannelId, usingGroup))
  131. } else {
  132. common.SysLog(fmt.Sprintf("[Distribute] model=%s default_channel=%d not enabled for group=%s, fallback", modelRequest.Model, defaultChannelId, usingGroup))
  133. }
  134. }
  135. }
  136. // 通道亲和性检查(仅在未选中默认通道时生效)
  137. if channel == nil {
  138. if preferredChannelID, found := service.GetPreferredChannelByAffinity(c, modelRequest.Model, usingGroup); found {
  139. preferred, err := model.CacheGetChannel(preferredChannelID)
  140. if err == nil && preferred != nil && preferred.Status == common.ChannelStatusEnabled {
  141. if usingGroup == "auto" {
  142. userGroup := common.GetContextKeyString(c, constant.ContextKeyUserGroup)
  143. autoGroups := service.GetUserAutoGroup(userGroup)
  144. for _, g := range autoGroups {
  145. if model.IsChannelEnabledForGroupModel(g, modelRequest.Model, preferred.Id) {
  146. selectGroup = g
  147. common.SetContextKey(c, constant.ContextKeyAutoGroup, g)
  148. channel = preferred
  149. service.MarkChannelAffinityUsed(c, g, preferred.Id)
  150. break
  151. }
  152. }
  153. } else if model.IsChannelEnabledForGroupModel(usingGroup, modelRequest.Model, preferred.Id) {
  154. channel = preferred
  155. selectGroup = usingGroup
  156. service.MarkChannelAffinityUsed(c, usingGroup, preferred.Id)
  157. }
  158. }
  159. }
  160. }
  161. if channel == nil {
  162. channel, selectGroup, err = service.CacheGetRandomSatisfiedChannel(&service.RetryParam{
  163. Ctx: c,
  164. ModelName: modelRequest.Model,
  165. TokenGroup: usingGroup,
  166. Retry: common.GetPointer(0),
  167. })
  168. if err != nil {
  169. showGroup := usingGroup
  170. if usingGroup == "auto" {
  171. showGroup = fmt.Sprintf("auto(%s)", selectGroup)
  172. }
  173. message := i18n.T(c, i18n.MsgDistributorGetChannelFailed, map[string]any{"Group": showGroup, "Model": modelRequest.Model, "Error": err.Error()})
  174. abortWithOpenAiMessage(c, http.StatusServiceUnavailable, message, types.ErrorCodeModelNotFound)
  175. return
  176. }
  177. if channel == nil {
  178. abortWithOpenAiMessage(c, http.StatusServiceUnavailable, i18n.T(c, i18n.MsgDistributorNoAvailableChannel, map[string]any{"Group": usingGroup, "Model": modelRequest.Model}), types.ErrorCodeModelNotFound)
  179. return
  180. }
  181. }
  182. }
  183. }
  184. common.SetContextKey(c, constant.ContextKeyRequestStartTime, time.Now())
  185. SetupContextForSelectedChannel(c, channel, modelRequest.Model)
  186. c.Next()
  187. if channel != nil && c.Writer != nil && c.Writer.Status() < http.StatusBadRequest {
  188. service.RecordChannelAffinity(c, channel.Id)
  189. }
  190. }
  191. }
  192. // getModelFromRequest 从请求中读取模型信息
  193. // 根据 Content-Type 自动处理:
  194. // - application/json
  195. // - application/x-www-form-urlencoded
  196. // - multipart/form-data
  197. func getModelFromRequest(c *gin.Context) (*ModelRequest, error) {
  198. var modelRequest ModelRequest
  199. err := common.UnmarshalBodyReusable(c, &modelRequest)
  200. if err != nil {
  201. return nil, errors.New(i18n.T(c, i18n.MsgDistributorInvalidRequest, map[string]any{"Error": err.Error()}))
  202. }
  203. return &modelRequest, nil
  204. }
  205. func getModelRequest(c *gin.Context) (*ModelRequest, bool, error) {
  206. var modelRequest ModelRequest
  207. shouldSelectChannel := true
  208. var err error
  209. if strings.Contains(c.Request.URL.Path, "/mj/") {
  210. relayMode := relayconstant.Path2RelayModeMidjourney(c.Request.URL.Path)
  211. if relayMode == relayconstant.RelayModeMidjourneyTaskFetch ||
  212. relayMode == relayconstant.RelayModeMidjourneyTaskFetchByCondition ||
  213. relayMode == relayconstant.RelayModeMidjourneyNotify ||
  214. relayMode == relayconstant.RelayModeMidjourneyTaskImageSeed {
  215. shouldSelectChannel = false
  216. } else {
  217. midjourneyRequest := dto.MidjourneyRequest{}
  218. err = common.UnmarshalBodyReusable(c, &midjourneyRequest)
  219. if err != nil {
  220. return nil, false, errors.New(i18n.T(c, i18n.MsgDistributorInvalidMidjourney, map[string]any{"Error": err.Error()}))
  221. }
  222. midjourneyModel, mjErr, success := service.GetMjRequestModel(relayMode, &midjourneyRequest)
  223. if mjErr != nil {
  224. return nil, false, fmt.Errorf("%s", mjErr.Description)
  225. }
  226. if midjourneyModel == "" {
  227. if !success {
  228. return nil, false, fmt.Errorf("%s", i18n.T(c, i18n.MsgDistributorInvalidParseModel))
  229. } else {
  230. // task fetch, task fetch by condition, notify
  231. shouldSelectChannel = false
  232. }
  233. }
  234. modelRequest.Model = midjourneyModel
  235. }
  236. c.Set("relay_mode", relayMode)
  237. } else if strings.Contains(c.Request.URL.Path, "/suno/") {
  238. relayMode := relayconstant.Path2RelaySuno(c.Request.Method, c.Request.URL.Path)
  239. if relayMode == relayconstant.RelayModeSunoFetch ||
  240. relayMode == relayconstant.RelayModeSunoFetchByID {
  241. shouldSelectChannel = false
  242. } else {
  243. modelName := service.CoverTaskActionToModelName(constant.TaskPlatformSuno, c.Param("action"))
  244. modelRequest.Model = modelName
  245. }
  246. c.Set("platform", string(constant.TaskPlatformSuno))
  247. c.Set("relay_mode", relayMode)
  248. } else if strings.Contains(c.Request.URL.Path, "/v1/videos/") && strings.HasSuffix(c.Request.URL.Path, "/remix") {
  249. relayMode := relayconstant.RelayModeVideoSubmit
  250. c.Set("relay_mode", relayMode)
  251. shouldSelectChannel = false
  252. } else if strings.Contains(c.Request.URL.Path, "/v1/videos") {
  253. //curl https://api.openai.com/v1/videos \
  254. // -H "Authorization: Bearer $OPENAI_API_KEY" \
  255. // -F "model=sora-2" \
  256. // -F "prompt=A calico cat playing a piano on stage"
  257. // -F input_reference="@image.jpg"
  258. relayMode := relayconstant.RelayModeUnknown
  259. if c.Request.Method == http.MethodPost {
  260. relayMode = relayconstant.RelayModeVideoSubmit
  261. req, err := getModelFromRequest(c)
  262. if err != nil {
  263. return nil, false, err
  264. }
  265. if req != nil {
  266. modelRequest.Model = req.Model
  267. }
  268. } else if c.Request.Method == http.MethodGet {
  269. relayMode = relayconstant.RelayModeVideoFetchByID
  270. shouldSelectChannel = false
  271. }
  272. c.Set("relay_mode", relayMode)
  273. } else if strings.Contains(c.Request.URL.Path, "/v1/video/generations") {
  274. relayMode := relayconstant.RelayModeUnknown
  275. if c.Request.Method == http.MethodPost {
  276. req, err := getModelFromRequest(c)
  277. if err != nil {
  278. return nil, false, err
  279. }
  280. modelRequest.Model = req.Model
  281. relayMode = relayconstant.RelayModeVideoSubmit
  282. } else if c.Request.Method == http.MethodGet {
  283. relayMode = relayconstant.RelayModeVideoFetchByID
  284. shouldSelectChannel = false
  285. }
  286. if _, ok := c.Get("relay_mode"); !ok {
  287. c.Set("relay_mode", relayMode)
  288. }
  289. } else if strings.HasPrefix(c.Request.URL.Path, "/v1beta/models/") || strings.HasPrefix(c.Request.URL.Path, "/v1/models/") {
  290. // Gemini API 路径处理: /v1beta/models/gemini-2.0-flash:generateContent
  291. relayMode := relayconstant.RelayModeGemini
  292. modelName := extractModelNameFromGeminiPath(c.Request.URL.Path)
  293. if modelName != "" {
  294. modelRequest.Model = modelName
  295. }
  296. c.Set("relay_mode", relayMode)
  297. } else if !strings.HasPrefix(c.Request.URL.Path, "/v1/audio/transcriptions") && !strings.Contains(c.Request.Header.Get("Content-Type"), "multipart/form-data") {
  298. req, err := getModelFromRequest(c)
  299. if err != nil {
  300. return nil, false, err
  301. }
  302. modelRequest.Model = req.Model
  303. }
  304. if strings.HasPrefix(c.Request.URL.Path, "/v1/realtime") {
  305. //wss://api.openai.com/v1/realtime?model=gpt-4o-realtime-preview-2024-10-01
  306. modelRequest.Model = c.Query("model")
  307. }
  308. if strings.HasPrefix(c.Request.URL.Path, "/v1/moderations") {
  309. if modelRequest.Model == "" {
  310. modelRequest.Model = "text-moderation-stable"
  311. }
  312. }
  313. if strings.HasSuffix(c.Request.URL.Path, "embeddings") {
  314. if modelRequest.Model == "" {
  315. modelRequest.Model = c.Param("model")
  316. }
  317. }
  318. if strings.HasPrefix(c.Request.URL.Path, "/v1/images/generations") {
  319. modelRequest.Model = common.GetStringIfEmpty(modelRequest.Model, "dall-e")
  320. } else if strings.HasPrefix(c.Request.URL.Path, "/v1/images/edits") {
  321. //modelRequest.Model = common.GetStringIfEmpty(c.PostForm("model"), "gpt-image-1")
  322. contentType := c.ContentType()
  323. if slices.Contains([]string{gin.MIMEPOSTForm, gin.MIMEMultipartPOSTForm}, contentType) {
  324. req, err := getModelFromRequest(c)
  325. if err == nil && req.Model != "" {
  326. modelRequest.Model = req.Model
  327. }
  328. }
  329. }
  330. if strings.HasPrefix(c.Request.URL.Path, "/v1/audio") {
  331. relayMode := relayconstant.RelayModeAudioSpeech
  332. if strings.HasPrefix(c.Request.URL.Path, "/v1/audio/speech") {
  333. modelRequest.Model = common.GetStringIfEmpty(modelRequest.Model, "tts-1")
  334. } else if strings.HasPrefix(c.Request.URL.Path, "/v1/audio/translations") {
  335. // 先尝试从请求读取
  336. if req, err := getModelFromRequest(c); err == nil && req.Model != "" {
  337. modelRequest.Model = req.Model
  338. }
  339. modelRequest.Model = common.GetStringIfEmpty(modelRequest.Model, "whisper-1")
  340. relayMode = relayconstant.RelayModeAudioTranslation
  341. } else if strings.HasPrefix(c.Request.URL.Path, "/v1/audio/transcriptions") {
  342. // 先尝试从请求读取
  343. if req, err := getModelFromRequest(c); err == nil && req.Model != "" {
  344. modelRequest.Model = req.Model
  345. }
  346. modelRequest.Model = common.GetStringIfEmpty(modelRequest.Model, "whisper-1")
  347. relayMode = relayconstant.RelayModeAudioTranscription
  348. }
  349. c.Set("relay_mode", relayMode)
  350. }
  351. if strings.HasPrefix(c.Request.URL.Path, "/pg/chat/completions") {
  352. // playground chat completions
  353. req, err := getModelFromRequest(c)
  354. if err != nil {
  355. return nil, false, err
  356. }
  357. modelRequest.Model = req.Model
  358. modelRequest.Group = req.Group
  359. common.SetContextKey(c, constant.ContextKeyTokenGroup, modelRequest.Group)
  360. if req.ChannelId > 0 {
  361. common.SetContextKey(c, constant.ContextKeyTokenSpecificChannelId, strconv.Itoa(req.ChannelId))
  362. }
  363. }
  364. if strings.HasPrefix(c.Request.URL.Path, "/v1/responses/compact") && modelRequest.Model != "" {
  365. modelRequest.Model = ratio_setting.WithCompactModelSuffix(modelRequest.Model)
  366. }
  367. return &modelRequest, shouldSelectChannel, nil
  368. }
  369. func SetupContextForSelectedChannel(c *gin.Context, channel *model.Channel, modelName string) *types.NewAPIError {
  370. c.Set("original_model", modelName) // for retry
  371. if channel == nil {
  372. return types.NewError(errors.New("channel is nil"), types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry())
  373. }
  374. common.SetContextKey(c, constant.ContextKeyChannelId, channel.Id)
  375. common.SetContextKey(c, constant.ContextKeyChannelName, channel.Name)
  376. common.SetContextKey(c, constant.ContextKeyChannelType, channel.Type)
  377. common.SetContextKey(c, constant.ContextKeyChannelCreateTime, channel.CreatedTime)
  378. common.SetContextKey(c, constant.ContextKeyChannelSetting, channel.GetSetting())
  379. common.SetContextKey(c, constant.ContextKeyChannelOtherSetting, channel.GetOtherSettings())
  380. common.SetContextKey(c, constant.ContextKeyChannelParamOverride, channel.GetParamOverride())
  381. common.SetContextKey(c, constant.ContextKeyChannelHeaderOverride, channel.GetHeaderOverride())
  382. if nil != channel.OpenAIOrganization && *channel.OpenAIOrganization != "" {
  383. common.SetContextKey(c, constant.ContextKeyChannelOrganization, *channel.OpenAIOrganization)
  384. }
  385. common.SetContextKey(c, constant.ContextKeyChannelAutoBan, channel.GetAutoBan())
  386. common.SetContextKey(c, constant.ContextKeyChannelModelMapping, channel.GetModelMapping())
  387. common.SetContextKey(c, constant.ContextKeyChannelStatusCodeMapping, channel.GetStatusCodeMapping())
  388. key, index, newAPIError := channel.GetNextEnabledKey()
  389. if newAPIError != nil {
  390. return newAPIError
  391. }
  392. if channel.ChannelInfo.IsMultiKey {
  393. common.SetContextKey(c, constant.ContextKeyChannelIsMultiKey, true)
  394. common.SetContextKey(c, constant.ContextKeyChannelMultiKeyIndex, index)
  395. } else {
  396. // 必须设置为 false,否则在重试到单个 key 的时候会导致日志显示错误
  397. common.SetContextKey(c, constant.ContextKeyChannelIsMultiKey, false)
  398. }
  399. // c.Request.Header.Set("Authorization", fmt.Sprintf("Bearer %s", key))
  400. common.SetContextKey(c, constant.ContextKeyChannelKey, key)
  401. common.SetContextKey(c, constant.ContextKeyChannelBaseUrl, channel.GetBaseURL())
  402. common.SetContextKey(c, constant.ContextKeySystemPromptOverride, false)
  403. // TODO: api_version统一
  404. switch channel.Type {
  405. case constant.ChannelTypeAzure:
  406. c.Set("api_version", channel.Other)
  407. case constant.ChannelTypeVertexAi:
  408. c.Set("region", channel.Other)
  409. case constant.ChannelTypeXunfei:
  410. c.Set("api_version", channel.Other)
  411. case constant.ChannelTypeGemini:
  412. c.Set("api_version", channel.Other)
  413. case constant.ChannelTypeAli:
  414. c.Set("plugin", channel.Other)
  415. case constant.ChannelCloudflare:
  416. c.Set("api_version", channel.Other)
  417. case constant.ChannelTypeMokaAI:
  418. c.Set("api_version", channel.Other)
  419. case constant.ChannelTypeCoze:
  420. c.Set("bot_id", channel.Other)
  421. }
  422. return nil
  423. }
  424. // extractModelNameFromGeminiPath 从 Gemini API URL 路径中提取模型名
  425. // 输入格式: /v1beta/models/gemini-2.0-flash:generateContent
  426. // 输出: gemini-2.0-flash
  427. func extractModelNameFromGeminiPath(path string) string {
  428. // 查找 "/models/" 的位置
  429. modelsPrefix := "/models/"
  430. modelsIndex := strings.Index(path, modelsPrefix)
  431. if modelsIndex == -1 {
  432. return ""
  433. }
  434. // 从 "/models/" 之后开始提取
  435. startIndex := modelsIndex + len(modelsPrefix)
  436. if startIndex >= len(path) {
  437. return ""
  438. }
  439. // 查找 ":" 的位置,模型名在 ":" 之前
  440. colonIndex := strings.Index(path[startIndex:], ":")
  441. if colonIndex == -1 {
  442. // 如果没有找到 ":",返回从 "/models/" 到路径结尾的部分
  443. return path[startIndex:]
  444. }
  445. // 返回模型名部分
  446. return path[startIndex : startIndex+colonIndex]
  447. }