feat: sync deployed v1.0.0-rc.38-custom.67

relay/helper/image_generation_resolution.go

@@ -0,0 +1,59 @@
+package helper
+
+import (
+	"fmt"
+	relaycommon "github.com/QuantumNous/new-api/relay/common"
+	"github.com/QuantumNous/new-api/relaykit/dto"
+	"github.com/QuantumNous/new-api/setting/model_setting"
+	"slices"
+	"strconv"
+	"strings"
+)
+
+func imageGenerationResolution(size string) (string, error) {
+	size = strings.ToLower(strings.TrimSpace(size))
+	if size == "" || size == "auto" || size == "1k" {
+		return "1k", nil
+	}
+	if size == "2k" || size == "4k" {
+		return size, nil
+	}
+	dimensions := strings.Split(size, "x")
+	if len(dimensions) != 2 {
+		return "", fmt.Errorf("size must be auto, 1k, 2k, 4k or WIDTHxHEIGHT")
+	}
+	width, errWidth := strconv.Atoi(dimensions[0])
+	height, errHeight := strconv.Atoi(dimensions[1])
+	if errWidth != nil || errHeight != nil || width <= 0 || height <= 0 {
+		return "", fmt.Errorf("invalid image size")
+	}
+	short := min(width, height)
+	if short <= 1024 {
+		return "1k", nil
+	}
+	if short <= 2048 {
+		return "2k", nil
+	}
+	return "4k", nil
+}
+
+func mapImageGenerationResolution(info *relaycommon.RelayInfo, request dto.Request) error {
+	cfg := model_setting.GetGlobalSettings()
+	image, ok := request.(*dto.ImageRequest)
+	if !ok || cfg.ImageGenerationGroup == "" || info.UsingGroup != cfg.ImageGenerationGroup || !slices.Contains(cfg.ImageGenerationModels, info.OriginModelName) {
+		return nil
+	}
+	name := info.OriginModelName
+	for _, suffix := range []string{"-1k", "-2k", "-4k"} {
+		if strings.HasSuffix(name, suffix) {
+			return nil
+		}
+	}
+	tier, err := imageGenerationResolution(image.Size)
+	if err != nil {
+		return err
+	}
+	info.UpstreamModelName = name + "-" + tier
+	info.IsModelMapped = true
+	return nil
+}

relay/helper/image_generation_resolution_test.go

@@ -0,0 +1,29 @@
+package helper
+
+import (
+	relaycommon "github.com/QuantumNous/new-api/relay/common"
+	"github.com/QuantumNous/new-api/relaykit/dto"
+	"github.com/QuantumNous/new-api/setting/model_setting"
+	"github.com/stretchr/testify/require"
+	"testing"
+)
+
+func TestImageGenerationResolutionMapping(t *testing.T) {
+	cfg := model_setting.GetGlobalSettings()
+	oldGroup, oldModels := cfg.ImageGenerationGroup, cfg.ImageGenerationModels
+	t.Cleanup(func() { cfg.ImageGenerationGroup, cfg.ImageGenerationModels = oldGroup, oldModels })
+	cfg.ImageGenerationGroup, cfg.ImageGenerationModels = "example-images", []string{"gpt-image-2", "gpt-image-2-4k"}
+	for _, tc := range []struct{ size, tier string }{{"", "1k"}, {"auto", "1k"}, {"1024x1536", "1k"}, {"1536x1024", "1k"}, {"2048x1152", "2k"}, {"2048x2048", "2k"}, {"3840x2160", "4k"}, {"2160x3840", "4k"}, {"4k", "4k"}} {
+		info := &relaycommon.RelayInfo{OriginModelName: "gpt-image-2", UsingGroup: "example-images", ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "gpt-image-2"}}
+		image := &dto.ImageRequest{Model: "gpt-image-2", Size: tc.size}
+		require.NoError(t, mapImageGenerationResolution(info, image))
+		require.Equal(t, "gpt-image-2-"+tc.tier, info.UpstreamModelName)
+	}
+	_, err := imageGenerationResolution("-1x1024")
+	require.Error(t, err)
+	_, err = imageGenerationResolution("garbage")
+	require.Error(t, err)
+	info := &relaycommon.RelayInfo{OriginModelName: "gpt-image-2", UsingGroup: "default", ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "gpt-image-2"}}
+	require.NoError(t, mapImageGenerationResolution(info, &dto.ImageRequest{Model: "gpt-image-2", Size: "3840x2160"}))
+	require.Equal(t, "gpt-image-2", info.UpstreamModelName)
+}

relay/helper/model_mapped.go

@@ -62,6 +62,10 @@ func ModelMappedHelper(c *gin.Context, info *relaycommon.RelayInfo, request dto.
 		}
 	}
 
+	if err := mapImageGenerationResolution(info, request); err != nil {
+		return err
+	}
+
 	if request != nil {
 		request.SetModelName(info.UpstreamModelName)
 	}

relay/helper/price.go

@@ -47,31 +47,11 @@ const defaultTieredPreConsumeMaxTokens = 8192
 
 // HandleGroupRatio checks for "auto_group" in the context and updates the group ratio and relayInfo.UsingGroup if present
 func HandleGroupRatio(ctx *gin.Context, relayInfo *relaycommon.RelayInfo) hosttypes.GroupRatioInfo {
-	groupRatioInfo := hosttypes.GroupRatioInfo{
-		GroupRatio:        1.0, // default ratio
-		GroupSpecialRatio: -1,
-	}
-
-	// check auto group
-	autoGroup, exists := ctx.Get("auto_group")
-	if exists {
+	if autoGroup, exists := ctx.Get("auto_group"); exists {
 		logger.LogDebug(ctx, "final group: %s", autoGroup)
 		relayInfo.UsingGroup = autoGroup.(string)
 	}
-
-	// check user group special ratio
-	userGroupRatio, ok := ratio_setting.GetGroupGroupRatio(relayInfo.UserGroup, relayInfo.UsingGroup)
-	if ok {
-		// user group special ratio
-		groupRatioInfo.GroupSpecialRatio = userGroupRatio
-		groupRatioInfo.GroupRatio = userGroupRatio
-		groupRatioInfo.HasSpecialRatio = true
-	} else {
-		// normal group ratio
-		groupRatioInfo.GroupRatio = ratio_setting.GetGroupRatio(relayInfo.UsingGroup)
-	}
-
-	return groupRatioInfo
+	return model.EffectiveGroupRatio(ctx, relayInfo.UserGroup, relayInfo.UsingGroup)
 }
 
 func ModelPriceHelper(c *gin.Context, info *relaycommon.RelayInfo, promptTokens int, meta *types.TokenCountMeta) (hosttypes.PriceData, error) {