fix: improve playground endpoint support / resolve issue #4503

controller/playground.go

@@ -4,15 +4,18 @@ import (
 	"errors"
 	"fmt"
 
+	"github.com/QuantumNous/new-api/common"
+	"github.com/QuantumNous/new-api/constant"
+	"github.com/QuantumNous/new-api/dto"
 	"github.com/QuantumNous/new-api/middleware"
 	"github.com/QuantumNous/new-api/model"
-	relaycommon "github.com/QuantumNous/new-api/relay/common"
+	"github.com/QuantumNous/new-api/service"
 	"github.com/QuantumNous/new-api/types"
 
 	"github.com/gin-gonic/gin"
 )
 
-func Playground(c *gin.Context) {
+func Playground(c *gin.Context, relayFormat types.RelayFormat) {
 	var newAPIError *types.NewAPIError
 
 	defer func() {
@@ -29,12 +32,6 @@ func Playground(c *gin.Context) {
 		return
 	}
 
-	relayInfo, err := relaycommon.GenRelayInfo(c, types.RelayFormatOpenAI, nil, nil)
-	if err != nil {
-		newAPIError = types.NewError(err, types.ErrorCodeInvalidRequest, types.ErrOptionWithSkipRetry())
-		return
-	}
-
 	userId := c.GetInt("id")
 
 	// Write user context to ensure acceptUnsetRatio is available
@@ -45,12 +42,28 @@ func Playground(c *gin.Context) {
 	}
 	userCache.WriteContext(c)
 
+	playgroundRequest := &dto.PlayGroundRequest{}
+	if err := common.UnmarshalBodyReusable(c, playgroundRequest); err != nil {
+		newAPIError = types.NewError(err, types.ErrorCodeInvalidRequest, types.ErrOptionWithSkipRetry())
+		return
+	}
+
+	usingGroup := common.GetContextKeyString(c, constant.ContextKeyUsingGroup)
+	if playgroundRequest.Group != "" {
+		userGroup := common.GetContextKeyString(c, constant.ContextKeyUserGroup)
+		if playgroundRequest.Group != userGroup && !service.GroupInUserUsableGroups(userGroup, playgroundRequest.Group) {
+			newAPIError = types.NewError(errors.New("group access denied"), types.ErrorCodeAccessDenied, types.ErrOptionWithSkipRetry())
+			return
+		}
+		usingGroup = playgroundRequest.Group
+	}
+
 	tempToken := &model.Token{
 		UserId: userId,
-		Name:   fmt.Sprintf("playground-%s", relayInfo.UsingGroup),
-		Group:  relayInfo.UsingGroup,
+		Name:   fmt.Sprintf("playground-%s", usingGroup),
+		Group:  usingGroup,
 	}
 	_ = middleware.SetupContextForToken(c, tempToken)
 
-	Relay(c, types.RelayFormatOpenAI)
+	Relay(c, relayFormat)
 }

middleware/distributor.go

@@ -82,7 +82,7 @@ func Distribute() func(c *gin.Context) {
 				var selectGroup string
 				usingGroup := common.GetContextKeyString(c, constant.ContextKeyUsingGroup)
 				// check path is /pg/chat/completions
-				if strings.HasPrefix(c.Request.URL.Path, "/pg/chat/completions") {
+				if isPlaygroundPath(c.Request.URL.Path) {
 					playgroundRequest := &dto.PlayGroundRequest{}
 					err = common.UnmarshalBodyReusable(c, playgroundRequest)
 					if err != nil {
@@ -325,14 +325,16 @@ func getModelRequest(c *gin.Context) (*ModelRequest, bool, error) {
 		}
 		c.Set("relay_mode", relayMode)
 	}
-	if strings.HasPrefix(c.Request.URL.Path, "/pg/chat/completions") {
-		// playground chat completions
+	if isPlaygroundPath(c.Request.URL.Path) {
 		req, err := getModelFromRequest(c)
 		if err != nil {
 			return nil, false, err
 		}
 		modelRequest.Model = req.Model
 		modelRequest.Group = req.Group
+		if strings.HasPrefix(c.Request.URL.Path, "/pg/images/generations") {
+			modelRequest.Model = common.GetStringIfEmpty(modelRequest.Model, "dall-e")
+		}
 		common.SetContextKey(c, constant.ContextKeyTokenGroup, modelRequest.Group)
 	}
 
@@ -342,6 +344,10 @@ func getModelRequest(c *gin.Context) (*ModelRequest, bool, error) {
 	return &modelRequest, shouldSelectChannel, nil
 }
 
+func isPlaygroundPath(path string) bool {
+	return strings.HasPrefix(path, "/pg/")
+}
+
 func SetupContextForSelectedChannel(c *gin.Context, channel *model.Channel, modelName string) *types.NewAPIError {
 	c.Set("original_model", modelName) // for retry
 	if channel == nil {

relay/constant/relay_mode.go

@@ -66,15 +66,15 @@ func Path2RelayMode(path string) int {
 		relayMode = RelayModeEmbeddings
 	} else if strings.HasPrefix(path, "/v1/moderations") {
 		relayMode = RelayModeModerations
-	} else if strings.HasPrefix(path, "/v1/images/generations") {
+	} else if strings.HasPrefix(path, "/v1/images/generations") || strings.HasPrefix(path, "/pg/images/generations") {
 		relayMode = RelayModeImagesGenerations
 	} else if strings.HasPrefix(path, "/v1/images/edits") {
 		relayMode = RelayModeImagesEdits
 	} else if strings.HasPrefix(path, "/v1/edits") {
 		relayMode = RelayModeEdits
 	} else if strings.HasPrefix(path, "/v1/responses/compact") {
 		relayMode = RelayModeResponsesCompact
-	} else if strings.HasPrefix(path, "/v1/responses") {
+	} else if strings.HasPrefix(path, "/v1/responses") || strings.HasPrefix(path, "/pg/responses") {
 		relayMode = RelayModeResponses
 	} else if strings.HasPrefix(path, "/v1/audio/speech") {
 		relayMode = RelayModeAudioSpeech

relay/helper/valid_request.go

@@ -4,6 +4,8 @@ import (
 	"errors"
 	"fmt"
 	"math"
+	"regexp"
+	"strconv"
 	"strings"
 
 	"github.com/QuantumNous/new-api/common"
@@ -138,6 +140,50 @@ func GetAndValidateResponsesCompactionRequest(c *gin.Context) (*dto.OpenAIRespon
 	return request, nil
 }
 
+var imageSizePattern = regexp.MustCompile(`^([1-9]\d*)[xX]([1-9]\d*)$`)
+
+var validImageQualities = map[string]struct{}{
+	"low":      {},
+	"medium":   {},
+	"high":     {},
+	"auto":     {},
+	"standard": {},
+	"hd":       {},
+}
+
+func validateOpenAIImageRequest(imageRequest *dto.ImageRequest) error {
+	if imageRequest.Quality != "" {
+		if _, ok := validImageQualities[imageRequest.Quality]; !ok {
+			return errors.New("invalid quality, must be one of: low, medium, high, auto, standard, hd")
+		}
+	}
+
+	if strings.Contains(imageRequest.Size, "×") {
+		return errors.New("size an unexpected error occurred in the parameter, please use 'x' instead of the multiplication sign '×'")
+	}
+
+	if imageRequest.Size != "" {
+		matches := imageSizePattern.FindStringSubmatch(imageRequest.Size)
+		if matches == nil {
+			return errors.New("size must use axb format with positive integer dimensions")
+		}
+
+		width, err := strconv.Atoi(matches[1])
+		if err != nil || width <= 0 {
+			return errors.New("size must use axb format with positive integer dimensions")
+		}
+		height, err := strconv.Atoi(matches[2])
+		if err != nil || height <= 0 {
+			return errors.New("size must use axb format with positive integer dimensions")
+		}
+		if width > 3840 || height > 3840 {
+			return errors.New("size width and height must be between 1 and 3840")
+		}
+	}
+
+	return nil
+}
+
 func GetAndValidOpenAIImageRequest(c *gin.Context, relayMode int) (*dto.ImageRequest, error) {
 	imageRequest := &dto.ImageRequest{}
 
@@ -158,6 +204,10 @@ func GetAndValidOpenAIImageRequest(c *gin.Context, relayMode int) (*dto.ImageReq
 				imageRequest.Image, _ = common.Marshal(imageValue)
 			}
 
+			if err := validateOpenAIImageRequest(imageRequest); err != nil {
+				return nil, err
+			}
+
 			if imageRequest.Model == "gpt-image-1" {
 				if imageRequest.Quality == "" {
 					imageRequest.Quality = "standard"
@@ -182,12 +232,11 @@ func GetAndValidOpenAIImageRequest(c *gin.Context, relayMode int) (*dto.ImageReq
 		}
 
 		if imageRequest.Model == "" {
-			//imageRequest.Model = "dall-e-3"
 			return nil, errors.New("model is required")
 		}
 
-		if strings.Contains(imageRequest.Size, "×") {
-			return nil, errors.New("size an unexpected error occurred in the parameter, please use 'x' instead of the multiplication sign '×'")
+		if err := validateOpenAIImageRequest(imageRequest); err != nil {
+			return nil, err
 		}
 
 		// Not "256x256", "512x512", or "1024x1024"
@@ -214,10 +263,6 @@ func GetAndValidOpenAIImageRequest(c *gin.Context, relayMode int) (*dto.ImageReq
 			}
 		}
 
-		//if imageRequest.Prompt == "" {
-		//	return nil, errors.New("prompt is required")
-		//}
-
 		if imageRequest.N == nil || *imageRequest.N == 0 {
 			imageRequest.N = common.GetPointer(uint(1))
 		}

router/relay-router.go

@@ -10,6 +10,12 @@ import (
 	"github.com/gin-gonic/gin"
 )
 
+func playgroundHandler(relayFormat types.RelayFormat) gin.HandlerFunc {
+	return func(c *gin.Context) {
+		controller.Playground(c, relayFormat)
+	}
+}
+
 func SetRelayRouter(router *gin.Engine) {
 	router.Use(middleware.CORS())
 	router.Use(middleware.DecompressRequestMiddleware())
@@ -64,7 +70,10 @@ func SetRelayRouter(router *gin.Engine) {
 	playgroundRouter.Use(middleware.SystemPerformanceCheck())
 	playgroundRouter.Use(middleware.UserAuth(), middleware.Distribute())
 	{
-		playgroundRouter.POST("/chat/completions", controller.Playground)
+		playgroundRouter.POST("/chat/completions", playgroundHandler(types.RelayFormatOpenAI))
+		playgroundRouter.POST("/responses", playgroundHandler(types.RelayFormatOpenAIResponses))
+		playgroundRouter.POST("/messages", playgroundHandler(types.RelayFormatClaude))
+		playgroundRouter.POST("/images/generations", playgroundHandler(types.RelayFormatOpenAIImage))
 	}
 	relayV1Router := router.Group("/v1")
 	relayV1Router.Use(middleware.RouteTag("relay"))