Skip to content

phsl / new-api

Go to a project
Commit b26333b2 by Seefs Committed by GitHub
Options

Merge pull request #2121 from QuantumNous/feat/special_group 

feat: add special user usable group setting
parents e7012e7e 6aec0886
Showing with 326 additions and 114 deletions
controller/group.go
...@@ -4,6 +4,7 @@ import ( ...@@ -4,6 +4,7 @@ import (
"net/http" "net/http"
"github.com/QuantumNous/new-api/model" "github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/setting" "github.com/QuantumNous/new-api/setting"
"github.com/QuantumNous/new-api/setting/ratio_setting" "github.com/QuantumNous/new-api/setting/ratio_setting"
...@@ -27,9 +28,9 @@ func GetUserGroups(c *gin.Context) { ...@@ -27,9 +28,9 @@ func GetUserGroups(c *gin.Context) {
userGroup := "" userGroup := ""
userId := c.GetInt("id") userId := c.GetInt("id")
userGroup, _ = model.GetUserGroup(userId, false) userGroup, _ = model.GetUserGroup(userId, false)
userUsableGroups := service.GetUserUsableGroups(userGroup)
for groupName, ratio := range ratio_setting.GetGroupRatioCopy() { for groupName, ratio := range ratio_setting.GetGroupRatioCopy() {
// UserUsableGroups contains the groups that the user can use // UserUsableGroups contains the groups that the user can use
userUsableGroups := setting.GetUserUsableGroups(userGroup)
if desc, ok := userUsableGroups[groupName]; ok { if desc, ok := userUsableGroups[groupName]; ok {
usableGroups[groupName] = map[string]interface{}{ usableGroups[groupName] = map[string]interface{}{
"ratio": ratio, "ratio": ratio,
...@@ -37,7 +38,7 @@ func GetUserGroups(c *gin.Context) { ...@@ -37,7 +38,7 @@ func GetUserGroups(c *gin.Context) {
} }
} }
} }
if setting.GroupInUserUsableGroups("auto") { if _, ok := userUsableGroups["auto"]; ok {
usableGroups["auto"] = map[string]interface{}{ usableGroups["auto"] = map[string]interface{}{
"ratio": "自动", "ratio": "自动",
"desc": setting.GetUsableGroupDescription("auto"), "desc": setting.GetUsableGroupDescription("auto"),
......
controller/model.go
...@@ -15,7 +15,7 @@ import ( ...@@ -15,7 +15,7 @@ import (
"github.com/QuantumNous/new-api/relay/channel/minimax" "github.com/QuantumNous/new-api/relay/channel/minimax"
"github.com/QuantumNous/new-api/relay/channel/moonshot" "github.com/QuantumNous/new-api/relay/channel/moonshot"
relaycommon "github.com/QuantumNous/new-api/relay/common" relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/setting" "github.com/QuantumNous/new-api/service"
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
"github.com/samber/lo" "github.com/samber/lo"
) )
...@@ -149,7 +149,7 @@ func ListModels(c *gin.Context, modelType int) { ...@@ -149,7 +149,7 @@ func ListModels(c *gin.Context, modelType int) {
} }
var models []string var models []string
if tokenGroup == "auto" { if tokenGroup == "auto" {
for _, autoGroup := range setting.AutoGroups { for _, autoGroup := range service.GetUserAutoGroup(userGroup) {
groupModels := model.GetGroupEnabledModels(autoGroup) groupModels := model.GetGroupEnabledModels(autoGroup)
for _, g := range groupModels { for _, g := range groupModels {
if !common.StringsContains(models, g) { if !common.StringsContains(models, g) {
......
controller/pricing.go
...@@ -2,7 +2,7 @@ package controller ...@@ -2,7 +2,7 @@ package controller
import ( import (
"github.com/QuantumNous/new-api/model" "github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/setting" "github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/setting/ratio_setting" "github.com/QuantumNous/new-api/setting/ratio_setting"
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
...@@ -30,7 +30,7 @@ func GetPricing(c *gin.Context) { ...@@ -30,7 +30,7 @@ func GetPricing(c *gin.Context) {
} }
} }
usableGroup = setting.GetUserUsableGroups(group) usableGroup = service.GetUserUsableGroups(group)
// check groupRatio contains usableGroup // check groupRatio contains usableGroup
for group := range ratio_setting.GetGroupRatioCopy() { for group := range ratio_setting.GetGroupRatioCopy() {
if _, ok := usableGroup[group]; !ok { if _, ok := usableGroup[group]; !ok {
...@@ -45,7 +45,7 @@ func GetPricing(c *gin.Context) { ...@@ -45,7 +45,7 @@ func GetPricing(c *gin.Context) {
"group_ratio": groupRatio, "group_ratio": groupRatio,
"usable_group": usableGroup, "usable_group": usableGroup,
"supported_endpoint": model.GetSupportedEndpointMap(), "supported_endpoint": model.GetSupportedEndpointMap(),
"auto_groups": setting.AutoGroups, "auto_groups": service.GetUserAutoGroup(group),
}) })
} }
......
controller/relay.go
...@@ -224,7 +224,7 @@ func getChannel(c *gin.Context, group, originalModel string, retryCount int) (*m ...@@ -224,7 +224,7 @@ func getChannel(c *gin.Context, group, originalModel string, retryCount int) (*m
AutoBan: &autoBanInt, AutoBan: &autoBanInt,
}, nil }, nil
} }
channel, selectGroup, err := model.CacheGetRandomSatisfiedChannel(c, group, originalModel, retryCount) channel, selectGroup, err := service.CacheGetRandomSatisfiedChannel(c, group, originalModel, retryCount)
if err != nil { if err != nil {
return nil, types.NewError(fmt.Errorf("获取分组 %s 下模型 %s 的可用渠道失败（retry）: %s", selectGroup, originalModel, err.Error()), types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry()) return nil, types.NewError(fmt.Errorf("获取分组 %s 下模型 %s 的可用渠道失败（retry）: %s", selectGroup, originalModel, err.Error()), types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry())
} }
......
controller/user.go
...@@ -13,6 +13,7 @@ import ( ...@@ -13,6 +13,7 @@ import (
"github.com/QuantumNous/new-api/dto" "github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/logger" "github.com/QuantumNous/new-api/logger"
"github.com/QuantumNous/new-api/model" "github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/setting" "github.com/QuantumNous/new-api/setting"
"github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/constant"
...@@ -579,7 +580,7 @@ func GetUserModels(c *gin.Context) { ...@@ -579,7 +580,7 @@ func GetUserModels(c *gin.Context) {
common.ApiError(c, err) common.ApiError(c, err)
return return
} }
groups := setting.GetUserUsableGroups(user.Group) groups := service.GetUserUsableGroups(user.Group)
var models []string var models []string
for group := range groups { for group := range groups {
for _, g := range model.GetGroupEnabledModels(group) { for _, g := range model.GetGroupEnabledModels(group) {
......
logger/logger.go
...@@ -66,7 +66,8 @@ func LogError(ctx context.Context, msg string) { ...@@ -66,7 +66,8 @@ func LogError(ctx context.Context, msg string) {
logHelper(ctx, loggerError, msg) logHelper(ctx, loggerError, msg)
} }
func LogDebug(ctx context.Context, msg string) { func LogDebug(ctx context.Context, msg string, args ...any) {
msg = fmt.Sprintf(msg, args...)
if common.DebugEnabled { if common.DebugEnabled {
logHelper(ctx, loggerDebug, msg) logHelper(ctx, loggerDebug, msg)
} }
......
middleware/auth.go
...@@ -9,7 +9,7 @@ import ( ...@@ -9,7 +9,7 @@ import (
"github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/model" "github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/setting" "github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/setting/ratio_setting" "github.com/QuantumNous/new-api/setting/ratio_setting"
"github.com/gin-contrib/sessions" "github.com/gin-contrib/sessions"
...@@ -266,8 +266,8 @@ func TokenAuth() func(c *gin.Context) { ...@@ -266,8 +266,8 @@ func TokenAuth() func(c *gin.Context) {
tokenGroup := token.Group tokenGroup := token.Group
if tokenGroup != "" { if tokenGroup != "" {
// check common.UserUsableGroups[userGroup] // check common.UserUsableGroups[userGroup]
if _, ok := setting.GetUserUsableGroups(userGroup)[tokenGroup]; !ok { if _, ok := service.GetUserUsableGroups(userGroup)[tokenGroup]; !ok {
abortWithOpenAiMessage(c, http.StatusForbidden, fmt.Sprintf("令牌分组 %s 已被禁用", tokenGroup)) abortWithOpenAiMessage(c, http.StatusForbidden, fmt.Sprintf("无权访问 %s 分组", tokenGroup))
return return
} }
// check group in common.GroupRatio // check group in common.GroupRatio
......
middleware/distributor.go
...@@ -15,7 +15,6 @@ import ( ...@@ -15,7 +15,6 @@ import (
"github.com/QuantumNous/new-api/model" "github.com/QuantumNous/new-api/model"
relayconstant "github.com/QuantumNous/new-api/relay/constant" relayconstant "github.com/QuantumNous/new-api/relay/constant"
"github.com/QuantumNous/new-api/service" "github.com/QuantumNous/new-api/service"
"github.com/QuantumNous/new-api/setting"
"github.com/QuantumNous/new-api/setting/ratio_setting" "github.com/QuantumNous/new-api/setting/ratio_setting"
"github.com/QuantumNous/new-api/types" "github.com/QuantumNous/new-api/types"
...@@ -80,7 +79,7 @@ func Distribute() func(c *gin.Context) { ...@@ -80,7 +79,7 @@ func Distribute() func(c *gin.Context) {
return return
} }
var selectGroup string var selectGroup string
userGroup := common.GetContextKeyString(c, constant.ContextKeyUsingGroup) usingGroup := common.GetContextKeyString(c, constant.ContextKeyUsingGroup)
// check path is /pg/chat/completions // check path is /pg/chat/completions
if strings.HasPrefix(c.Request.URL.Path, "/pg/chat/completions") { if strings.HasPrefix(c.Request.URL.Path, "/pg/chat/completions") {
playgroundRequest := &dto.PlayGroundRequest{} playgroundRequest := &dto.PlayGroundRequest{}
...@@ -90,17 +89,17 @@ func Distribute() func(c *gin.Context) { ...@@ -90,17 +89,17 @@ func Distribute() func(c *gin.Context) {
return return
} }
if playgroundRequest.Group != "" { if playgroundRequest.Group != "" {
if !setting.GroupInUserUsableGroups(playgroundRequest.Group) && playgroundRequest.Group != userGroup { if !service.GroupInUserUsableGroups(usingGroup, playgroundRequest.Group) && playgroundRequest.Group != usingGroup {
abortWithOpenAiMessage(c, http.StatusForbidden, "无权访问该分组") abortWithOpenAiMessage(c, http.StatusForbidden, "无权访问该分组")
return return
} }
userGroup = playgroundRequest.Group usingGroup = playgroundRequest.Group
} }
} }
channel, selectGroup, err = model.CacheGetRandomSatisfiedChannel(c, userGroup, modelRequest.Model, 0) channel, selectGroup, err = service.CacheGetRandomSatisfiedChannel(c, usingGroup, modelRequest.Model, 0)
if err != nil { if err != nil {
showGroup := userGroup showGroup := usingGroup
if userGroup == "auto" { if usingGroup == "auto" {
showGroup = fmt.Sprintf("auto(%s)", selectGroup) showGroup = fmt.Sprintf("auto(%s)", selectGroup)
} }
message := fmt.Sprintf("获取分组 %s 下模型 %s 的可用渠道失败（distributor）: %s", showGroup, modelRequest.Model, err.Error()) message := fmt.Sprintf("获取分组 %s 下模型 %s 的可用渠道失败（distributor）: %s", showGroup, modelRequest.Model, err.Error())
...@@ -113,7 +112,7 @@ func Distribute() func(c *gin.Context) { ...@@ -113,7 +112,7 @@ func Distribute() func(c *gin.Context) {
return return
} }
if channel == nil { if channel == nil {
abortWithOpenAiMessage(c, http.StatusServiceUnavailable, fmt.Sprintf("分组 %s 下模型 %s 无可用渠道（distributor）", userGroup, modelRequest.Model), string(types.ErrorCodeModelNotFound)) abortWithOpenAiMessage(c, http.StatusServiceUnavailable, fmt.Sprintf("分组 %s 下模型 %s 无可用渠道（distributor）", usingGroup, modelRequest.Model), string(types.ErrorCodeModelNotFound))
return return
} }
} }
......
model/ability.go
...@@ -103,7 +103,7 @@ func getChannelQuery(group string, model string, retry int) (*gorm.DB, error) { ...@@ -103,7 +103,7 @@ func getChannelQuery(group string, model string, retry int) (*gorm.DB, error) {
return channelQuery, nil return channelQuery, nil
} }
func GetRandomSatisfiedChannel(group string, model string, retry int) (*Channel, error) { func GetChannel(group string, model string, retry int) (*Channel, error) {
var abilities []Ability var abilities []Ability
var err error = nil var err error = nil
......
model/channel_cache.go
...@@ -11,10 +11,7 @@ import ( ...@@ -11,10 +11,7 @@ import (
"github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/setting"
"github.com/QuantumNous/new-api/setting/ratio_setting" "github.com/QuantumNous/new-api/setting/ratio_setting"
"github.com/gin-gonic/gin"
) )
var group2model2channels map[string]map[string][]int // enabled channel var group2model2channels map[string]map[string][]int // enabled channel
...@@ -96,43 +93,10 @@ func SyncChannelCache(frequency int) { ...@@ -96,43 +93,10 @@ func SyncChannelCache(frequency int) {
} }
} }
func CacheGetRandomSatisfiedChannel(c *gin.Context, group string, model string, retry int) (*Channel, string, error) { func GetRandomSatisfiedChannel(group string, model string, retry int) (*Channel, error) {
var channel *Channel
var err error
selectGroup := group
if group == "auto" {
if len(setting.AutoGroups) == 0 {
return nil, selectGroup, errors.New("auto groups is not enabled")
}
for _, autoGroup := range setting.AutoGroups {
if common.DebugEnabled {
println("autoGroup:", autoGroup)
}
channel, _ = getRandomSatisfiedChannel(autoGroup, model, retry)
if channel == nil {
continue
} else {
c.Set("auto_group", autoGroup)
selectGroup = autoGroup
if common.DebugEnabled {
println("selectGroup:", selectGroup)
}
break
}
}
} else {
channel, err = getRandomSatisfiedChannel(group, model, retry)
if err != nil {
return nil, group, err
}
}
return channel, selectGroup, nil
}
func getRandomSatisfiedChannel(group string, model string, retry int) (*Channel, error) {
// if memory cache is disabled, get channel directly from database // if memory cache is disabled, get channel directly from database
if !common.MemoryCacheEnabled { if !common.MemoryCacheEnabled {
return GetRandomSatisfiedChannel(group, model, retry) return GetChannel(group, model, retry)
} }
channelSyncLock.RLock() channelSyncLock.RLock()
...@@ -178,30 +142,36 @@ func getRandomSatisfiedChannel(group string, model string, retry int) (*Channel, ...@@ -178,30 +142,36 @@ func getRandomSatisfiedChannel(group string, model string, retry int) (*Channel,
targetPriority := int64(sortedUniquePriorities[retry]) targetPriority := int64(sortedUniquePriorities[retry])
// get the priority for the given retry number // get the priority for the given retry number
var shouldSmooth = false
var sumWeight = 0
var targetChannels []*Channel var targetChannels []*Channel
for _, channelId := range channels { for _, channelId := range channels {
if channel, ok := channelsIDM[channelId]; ok { if channel, ok := channelsIDM[channelId]; ok {
if channel.GetPriority() == targetPriority { if channel.GetPriority() == targetPriority {
sumWeight += channel.GetWeight()
targetChannels = append(targetChannels, channel) targetChannels = append(targetChannels, channel)
} }
} else { } else {
return nil, fmt.Errorf("数据库一致性错误，渠道# %d 不存在，请联系管理员修复", channelId) return nil, fmt.Errorf("数据库一致性错误，渠道# %d 不存在，请联系管理员修复", channelId)
} }
} }
if sumWeight/len(targetChannels) < 10 {
shouldSmooth = true
}
// 平滑系数 // 平滑系数
smoothingFactor := 10 smoothingFactor := 1
// Calculate the total weight of all channels up to endIdx if shouldSmooth {
totalWeight := 0 smoothingFactor = 100
for _, channel := range targetChannels {
totalWeight += channel.GetWeight() + smoothingFactor
} }
// Calculate the total weight of all channels up to endIdx
totalWeight := sumWeight * smoothingFactor
// Generate a random value in the range [0, totalWeight) // Generate a random value in the range [0, totalWeight)
randomWeight := rand.Intn(totalWeight) randomWeight := rand.Intn(totalWeight)
// Find a channel based on its weight // Find a channel based on its weight
for _, channel := range targetChannels { for _, channel := range targetChannels {
randomWeight -= channel.GetWeight() + smoothingFactor randomWeight -= channel.GetWeight() * smoothingFactor
if randomWeight < 0 { if randomWeight < 0 {
return channel, nil return channel, nil
} }
......
relay/helper/price.go
...@@ -4,6 +4,7 @@ import ( ...@@ -4,6 +4,7 @@ import (
"fmt" "fmt"
"github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/logger"
relaycommon "github.com/QuantumNous/new-api/relay/common" relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/setting/operation_setting" "github.com/QuantumNous/new-api/setting/operation_setting"
"github.com/QuantumNous/new-api/setting/ratio_setting" "github.com/QuantumNous/new-api/setting/ratio_setting"
...@@ -22,9 +23,7 @@ func HandleGroupRatio(ctx *gin.Context, relayInfo *relaycommon.RelayInfo) types. ...@@ -22,9 +23,7 @@ func HandleGroupRatio(ctx *gin.Context, relayInfo *relaycommon.RelayInfo) types.
// check auto group // check auto group
autoGroup, exists := ctx.Get("auto_group") autoGroup, exists := ctx.Get("auto_group")
if exists { if exists {
if common.DebugEnabled { logger.LogDebug(ctx, fmt.Sprintf("final group: %s", autoGroup))
println(fmt.Sprintf("final group: %s", autoGroup))
}
relayInfo.UsingGroup = autoGroup.(string) relayInfo.UsingGroup = autoGroup.(string)
} }
......
service/channel_select.go 0 → 100644
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 → 100644
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
}
setting/auto_group.go
package setting package setting
import "encoding/json" import (
"github.com/QuantumNous/new-api/common"
)
var AutoGroups = []string{ var autoGroups = []string{
"default", "default",
} }
var DefaultUseAutoGroup = false var DefaultUseAutoGroup = false
func ContainsAutoGroup(group string) bool { func ContainsAutoGroup(group string) bool {
for _, autoGroup := range AutoGroups { for _, autoGroup := range autoGroups {
if autoGroup == group { if autoGroup == group {
return true return true
} }
...@@ -18,14 +20,18 @@ func ContainsAutoGroup(group string) bool { ...@@ -18,14 +20,18 @@ func ContainsAutoGroup(group string) bool {
} }
func UpdateAutoGroupsByJsonString(jsonString string) error { func UpdateAutoGroupsByJsonString(jsonString string) error {
AutoGroups = make([]string, 0) autoGroups = make([]string, 0)
return json.Unmarshal([]byte(jsonString), &AutoGroups) return common.Unmarshal([]byte(jsonString), &autoGroups)
} }
func AutoGroups2JsonString() string { func AutoGroups2JsonString() string {
jsonBytes, err := json.Marshal(AutoGroups) jsonBytes, err := common.Marshal(autoGroups)
if err != nil { if err != nil {
return "[]" return "[]"
} }
return string(jsonBytes) return string(jsonBytes)
} }
func GetAutoGroups() []string {
return autoGroups
}
setting/config/config.go
...@@ -131,6 +131,18 @@ func configToMap(config interface{}) (map[string]string, error) { ...@@ -131,6 +131,18 @@ func configToMap(config interface{}) (map[string]string, error) {
strValue = strconv.FormatUint(field.Uint(), 10) strValue = strconv.FormatUint(field.Uint(), 10)
case reflect.Float32, reflect.Float64: case reflect.Float32, reflect.Float64:
strValue = strconv.FormatFloat(field.Float(), 'f', -1, 64) strValue = strconv.FormatFloat(field.Float(), 'f', -1, 64)
case reflect.Ptr:
// 处理指针类型：如果非 nil，序列化指向的值
if !field.IsNil() {
bytes, err := json.Marshal(field.Interface())
if err != nil {
return nil, err
}
strValue = string(bytes)
} else {
// nil 指针序列化为 "null"
strValue = "null"
}
case reflect.Map, reflect.Slice, reflect.Struct: case reflect.Map, reflect.Slice, reflect.Struct:
// 复杂类型使用JSON序列化 // 复杂类型使用JSON序列化
bytes, err := json.Marshal(field.Interface()) bytes, err := json.Marshal(field.Interface())
...@@ -215,6 +227,21 @@ func updateConfigFromMap(config interface{}, configMap map[string]string) error ...@@ -215,6 +227,21 @@ func updateConfigFromMap(config interface{}, configMap map[string]string) error
continue continue
} }
field.SetFloat(floatValue) field.SetFloat(floatValue)
case reflect.Ptr:
// 处理指针类型
if strValue == "null" {
field.Set(reflect.Zero(field.Type()))
} else {
// 如果指针是 nil，需要先初始化
if field.IsNil() {
field.Set(reflect.New(field.Type().Elem()))
}
// 反序列化到指针指向的值
err := json.Unmarshal([]byte(strValue), field.Interface())
if err != nil {
continue
}
}
case reflect.Map, reflect.Slice, reflect.Struct: case reflect.Map, reflect.Slice, reflect.Struct:
// 复杂类型使用JSON反序列化 // 复杂类型使用JSON反序列化
err := json.Unmarshal([]byte(strValue), field.Addr().Interface()) err := json.Unmarshal([]byte(strValue), field.Addr().Interface())
......
setting/ratio_setting/group_ratio.go
...@@ -6,6 +6,8 @@ import ( ...@@ -6,6 +6,8 @@ import (
"sync" "sync"
"github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/setting/config"
"github.com/QuantumNous/new-api/types"
) )
var groupRatio = map[string]float64{ var groupRatio = map[string]float64{
...@@ -13,6 +15,7 @@ var groupRatio = map[string]float64{ ...@@ -13,6 +15,7 @@ var groupRatio = map[string]float64{
"vip": 1, "vip": 1,
"svip": 1, "svip": 1,
} }
var groupRatioMutex sync.RWMutex var groupRatioMutex sync.RWMutex
var ( var (
...@@ -24,6 +27,42 @@ var ( ...@@ -24,6 +27,42 @@ var (
groupGroupRatioMutex sync.RWMutex groupGroupRatioMutex sync.RWMutex
) )
var defaultGroupSpecialUsableGroup = map[string]map[string]string{
"vip": {
"append_1": "vip_special_group_1",
"-:remove_1": "vip_removed_group_1",
},
}
type GroupRatioSetting struct {
GroupRatio map[string]float64 `json:"group_ratio"`
GroupGroupRatio map[string]map[string]float64 `json:"group_group_ratio"`
GroupSpecialUsableGroup *types.RWMap[string, map[string]string] `json:"group_special_usable_group"`
}
var groupRatioSetting GroupRatioSetting
func init() {
groupSpecialUsableGroup := types.NewRWMap[string, map[string]string]()
groupSpecialUsableGroup.AddAll(defaultGroupSpecialUsableGroup)
groupRatioSetting = GroupRatioSetting{
GroupSpecialUsableGroup: groupSpecialUsableGroup,
GroupRatio: groupRatio,
GroupGroupRatio: GroupGroupRatio,
}
config.GlobalConfig.Register("group_ratio_setting", &groupRatioSetting)
}
func GetGroupRatioSetting() *GroupRatioSetting {
if groupRatioSetting.GroupSpecialUsableGroup == nil {
groupRatioSetting.GroupSpecialUsableGroup = types.NewRWMap[string, map[string]string]()
groupRatioSetting.GroupSpecialUsableGroup.AddAll(defaultGroupSpecialUsableGroup)
}
return &groupRatioSetting
}
func GetGroupRatioCopy() map[string]float64 { func GetGroupRatioCopy() map[string]float64 {
groupRatioMutex.RLock() groupRatioMutex.RLock()
defer groupRatioMutex.RUnlock() defer groupRatioMutex.RUnlock()
......
setting/user_usable_group.go
...@@ -43,29 +43,6 @@ func UpdateUserUsableGroupsByJSONString(jsonStr string) error { ...@@ -43,29 +43,6 @@ func UpdateUserUsableGroupsByJSONString(jsonStr string) error {
return json.Unmarshal([]byte(jsonStr), &userUsableGroups) return json.Unmarshal([]byte(jsonStr), &userUsableGroups)
} }
func GetUserUsableGroups(userGroup string) map[string]string {
groupsCopy := GetUserUsableGroupsCopy()
if userGroup == "" {
if _, ok := groupsCopy["default"]; !ok {
groupsCopy["default"] = "default"
}
}
// 如果userGroup不在UserUsableGroups中，返回UserUsableGroups + userGroup
if _, ok := groupsCopy[userGroup]; !ok {
groupsCopy[userGroup] = "用户分组"
}
// 如果userGroup在UserUsableGroups中，返回UserUsableGroups
return groupsCopy
}
func GroupInUserUsableGroups(groupName string) bool {
userUsableGroupsMutex.RLock()
defer userUsableGroupsMutex.RUnlock()
_, ok := userUsableGroups[groupName]
return ok
}
func GetUsableGroupDescription(groupName string) string { func GetUsableGroupDescription(groupName string) string {
userUsableGroupsMutex.RLock() userUsableGroupsMutex.RLock()
defer userUsableGroupsMutex.RUnlock() defer userUsableGroupsMutex.RUnlock()
......
types/rw_map.go 0 → 100644
package types
import (
"sync"
"github.com/QuantumNous/new-api/common"
)
type RWMap[K comparable, V any] struct {
data map[K]V
mutex sync.RWMutex
}
func (m *RWMap[K, V]) UnmarshalJSON(b []byte) error {
m.mutex.Lock()
defer m.mutex.Unlock()
m.data = make(map[K]V)
return common.Unmarshal(b, &m.data)
}
func (m *RWMap[K, V]) MarshalJSON() ([]byte, error) {
m.mutex.RLock()
defer m.mutex.RUnlock()
return common.Marshal(m.data)
}
func NewRWMap[K comparable, V any]() *RWMap[K, V] {
return &RWMap[K, V]{
data: make(map[K]V),
}
}
func (m *RWMap[K, V]) Get(key K) (V, bool) {
m.mutex.RLock()
defer m.mutex.RUnlock()
value, exists := m.data[key]
return value, exists
}
func (m *RWMap[K, V]) Set(key K, value V) {
m.mutex.Lock()
defer m.mutex.Unlock()
m.data[key] = value
}
func (m *RWMap[K, V]) AddAll(other map[K]V) {
m.mutex.Lock()
defer m.mutex.Unlock()
for k, v := range other {
m.data[k] = v
}
}
func (m *RWMap[K, V]) Clear() {
m.mutex.Lock()
defer m.mutex.Unlock()
m.data = make(map[K]V)
}
// ReadAll returns a copy of the entire map.
func (m *RWMap[K, V]) ReadAll() map[K]V {
m.mutex.RLock()
defer m.mutex.RUnlock()
copiedMap := make(map[K]V)
for k, v := range m.data {
copiedMap[k] = v
}
return copiedMap
}
func (m *RWMap[K, V]) Len() int {
m.mutex.RLock()
defer m.mutex.RUnlock()
return len(m.data)
}
func LoadFromJsonString[K comparable, V any](m *RWMap[K, V], jsonStr string) error {
m.mutex.Lock()
defer m.mutex.Unlock()
m.data = make(map[K]V)
return common.Unmarshal([]byte(jsonStr), &m.data)
}
web/src/components/settings/RatioSetting.jsx
...@@ -46,6 +46,7 @@ const RatioSetting = () => { ...@@ -46,6 +46,7 @@ const RatioSetting = () => {
DefaultUseAutoGroup: false, DefaultUseAutoGroup: false,
ExposeRatioEnabled: false, ExposeRatioEnabled: false,
UserUsableGroups: '', UserUsableGroups: '',
'group_ratio_setting.group_special_usable_group': '',
}); });
const [loading, setLoading] = useState(false); const [loading, setLoading] = useState(false);
...@@ -57,17 +58,7 @@ const RatioSetting = () => { ...@@ -57,17 +58,7 @@ const RatioSetting = () => {
let newInputs = {}; let newInputs = {};
data.forEach((item) => { data.forEach((item) => {
if ( if (
item.key === 'ModelRatio' || item.value.startsWith('{') || item.value.startsWith('[')
item.key === 'GroupRatio' ||
item.key === 'GroupGroupRatio' ||
item.key === 'AutoGroups' ||
item.key === 'UserUsableGroups' ||
item.key === 'CompletionRatio' ||
item.key === 'ModelPrice' ||
item.key === 'CacheRatio' ||
item.key === 'ImageRatio' ||
item.key === 'AudioRatio' ||
item.key === 'AudioCompletionRatio'
) { ) {
try { try {
item.value = JSON.stringify(JSON.parse(item.value), null, 2); item.value = JSON.stringify(JSON.parse(item.value), null, 2);
......
web/src/components/table/tokens/modals/EditTokenModal.jsx
...@@ -137,14 +137,12 @@ const EditTokenModal = (props) => { ...@@ -137,14 +137,12 @@ const EditTokenModal = (props) => {
if (statusState?.status?.default_use_auto_group) { if (statusState?.status?.default_use_auto_group) {
if (localGroupOptions.some((group) => group.value === 'auto')) { if (localGroupOptions.some((group) => group.value === 'auto')) {
localGroupOptions.sort((a, b) => (a.value === 'auto' ? -1 : 1)); localGroupOptions.sort((a, b) => (a.value === 'auto' ? -1 : 1));
} else {
localGroupOptions.unshift({ label: t('自动选择'), value: 'auto' });
} }
} }
setGroups(localGroupOptions); setGroups(localGroupOptions);
if (statusState?.status?.default_use_auto_group && formApiRef.current) { // if (statusState?.status?.default_use_auto_group && formApiRef.current) {
formApiRef.current.setValue('group', 'auto'); // formApiRef.current.setValue('group', 'auto');
} // }
} else { } else {
showError(t(message)); showError(t(message));
} }
......
web/src/pages/Setting/Ratio/GroupRatioSettings.jsx
...@@ -36,6 +36,7 @@ export default function GroupRatioSettings(props) { ...@@ -36,6 +36,7 @@ export default function GroupRatioSettings(props) {
GroupRatio: '', GroupRatio: '',
UserUsableGroups: '', UserUsableGroups: '',
GroupGroupRatio: '', GroupGroupRatio: '',
'group_ratio_setting.group_special_usable_group': '',
AutoGroups: '', AutoGroups: '',
DefaultUseAutoGroup: false, DefaultUseAutoGroup: false,
}); });
...@@ -188,6 +189,30 @@ export default function GroupRatioSettings(props) { ...@@ -188,6 +189,30 @@ export default function GroupRatioSettings(props) {
<Row gutter={16}> <Row gutter={16}>
<Col xs={24} sm={16}> <Col xs={24} sm={16}>
<Form.TextArea <Form.TextArea
label={t('分组特殊可用分组')}
placeholder={t('为一个 JSON 文本')}
extraText={t(
'键为用户分组名称，值为操作映射对象。内层键以"+:"开头表示添加指定分组（键值为分组名称，值为描述），以"-:"开头表示移除指定分组（键值为分组名称），不带前缀的键直接添加该分组。例如：{"vip": {"+:premium": "高级分组", "special": "特殊分组", "-:default": "默认分组"}}，表示 vip 分组的用户可以使用 premium 和 special 分组，同时移除 default 分组的访问权限',
)}
field={'group_ratio_setting.group_special_usable_group'}
autosize={{ minRows: 6, maxRows: 12 }}
trigger='blur'
stopValidateWithError
rules={[
{
validator: (rule, value) => verifyJSON(value),
message: t('不是合法的 JSON 字符串'),
},
]}
onChange={(value) =>
setInputs({ ...inputs, 'group_ratio_setting.group_special_usable_group': value })
}
/>
</Col>
</Row>
<Row gutter={16}>
<Col xs={24} sm={16}>
<Form.TextArea
label={t('自动分组auto，从第一个开始选择')} label={t('自动分组auto，从第一个开始选择')}
placeholder={t('为一个 JSON 文本')} placeholder={t('为一个 JSON 文本')}
field={'AutoGroups'} field={'AutoGroups'}
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or sign in to comment