feat(advanced-custom): remove fallback option and enhance path matching for advanced custom routes
This commit is contained in:
+50
-1
@@ -7,6 +7,8 @@ import (
|
||||
"sync"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/dto"
|
||||
|
||||
"github.com/samber/lo"
|
||||
"gorm.io/gorm"
|
||||
@@ -103,7 +105,7 @@ func getChannelQuery(group string, model string, retry int) (*gorm.DB, error) {
|
||||
return channelQuery, nil
|
||||
}
|
||||
|
||||
func GetChannel(group string, model string, retry int) (*Channel, error) {
|
||||
func GetChannel(group string, model string, retry int, requestPath string) (*Channel, error) {
|
||||
var abilities []Ability
|
||||
|
||||
var err error = nil
|
||||
@@ -119,6 +121,7 @@ func GetChannel(group string, model string, retry int) (*Channel, error) {
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
abilities = filterAbilitiesByRequestPath(abilities, requestPath)
|
||||
channel := Channel{}
|
||||
if len(abilities) > 0 {
|
||||
// Randomly choose one
|
||||
@@ -143,6 +146,52 @@ func GetChannel(group string, model string, retry int) (*Channel, error) {
|
||||
return &channel, err
|
||||
}
|
||||
|
||||
// filterAbilitiesByRequestPath restricts candidates by request path for the DB
|
||||
// (non-memory-cache) selection path. Only Advanced Custom (type 58) channels are
|
||||
// path-checked: kept only when one of their routes matches requestPath; all other
|
||||
// channel types always pass. When requestPath is empty, filtering is skipped.
|
||||
func filterAbilitiesByRequestPath(abilities []Ability, requestPath string) []Ability {
|
||||
if requestPath == "" || len(abilities) == 0 {
|
||||
return abilities
|
||||
}
|
||||
|
||||
channelIds := make([]int, 0, len(abilities))
|
||||
seen := make(map[int]struct{}, len(abilities))
|
||||
for _, ability := range abilities {
|
||||
if _, ok := seen[ability.ChannelId]; ok {
|
||||
continue
|
||||
}
|
||||
seen[ability.ChannelId] = struct{}{}
|
||||
channelIds = append(channelIds, ability.ChannelId)
|
||||
}
|
||||
|
||||
var channels []*Channel
|
||||
if err := DB.Where("id IN ?", channelIds).Find(&channels).Error; err != nil {
|
||||
// On error, fall back to unfiltered candidates to avoid blocking selection
|
||||
return abilities
|
||||
}
|
||||
|
||||
advancedConfigs := make(map[int]*dto.AdvancedCustomConfig)
|
||||
for _, channel := range channels {
|
||||
if channel.Type == constant.ChannelTypeAdvancedCustom {
|
||||
advancedConfigs[channel.Id] = channel.GetOtherSettings().AdvancedCustom
|
||||
}
|
||||
}
|
||||
|
||||
filtered := make([]Ability, 0, len(abilities))
|
||||
for _, ability := range abilities {
|
||||
config, isAdvancedCustom := advancedConfigs[ability.ChannelId]
|
||||
if !isAdvancedCustom {
|
||||
filtered = append(filtered, ability)
|
||||
continue
|
||||
}
|
||||
if config != nil && config.SupportsPath(requestPath) {
|
||||
filtered = append(filtered, ability)
|
||||
}
|
||||
}
|
||||
return filtered
|
||||
}
|
||||
|
||||
func (channel *Channel) AddAbilities(tx *gorm.DB) error {
|
||||
models_ := strings.Split(channel.Models, ",")
|
||||
groups_ := strings.Split(channel.Group, ",")
|
||||
|
||||
@@ -956,9 +956,6 @@ func (channel *Channel) ValidateSettings() error {
|
||||
if channelOtherSettings.AdvancedCustom == nil {
|
||||
return fmt.Errorf("advanced_custom is required")
|
||||
}
|
||||
if channelOtherSettings.AdvancedCustom.Fallback.Enabled && (channel.BaseURL == nil || strings.TrimSpace(*channel.BaseURL) == "") {
|
||||
return fmt.Errorf("base_url is required when advanced_custom advanced_fallback is enabled")
|
||||
}
|
||||
}
|
||||
if channelOtherSettings.AdvancedCustom != nil {
|
||||
if err := channelOtherSettings.AdvancedCustom.Validate(); err != nil {
|
||||
|
||||
+43
-4
@@ -11,12 +11,16 @@ import (
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/dto"
|
||||
"github.com/QuantumNous/new-api/logger"
|
||||
"github.com/QuantumNous/new-api/setting/ratio_setting"
|
||||
)
|
||||
|
||||
var group2model2channels map[string]map[string][]int // enabled channel
|
||||
var channelsIDM map[int]*Channel // all channels include disabled
|
||||
// channel2advancedCustomConfig caches parsed Advanced Custom (type 58) configs so
|
||||
// path-aware selection avoids re-parsing JSON per request. Refreshed on full sync.
|
||||
var channel2advancedCustomConfig map[int]*dto.AdvancedCustomConfig
|
||||
var channelSyncLock sync.RWMutex
|
||||
|
||||
func InitChannelCache() {
|
||||
@@ -24,10 +28,16 @@ func InitChannelCache() {
|
||||
return
|
||||
}
|
||||
newChannelId2channel := make(map[int]*Channel)
|
||||
newChannel2advancedCustomConfig := make(map[int]*dto.AdvancedCustomConfig)
|
||||
var channels []*Channel
|
||||
DB.Find(&channels)
|
||||
for _, channel := range channels {
|
||||
newChannelId2channel[channel.Id] = channel
|
||||
if channel.Type == constant.ChannelTypeAdvancedCustom {
|
||||
if config := channel.GetOtherSettings().AdvancedCustom; config != nil {
|
||||
newChannel2advancedCustomConfig[channel.Id] = config
|
||||
}
|
||||
}
|
||||
}
|
||||
var abilities []*Ability
|
||||
DB.Find(&abilities)
|
||||
@@ -82,6 +92,7 @@ func InitChannelCache() {
|
||||
}
|
||||
}
|
||||
channelsIDM = newChannelId2channel
|
||||
channel2advancedCustomConfig = newChannel2advancedCustomConfig
|
||||
channelSyncLock.Unlock()
|
||||
common.SysLog("channels synced from database")
|
||||
}
|
||||
@@ -94,22 +105,22 @@ func SyncChannelCache(frequency int) {
|
||||
}
|
||||
}
|
||||
|
||||
func GetRandomSatisfiedChannel(group string, model string, retry int) (*Channel, error) {
|
||||
func GetRandomSatisfiedChannel(group string, model string, retry int, requestPath string) (*Channel, error) {
|
||||
// if memory cache is disabled, get channel directly from database
|
||||
if !common.MemoryCacheEnabled {
|
||||
return GetChannel(group, model, retry)
|
||||
return GetChannel(group, model, retry, requestPath)
|
||||
}
|
||||
|
||||
channelSyncLock.RLock()
|
||||
defer channelSyncLock.RUnlock()
|
||||
|
||||
// First, try to find channels with the exact model name.
|
||||
channels := group2model2channels[group][model]
|
||||
channels := filterChannelsByRequestPath(group2model2channels[group][model], requestPath)
|
||||
|
||||
// If no channels found, try to find channels with the normalized model name.
|
||||
if len(channels) == 0 {
|
||||
normalizedModel := ratio_setting.FormatMatchingModelName(model)
|
||||
channels = group2model2channels[group][normalizedModel]
|
||||
channels = filterChannelsByRequestPath(group2model2channels[group][normalizedModel], requestPath)
|
||||
}
|
||||
|
||||
if len(channels) == 0 {
|
||||
@@ -191,6 +202,34 @@ func GetRandomSatisfiedChannel(group string, model string, retry int) (*Channel,
|
||||
return nil, errors.New("channel not found")
|
||||
}
|
||||
|
||||
// filterChannelsByRequestPath restricts candidates by request path. Only Advanced
|
||||
// Custom (type 58) channels are path-checked: they are kept only when one of their
|
||||
// configured routes matches requestPath. All other channel types always pass.
|
||||
// When requestPath is empty (non-relay callers) filtering is skipped.
|
||||
// Caller must hold channelSyncLock (read lock). The cached slice is never mutated.
|
||||
func filterChannelsByRequestPath(channels []int, requestPath string) []int {
|
||||
if requestPath == "" || len(channels) == 0 {
|
||||
return channels
|
||||
}
|
||||
filtered := make([]int, 0, len(channels))
|
||||
for _, channelId := range channels {
|
||||
channel, ok := channelsIDM[channelId]
|
||||
if !ok {
|
||||
// keep it so the downstream consistency error is raised as before
|
||||
filtered = append(filtered, channelId)
|
||||
continue
|
||||
}
|
||||
if channel.Type != constant.ChannelTypeAdvancedCustom {
|
||||
filtered = append(filtered, channelId)
|
||||
continue
|
||||
}
|
||||
if config := channel2advancedCustomConfig[channelId]; config != nil && config.SupportsPath(requestPath) {
|
||||
filtered = append(filtered, channelId)
|
||||
}
|
||||
}
|
||||
return filtered
|
||||
}
|
||||
|
||||
func CacheGetChannel(id int) (*Channel, error) {
|
||||
if !common.MemoryCacheEnabled {
|
||||
return GetChannelById(id, true)
|
||||
|
||||
Reference in New Issue
Block a user