feat: add special user usable group setting

service/channel_select.go

@@ -0,0 +1,42 @@
+package service
+
+import (
+	"errors"
+
+	"github.com/QuantumNous/new-api/common"
+	"github.com/QuantumNous/new-api/constant"
+	"github.com/QuantumNous/new-api/logger"
+	"github.com/QuantumNous/new-api/model"
+	"github.com/QuantumNous/new-api/setting"
+	"github.com/gin-gonic/gin"
+)
+
+func CacheGetRandomSatisfiedChannel(c *gin.Context, group string, modelName string, retry int) (*model.Channel, string, error) {
+	var channel *model.Channel
+	var err error
+	selectGroup := group
+	userGroup := common.GetContextKeyString(c, constant.ContextKeyUserGroup)
+	if group == "auto" {
+		if len(setting.GetAutoGroups()) == 0 {
+			return nil, selectGroup, errors.New("auto groups is not enabled")
+		}
+		for _, autoGroup := range GetUserAutoGroup(userGroup) {
+			logger.LogDebug(c, "Auto selecting group:", autoGroup)
+			channel, _ = model.GetRandomSatisfiedChannel(autoGroup, modelName, retry)
+			if channel == nil {
+				continue
+			} else {
+				c.Set("auto_group", autoGroup)
+				selectGroup = autoGroup
+				logger.LogDebug(c, "Auto selected group:", autoGroup)
+				break
+			}
+		}
+	} else {
+		channel, err = model.GetRandomSatisfiedChannel(group, modelName, retry)
+		if err != nil {
+			return nil, group, err
+		}
+	}
+	return channel, selectGroup, nil
+}

service/group.go

@@ -0,0 +1,54 @@
+package service
+
+import (
+	"strings"
+
+	"github.com/QuantumNous/new-api/setting"
+	"github.com/QuantumNous/new-api/setting/ratio_setting"
+)
+
+func GetUserUsableGroups(userGroup string) map[string]string {
+	groupsCopy := setting.GetUserUsableGroupsCopy()
+	if userGroup != "" {
+		specialSettings, b := ratio_setting.GetGroupRatioSetting().GroupSpecialUsableGroup.Get(userGroup)
+		if b {
+			// 处理特殊可用分组
+			for specialGroup, desc := range specialSettings {
+				if strings.HasPrefix(specialGroup, "-:") {
+					// 移除分组
+					groupToRemove := strings.TrimPrefix(specialGroup, "-:")
+					delete(groupsCopy, groupToRemove)
+				} else if strings.HasPrefix(specialGroup, "+:") {
+					// 添加分组
+					groupToAdd := strings.TrimPrefix(specialGroup, "+:")
+					groupsCopy[groupToAdd] = desc
+				} else {
+					// 直接添加分组
+					groupsCopy[specialGroup] = desc
+				}
+			}
+		}
+		// 如果userGroup不在UserUsableGroups中，返回UserUsableGroups + userGroup
+		if _, ok := groupsCopy[userGroup]; !ok {
+			groupsCopy[userGroup] = "用户分组"
+		}
+	}
+	return groupsCopy
+}
+
+func GroupInUserUsableGroups(userGroup, groupName string) bool {
+	_, ok := GetUserUsableGroups(userGroup)[groupName]
+	return ok
+}
+
+// GetUserAutoGroup 根据用户分组获取自动分组设置
+func GetUserAutoGroup(userGroup string) []string {
+	groups := GetUserUsableGroups(userGroup)
+	autoGroups := make([]string, 0)
+	for _, group := range setting.GetAutoGroups() {
+		if _, ok := groups[group]; ok {
+			autoGroups = append(autoGroups, group)
+		}
+	}
+	return autoGroups
+}