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.
 
 
 

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