github/new-api
1
0
Fork 0
mirror of https://github.com/QuantumNous/new-api.git synced 2026-03-30 04:22:58 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #2121 from QuantumNous/feat/special_group

Browse Source 
feat: add special user usable group setting
This commit is contained in:
Seefs
2025-10-29 18:54:51 +09:00
committed by GitHub
parent c4ea095ae0 1ded19795a
commit 1f419a3c71
22 changed files with 521 additions and 142 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
model/ability.go
View File

@@ -103,7 +103,7 @@ func getChannelQuery(group string, model string, retry int) (*gorm.DB, error) {
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 err error = nil

58
model/channel_cache.go
View File

@@ -11,10 +11,7 @@ import (
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/setting"
"github.com/QuantumNous/new-api/setting/ratio_setting"
"github.com/gin-gonic/gin"
)
var group2model2channels map[string]map[string][]int // enabled channel
@@ -96,43 +93,10 @@ func SyncChannelCache(frequency int) {
}
}
func CacheGetRandomSatisfiedChannel(c *gin.Context, group string, model string, retry int) (*Channel, string, 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) {
func GetRandomSatisfiedChannel(group string, model string, retry int) (*Channel, error) {
// if memory cache is disabled, get channel directly from database
if !common.MemoryCacheEnabled {
return GetRandomSatisfiedChannel(group, model, retry)
return GetChannel(group, model, retry)
}
channelSyncLock.RLock()
@@ -178,30 +142,36 @@ func getRandomSatisfiedChannel(group string, model string, retry int) (*Channel,
targetPriority := int64(sortedUniquePriorities[retry])
// get the priority for the given retry number
var shouldSmooth = false
var sumWeight = 0
var targetChannels []*Channel
for _, channelId := range channels {
if channel, ok := channelsIDM[channelId]; ok {
if channel.GetPriority() == targetPriority {
sumWeight += channel.GetWeight()
targetChannels = append(targetChannels, channel)
}
} else {
return nil, fmt.Errorf("数据库一致性错误，渠道# %d 不存在，请联系管理员修复", channelId)
}
}
if sumWeight/len(targetChannels) < 10 {
shouldSmooth = true
}
// 平滑系数
smoothingFactor := 10
// Calculate the total weight of all channels up to endIdx
totalWeight := 0
for _, channel := range targetChannels {
totalWeight += channel.GetWeight() + smoothingFactor
smoothingFactor := 1
if shouldSmooth {
smoothingFactor = 100
}
// Calculate the total weight of all channels up to endIdx
totalWeight := sumWeight * smoothingFactor
// Generate a random value in the range [0, totalWeight)
randomWeight := rand.Intn(totalWeight)
// Find a channel based on its weight
for _, channel := range targetChannels {
randomWeight -= channel.GetWeight() + smoothingFactor
randomWeight -= channel.GetWeight() * smoothingFactor
if randomWeight < 0 {
return channel, nil
}