|
- package middleware
-
- import (
- "fmt"
- "net/http"
- "slices"
- "strings"
-
- "github.com/QuantumNous/new-api/common"
- "github.com/QuantumNous/new-api/constant"
- "github.com/QuantumNous/new-api/setting/model_setting"
-
- "github.com/gin-gonic/gin"
- )
-
- // pathToFamily maps URL path prefixes to format families.
- // The middleware uses this to determine the request's protocol family
- // directly from the request path, without relying on context values.
- var pathToFamily = []struct {
- prefix string
- family string
- }{
- {"/v1/messages", "claude"},
- {"/v1beta/", "gemini"},
- {"/v1/models/", "gemini"}, // Gemini relay via /v1/models/*path
- {"/v1/engines/", "gemini"}, // Gemini relay via /v1/engines/:model/embeddings
- // All other /v1/ paths default to "openai"
- }
-
- // getFamilyFromPath returns the format family for a given request path.
- func getFamilyFromPath(path string) string {
- for _, rule := range pathToFamily {
- if strings.HasPrefix(path, rule.prefix) {
- return rule.family
- }
- }
- if strings.HasPrefix(path, "/v1/") {
- return "openai"
- }
- return ""
- }
-
- var channelNativeFamilies = map[int][]string{
- constant.ChannelTypeOpenAI: {"openai"},
- constant.ChannelTypeAzure: {"openai"},
- constant.ChannelTypeOllama: {"openai"},
- constant.ChannelTypeDeepSeek: {"openai"},
- constant.ChannelTypeOpenRouter: {"openai"},
- constant.ChannelTypeXinference: {"openai"},
- constant.ChannelTypeXai: {"openai"},
- constant.ChannelTypeMistral: {"openai"},
- constant.ChannelTypeSiliconFlow: {"openai"},
- constant.ChannelTypePerplexity: {"openai"},
- constant.ChannelTypeCodex: {"openai"},
- constant.ChannelTypeMiniMax: {"openai"},
- constant.ChannelTypeMoonshot: {"openai"},
- constant.ChannelTypeBaidu: {"openai"},
- constant.ChannelTypeBaiduV2: {"openai"},
- constant.ChannelTypeZhipu: {"openai"},
- constant.ChannelTypeZhipu_v4: {"openai"},
- constant.ChannelTypeTencent: {"openai"},
- constant.ChannelTypeXunfei: {"openai"},
- constant.ChannelTypeCohere: {"openai"},
- constant.ChannelTypeDify: {"openai"},
- constant.ChannelTypeCoze: {"openai"},
- constant.ChannelTypeVolcEngine: {"openai"},
- constant.ChannelTypeReplicate: {"openai"},
- constant.ChannelTypeSubmodel: {"openai"},
- constant.ChannelTypeJina: {"openai"},
- constant.ChannelTypePaLM: {"openai"},
- constant.ChannelTypeAnthropic: {"claude"},
- constant.ChannelTypeAws: {"claude"},
- constant.ChannelTypeGemini: {"gemini"},
- constant.ChannelTypeVertexAi: {"openai", "claude", "gemini"},
- constant.ChannelTypeAli: {"openai", "claude"},
- }
-
- var nativeEndpoints = map[string]string{
- "openai": "/v1/chat/completions",
- "claude": "/v1/messages",
- "gemini": "/v1beta/models/{model}:generateContent",
- }
-
- func isChannelAllowed(channelType int, family string) bool {
- families, ok := channelNativeFamilies[channelType]
- if !ok {
- return true
- }
- return slices.Contains(families, family)
- }
-
- func EndpointFormatGuard() gin.HandlerFunc {
- return func(c *gin.Context) {
- if !model_setting.IsEndpointFormatRestrictionEnabled() {
- c.Next()
- return
- }
-
- family := getFamilyFromPath(c.Request.URL.Path)
- if family == "" {
- c.Next()
- return
- }
-
- channelType := common.GetContextKeyInt(c, constant.ContextKeyChannelType)
- if channelType == 0 {
- c.Next()
- return
- }
-
- if isChannelAllowed(channelType, family) {
- c.Next()
- return
- }
-
- channelTypeName := constant.GetChannelTypeName(channelType)
-
- var nativeEndpoint string
- if families, ok := channelNativeFamilies[channelType]; ok && len(families) > 0 {
- nativeEndpoint = nativeEndpoints[families[0]]
- }
-
- c.JSON(http.StatusBadRequest, gin.H{
- "error": gin.H{
- "message": fmt.Sprintf(
- "Model is on a %s channel and cannot be accessed via %s endpoints. Please use the native %s endpoint (%s) instead.",
- channelTypeName, family, channelTypeName, nativeEndpoint,
- ),
- "type": "new_api_error",
- "code": "endpoint_format_restricted",
- },
- })
- c.Abort()
- }
- }
-
- // setTestRestrictionEnabled is a test helper.
- func setTestRestrictionEnabled(enabled bool) bool {
- prev := model_setting.IsEndpointFormatRestrictionEnabled()
- model_setting.GetGlobalSettings().EndpointFormatRestrictionEnabled = enabled
- return prev
- }
|