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

service/image_generation_group.go

@@ -0,0 +1,58 @@
+package service
+
+import (
+	"github.com/QuantumNous/new-api/common"
+	"github.com/QuantumNous/new-api/constant"
+	"github.com/QuantumNous/new-api/model"
+	"github.com/QuantumNous/new-api/setting/model_setting"
+	"github.com/gin-gonic/gin"
+	"slices"
+)
+
+func imageGenerationPoolGroup(c *gin.Context) string {
+	settings := model_setting.GetGlobalSettings()
+	group := settings.ImageGenerationGroup
+	if group == "" || !IsUserSelectableGroup(common.GetContextKeyString(c, constant.ContextKeyUserGroup), group) {
+		return ""
+	}
+	policy := RequestRoutingPolicy(c)
+	if policy != nil {
+		if policy.Mode == "user" {
+			return ""
+		}
+		return group
+	}
+	if common.GetContextKeyString(c, constant.ContextKeyTokenGroup) == "auto" || common.GetContextKeyBool(c, constant.ContextKeyTokenCrossGroupRetry) {
+		return group
+	}
+	return ""
+}
+
+func imageGenerationRequestGroup(c *gin.Context, modelName string) string {
+	if !slices.Contains(model_setting.GetGlobalSettings().ImageGenerationModels, modelName) {
+		return ""
+	}
+	return imageGenerationPoolGroup(c)
+}
+
+// The dedicated route is shared by selection, affinity and the routing preview.
+func WithImageGenerationGroup(c *gin.Context, modelName string, entries []model.TokenGroupEntry) []model.TokenGroupEntry {
+	group := imageGenerationPoolGroup(c)
+	if group == "" {
+		return entries
+	}
+	if modelName != "" {
+		if imageGenerationRequestGroup(c, modelName) != "" {
+			return []model.TokenGroupEntry{{Group: group, Priority: 0}}
+		}
+		return entries
+	}
+	priority := 0
+	for _, entry := range entries {
+		if entry.Group == group {
+			return entries
+		}
+		priority = max(priority, entry.Priority+1)
+	}
+	return append(slices.Clone(entries), model.TokenGroupEntry{Group: group, Priority: priority})
+}