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

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)
+}