Вы не можете выбрать более 25 тем Темы должны начинаться с буквы или цифры, могут содержать дефисы(-) и должны содержать не более 35 символов.
 
 
 

411 строки
16 KiB

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