Vous ne pouvez pas sélectionner plus de 25 sujets Les noms de sujets doivent commencer par une lettre ou un nombre, peuvent contenir des tirets ('-') et peuvent comporter jusqu'à 35 caractères.
 
 
 

143 lignes
4.2 KiB

  1. package middleware
  2. import (
  3. "fmt"
  4. "net/http"
  5. "slices"
  6. "strings"
  7. "github.com/QuantumNous/new-api/common"
  8. "github.com/QuantumNous/new-api/constant"
  9. "github.com/QuantumNous/new-api/setting/model_setting"
  10. "github.com/gin-gonic/gin"
  11. )
  12. // pathToFamily maps URL path prefixes to format families.
  13. // The middleware uses this to determine the request's protocol family
  14. // directly from the request path, without relying on context values.
  15. var pathToFamily = []struct {
  16. prefix string
  17. family string
  18. }{
  19. {"/v1/messages", "claude"},
  20. {"/v1beta/", "gemini"},
  21. {"/v1/models/", "gemini"}, // Gemini relay via /v1/models/*path
  22. {"/v1/engines/", "gemini"}, // Gemini relay via /v1/engines/:model/embeddings
  23. // All other /v1/ paths default to "openai"
  24. }
  25. // getFamilyFromPath returns the format family for a given request path.
  26. func getFamilyFromPath(path string) string {
  27. for _, rule := range pathToFamily {
  28. if strings.HasPrefix(path, rule.prefix) {
  29. return rule.family
  30. }
  31. }
  32. if strings.HasPrefix(path, "/v1/") {
  33. return "openai"
  34. }
  35. return ""
  36. }
  37. var channelNativeFamilies = map[int][]string{
  38. constant.ChannelTypeOpenAI: {"openai"},
  39. constant.ChannelTypeAzure: {"openai"},
  40. constant.ChannelTypeOllama: {"openai"},
  41. constant.ChannelTypeDeepSeek: {"openai"},
  42. constant.ChannelTypeOpenRouter: {"openai"},
  43. constant.ChannelTypeXinference: {"openai"},
  44. constant.ChannelTypeXai: {"openai"},
  45. constant.ChannelTypeMistral: {"openai"},
  46. constant.ChannelTypeSiliconFlow: {"openai"},
  47. constant.ChannelTypePerplexity: {"openai"},
  48. constant.ChannelTypeCodex: {"openai"},
  49. constant.ChannelTypeMiniMax: {"openai"},
  50. constant.ChannelTypeMoonshot: {"openai"},
  51. constant.ChannelTypeBaidu: {"openai"},
  52. constant.ChannelTypeBaiduV2: {"openai"},
  53. constant.ChannelTypeZhipu: {"openai"},
  54. constant.ChannelTypeZhipu_v4: {"openai"},
  55. constant.ChannelTypeTencent: {"openai"},
  56. constant.ChannelTypeXunfei: {"openai"},
  57. constant.ChannelTypeCohere: {"openai"},
  58. constant.ChannelTypeDify: {"openai"},
  59. constant.ChannelTypeCoze: {"openai"},
  60. constant.ChannelTypeVolcEngine: {"openai"},
  61. constant.ChannelTypeReplicate: {"openai"},
  62. constant.ChannelTypeSubmodel: {"openai"},
  63. constant.ChannelTypeJina: {"openai"},
  64. constant.ChannelTypePaLM: {"openai"},
  65. constant.ChannelTypeAnthropic: {"claude"},
  66. constant.ChannelTypeAws: {"claude"},
  67. constant.ChannelTypeGemini: {"gemini"},
  68. constant.ChannelTypeVertexAi: {"openai", "claude", "gemini"},
  69. constant.ChannelTypeAli: {"openai", "claude"},
  70. }
  71. var nativeEndpoints = map[string]string{
  72. "openai": "/v1/chat/completions",
  73. "claude": "/v1/messages",
  74. "gemini": "/v1beta/models/{model}:generateContent",
  75. }
  76. func isChannelAllowed(channelType int, family string) bool {
  77. families, ok := channelNativeFamilies[channelType]
  78. if !ok {
  79. return true
  80. }
  81. return slices.Contains(families, family)
  82. }
  83. func EndpointFormatGuard() gin.HandlerFunc {
  84. return func(c *gin.Context) {
  85. if !model_setting.IsEndpointFormatRestrictionEnabled() {
  86. c.Next()
  87. return
  88. }
  89. family := getFamilyFromPath(c.Request.URL.Path)
  90. if family == "" {
  91. c.Next()
  92. return
  93. }
  94. channelType := common.GetContextKeyInt(c, constant.ContextKeyChannelType)
  95. if channelType == 0 {
  96. c.Next()
  97. return
  98. }
  99. if isChannelAllowed(channelType, family) {
  100. c.Next()
  101. return
  102. }
  103. channelTypeName := constant.GetChannelTypeName(channelType)
  104. var nativeEndpoint string
  105. if families, ok := channelNativeFamilies[channelType]; ok && len(families) > 0 {
  106. nativeEndpoint = nativeEndpoints[families[0]]
  107. }
  108. c.JSON(http.StatusBadRequest, gin.H{
  109. "error": gin.H{
  110. "message": fmt.Sprintf(
  111. "Model is on a %s channel and cannot be accessed via %s endpoints. Please use the native %s endpoint (%s) instead.",
  112. channelTypeName, family, channelTypeName, nativeEndpoint,
  113. ),
  114. "type": "new_api_error",
  115. "code": "endpoint_format_restricted",
  116. },
  117. })
  118. c.Abort()
  119. }
  120. }
  121. // setTestRestrictionEnabled is a test helper.
  122. func setTestRestrictionEnabled(enabled bool) bool {
  123. prev := model_setting.IsEndpointFormatRestrictionEnabled()
  124. model_setting.GetGlobalSettings().EndpointFormatRestrictionEnabled = enabled
  125. return prev
  126. }