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

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