25개 이상의 토픽을 선택하실 수 없습니다. Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 

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