No puede seleccionar más de 25 temas Los temas deben comenzar con una letra o número, pueden incluir guiones ('-') y pueden tener hasta 35 caracteres de largo.
 
 
 

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