59 lines
1.7 KiB
Go
59 lines
1.7 KiB
Go
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})
|
|
}
|