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) {

relay/responses_websocket.go

@@ -281,7 +281,8 @@ func (s *responsesWSSession) runCall(c *gin.Context, state *responsesWSCallState
 		}
 	} else {
 		retry := &service.RetryParam{Ctx: c, TokenGroup: common.GetContextKeyString(c, appconstant.ContextKeyUsingGroup), ModelName: modelName, RequestPath: c.Request.URL.Path, Retry: common.GetPointer(0)}
-		for ; retry.GetRetry() <= common.RetryTimes; retry.IncreaseRetry() {
+		maxRetries := max(common.RetryTimes, service.PoolRetryBudget(c, modelName))
+		for ; retry.GetRetry() <= maxRetries; retry.IncreaseRetry() {
 			var channel *appmodel.Channel
 			channel, apiErr = selectResponsesWSChannel(c, modelName, retry)
 			if apiErr != nil {
@@ -316,7 +317,15 @@ func (s *responsesWSSession) runCall(c *gin.Context, state *responsesWSCallState
 				apiErr = service.NormalizeViolationFeeError(types.NewError(dialErr, types.ErrorCodeDoRequestFailed))
 				service.ResetStatusCode(apiErr, c.GetString("status_code_mapping"))
 				info.LastError = apiErr
-				decision := service.DecideRelayRetry(c, apiErr, common.RetryTimes-retry.GetRetry())
+				remaining := common.RetryTimes - retry.GetRetry()
+				if service.IsPoolFailoverError(c, apiErr) {
+					remaining = maxRetries - retry.GetRetry()
+					service.RecordPoolRetryableFailure(c, modelName, info.UsingGroup)
+				}
+				decision := service.DecideRelayRetry(c, apiErr, remaining)
+				if service.IsPoolFailoverError(c, apiErr) && service.PoolRetryExhausted(c, modelName, info.UsingGroup) {
+					decision = service.PolicyDecision{Action: "stop", Reason: "pool_no_backup", Source: "routing"}
+				}
 				service.RecordPolicyFailure(c, channel.Id, apiErr, decision)
 				service.ProcessChannelError(c, *types.NewChannelError(channel.Id, channel.Type, channel.Name, channel.ChannelInfo.IsMultiKey, info.ApiKey, channel.GetAutoBan()), apiErr, info)
 				if decision.Action == "retry" {
@@ -571,6 +580,13 @@ func (s *responsesWSSession) restoreConnectionContext(c *gin.Context, model stri
 		}
 	} else {
 		group := common.GetContextKeyString(c, appconstant.ContextKeyUsingGroup)
+		if service.RequestRoutingPolicy(c) != nil {
+			if !slices.ContainsFunc(service.ResolveRequestPool(c, model), func(entry appmodel.TokenGroupEntry) bool { return entry.Group == s.lockedGroup }) {
+				return types.NewErrorWithStatusCode(errors.New("the connection group is no longer allowed"), types.ErrorCodeAccessDenied, http.StatusForbidden, types.ErrOptionWithSkipRetry())
+			}
+			group = s.lockedGroup
+			common.SetContextKey(c, appconstant.ContextKeyAutoGroup, group)
+		}
 		if group == "auto" {
 			if !slices.Contains(service.GetRequestAutoGroups(c, common.GetContextKeyString(c, appconstant.ContextKeyUserGroup)), s.lockedGroup) {
 				return types.NewErrorWithStatusCode(errors.New("the connection group is no longer allowed"), types.ErrorCodeAccessDenied, http.StatusForbidden, types.ErrOptionWithSkipRetry())