From f6321b391046881ec044b94fdbb9c9baecedf2e4 Mon Sep 17 00:00:00 2001 From: fengsilin Date: Thu, 9 Jul 2026 16:45:45 +0800 Subject: [PATCH] feat(channel): support dynamic video asset channel selection Co-Authored-By: Codex --- controller/channel_test_tianyiyun_test.go | 11 +- controller/doubao_aiping_video.go | 4 - controller/kling_aiping_native.go | 6 +- controller/kling_aiping_native_test.go | 13 +- controller/relay.go | 49 ++--- middleware/distributor.go | 48 +++-- middleware/doubao_asset_binding_test.go | 41 ++++- model/user_asset_channel.go | 25 ++- model/user_asset_channel_test.go | 69 +++++++ .../task/doubao_tianyiyun/adaptor_test.go | 24 +++ .../channel/task/kling/aiping/adaptor_test.go | 33 ++++ service/channel_select.go | 18 +- service/channel_select_test.go | 102 +++++++++++ service/doubao_asset_channel.go | 132 +------------- service/doubao_asset_channel_test.go | 21 ++- service/kling_asset_channel.go | 92 ++-------- service/kling_asset_channel_test.go | 26 ++- service/video_asset_channel.go | 171 ++++++++++++++++++ service/video_asset_channel_test.go | 107 +++++++++++ 19 files changed, 709 insertions(+), 283 deletions(-) create mode 100644 service/video_asset_channel.go create mode 100644 service/video_asset_channel_test.go diff --git a/controller/channel_test_tianyiyun_test.go b/controller/channel_test_tianyiyun_test.go index 31b9b9d..c6b6c65 100644 --- a/controller/channel_test_tianyiyun_test.go +++ b/controller/channel_test_tianyiyun_test.go @@ -7,6 +7,7 @@ import ( "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/service" "github.com/stretchr/testify/require" ) @@ -24,8 +25,12 @@ func TestTestChannelRejectsAsyncVideoChannels(t *testing.T) { } } -func TestRequiredTaskChannelTypeForTianyiYunSeedanceModel(t *testing.T) { - c := newControllerJSONContext(t, "/api/v3/contents/generations/tasks", `{"model":"cdance2.0-0611"}`) +func TestRequiredTaskChannelTypeForTianyiYunSeedanceModelUsesAllowedFamily(t *testing.T) { + c := newControllerJSONContext(t, "/api/v3/contents/generations/tasks", `{"model":"Doubao-Seedance-2.0"}`) - require.Equal(t, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, requiredTaskChannelTypeForRequest(c)) + require.Equal(t, 0, requiredTaskChannelTypeForRequest(c)) + require.ElementsMatch(t, + service.VideoAssetChannelTypesForFamily(service.VideoAssetFamilySeedance), + allowedTaskChannelTypesForRequest(c), + ) } diff --git a/controller/doubao_aiping_video.go b/controller/doubao_aiping_video.go index 515d238..dae543a 100644 --- a/controller/doubao_aiping_video.go +++ b/controller/doubao_aiping_video.go @@ -10,7 +10,6 @@ import ( "github.com/QuantumNous/new-api/dto" "github.com/QuantumNous/new-api/model" relaycommon "github.com/QuantumNous/new-api/relay/common" - "github.com/QuantumNous/new-api/service" "github.com/QuantumNous/new-api/types" "github.com/gin-gonic/gin" ) @@ -94,9 +93,6 @@ func AipingNativeVideoSubmit(c *gin.Context) { }) return } - channelType := service.DoubaoAssetChannelTypeForModel(relayInfo.OriginModelName) - c.Set("channel_type", channelType) - relayInfo.ChannelType = channelType relayTaskWithInfo(c, relayInfo) } diff --git a/controller/kling_aiping_native.go b/controller/kling_aiping_native.go index ebc62ad..a636c45 100644 --- a/controller/kling_aiping_native.go +++ b/controller/kling_aiping_native.go @@ -54,7 +54,7 @@ func resolveKlingAipingBoundChannel(c *gin.Context, relayInfo *relaycommon.Relay if group == "" { return nil } - channel, err := service.ResolveKlingAssetChannel(c.GetInt("id"), group) + channel, err := service.GetBoundKlingAssetChannelForModel(c.GetInt("id"), group, relayInfo.OriginModelName) if err != nil { logger.LogError(c, fmt.Sprintf("resolve kling aiping bound channel failed: %v", err)) return nil @@ -141,7 +141,7 @@ func KlingAipingNativeProxy(c *gin.Context) { TokenGroup: group, ModelName: route.BillingModel, Retry: common.GetPointer(0), - RequiredChannelType: constant.ChannelTypeKlingAiping, + AllowedChannelTypes: service.VideoAssetChannelTypesForFamily(service.VideoAssetFamilyKling), }) if err != nil { c.JSON(http.StatusServiceUnavailable, gin.H{"code": 503, "message": err.Error()}) @@ -167,7 +167,7 @@ func resolveKlingAipingBoundChannelForProxy(c *gin.Context, group, billingModel if group == "" || group == "auto" || strings.TrimSpace(billingModel) == "" { return nil } - channel, err := service.ResolveKlingAssetChannel(c.GetInt("id"), group) + channel, err := service.GetBoundKlingAssetChannelForModel(c.GetInt("id"), group, billingModel) if err != nil { logger.LogError(c, fmt.Sprintf("resolve kling aiping bound channel (proxy) failed: %v", err)) return nil diff --git a/controller/kling_aiping_native_test.go b/controller/kling_aiping_native_test.go index 9096534..9018d54 100644 --- a/controller/kling_aiping_native_test.go +++ b/controller/kling_aiping_native_test.go @@ -86,15 +86,24 @@ func TestConfigureKlingAipingTaskRelayInfoForcesChannelSelection(t *testing.T) { func TestRequiredTaskChannelTypeOnlyMatchesKlingAipingNativePaths(t *testing.T) { c := newControllerJSONContext(t, "/v1/videos/text2video", `{}`) c.Request.URL.Path = "/v1/videos/text2video" - require.Equal(t, constant.ChannelTypeKlingAiping, requiredTaskChannelTypeForRequest(c)) + require.Zero(t, requiredTaskChannelTypeForRequest(c)) + require.ElementsMatch(t, + service.VideoAssetChannelTypesForFamily(service.VideoAssetFamilyKling), + allowedTaskChannelTypesForRequest(c), + ) c = newControllerJSONContext(t, "/v1/videos/video-extend", `{}`) c.Request.URL.Path = "/v1/videos/video-extend" - require.Equal(t, constant.ChannelTypeKlingAiping, requiredTaskChannelTypeForRequest(c)) + require.Zero(t, requiredTaskChannelTypeForRequest(c)) + require.ElementsMatch(t, + service.VideoAssetChannelTypesForFamily(service.VideoAssetFamilyKling), + allowedTaskChannelTypesForRequest(c), + ) c = newControllerJSONContext(t, "/v1/videos/video_123/remix", `{}`) c.Request.URL.Path = "/v1/videos/video_123/remix" require.Zero(t, requiredTaskChannelTypeForRequest(c)) + require.Nil(t, allowedTaskChannelTypesForRequest(c)) } func TestKlingAipingTaskDataObjectUsesPublicTaskIDAndWatermarkURL(t *testing.T) { diff --git a/controller/relay.go b/controller/relay.go index bbe50bb..d7266f7 100644 --- a/controller/relay.go +++ b/controller/relay.go @@ -316,19 +316,14 @@ func persistTaskVideoBindingIfNeeded(userId int, group string, channel *model.Ch if group == "" || group == "auto" || channel == nil { return nil } - switch channel.Type { - case service.DoubaoAssetChannelType, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun: - if !service.IsUsableDoubaoAssetChannel(channel, group) { - return nil - } - return model.BindUserAssetChannel(userId, channel.Type, group, channel.Id) - case constant.ChannelTypeKlingAiping: - if !service.IsUsableKlingAssetChannel(channel, group) { - return nil - } - return model.BindUserAssetChannel(userId, constant.ChannelTypeKlingAiping, group, channel.Id) + family, ok := service.VideoAssetFamilyForChannelType(channel.Type) + if !ok { + return nil } - return nil + if !service.IsUsableVideoAssetChannelForFamily(channel, group, "", family) { + return nil + } + return service.BindVideoAssetChannel(userId, group, channel, family) } func getChannel(c *gin.Context, info *relaycommon.RelayInfo, retryParam *service.RetryParam) (*model.Channel, string, *types.NewAPIError) { @@ -367,34 +362,17 @@ func getChannel(c *gin.Context, info *relaycommon.RelayInfo, retryParam *service } func requiredTaskChannelTypeForRequest(c *gin.Context) int { - if strings.HasPrefix(c.Request.URL.Path, "/api/v3/contents/generations/tasks") { - return service.DoubaoAssetChannelTypeForModel(taskModelNameForRequiredChannelType(c)) - } - if isKlingAipingNativePath(c.Request.URL.Path) { - return constant.ChannelTypeKlingAiping - } return 0 } -func taskModelNameForRequiredChannelType(c *gin.Context) string { - if req, err := relaycommon.GetTaskRequest(c); err == nil { - return req.Model - } - bodyStorage, err := common.GetBodyStorage(c) - if err != nil { - return "" - } - body, err := bodyStorage.Bytes() - if err != nil { - return "" - } - var raw struct { - Model string `json:"model"` +func allowedTaskChannelTypesForRequest(c *gin.Context) []int { + if strings.HasPrefix(c.Request.URL.Path, "/api/v3/contents/generations/tasks") { + return service.VideoAssetChannelTypesForFamily(service.VideoAssetFamilySeedance) } - if err := common.Unmarshal(body, &raw); err != nil { - return "" + if isKlingAipingNativePath(c.Request.URL.Path) { + return service.VideoAssetChannelTypesForFamily(service.VideoAssetFamilyKling) } - return raw.Model + return nil } func shouldRetry(c *gin.Context, openaiErr *types.NewAPIError, retryTimes int) bool { @@ -690,6 +668,7 @@ func relayTaskWithInfo(c *gin.Context, relayInfo *relaycommon.RelayInfo) { ModelName: relayInfo.OriginModelName, Retry: common.GetPointer(0), RequiredChannelType: requiredTaskChannelTypeForRequest(c), + AllowedChannelTypes: allowedTaskChannelTypesForRequest(c), } for ; retryParam.GetRetry() <= common.RetryTimes; retryParam.IncreaseRetry() { diff --git a/middleware/distributor.go b/middleware/distributor.go index a4516ee..7f577fe 100644 --- a/middleware/distributor.go +++ b/middleware/distributor.go @@ -65,6 +65,7 @@ func Distribute() func(c *gin.Context) { } var selectGroup string usingGroup := common.GetContextKeyString(c, constant.ContextKeyUsingGroup) + allowedChannelTypes := allowedChannelTypesForRequest(c, modelRequest.Model) requiredChannelType := requiredChannelTypeForRequest(c, modelRequest.Model) // check path is /pg/chat/completions if strings.HasPrefix(c.Request.URL.Path, "/pg/chat/completions") { @@ -88,7 +89,7 @@ func Distribute() func(c *gin.Context) { if preferredChannelID, found := service.GetPreferredChannelByAffinity(c, modelRequest.Model, usingGroup); found { preferred, err := model.CacheGetChannel(preferredChannelID) if err == nil && preferred != nil && preferred.Status == common.ChannelStatusEnabled && - (requiredChannelType == 0 || preferred.Type == requiredChannelType) { + channelTypeMatchesRequest(preferred.Type, requiredChannelType, allowedChannelTypes) { if usingGroup == "auto" { userGroup := common.GetContextKeyString(c, constant.ContextKeyUserGroup) autoGroups := service.GetUserAutoGroup(userGroup) @@ -109,11 +110,11 @@ func Distribute() func(c *gin.Context) { } } - if channel == nil && shouldCheckDoubaoVideoBindingPreference(c, usingGroup, modelRequest.Model) { + if channel == nil && shouldCheckSeedanceVideoBindingPreference(c, usingGroup, modelRequest.Model) { bound, bindErr := service.GetBoundDoubaoAssetChannelForModel(c.GetInt("id"), usingGroup, modelRequest.Model) if bindErr != nil { logger.LogError(c, fmt.Sprintf("get bound doubao video channel failed: %v", bindErr)) - } else if shouldUseBoundDoubaoVideoChannel(c, bound, usingGroup, modelRequest.Model) { + } else if shouldUseBoundSeedanceVideoChannel(c, bound, usingGroup, modelRequest.Model) { channel = bound selectGroup = usingGroup } @@ -126,6 +127,7 @@ func Distribute() func(c *gin.Context) { TokenGroup: usingGroup, Retry: common.GetPointer(0), RequiredChannelType: requiredChannelType, + AllowedChannelTypes: allowedChannelTypes, }) if err != nil { showGroup := usingGroup @@ -341,21 +343,37 @@ func getModelRequest(c *gin.Context) (*ModelRequest, bool, error) { return &modelRequest, shouldSelectChannel, nil } -func shouldUseBoundDoubaoVideoChannel(_ *gin.Context, ch *model.Channel, group string, modelName string) bool { - if !service.IsUsableDoubaoAssetChannel(ch, group) { - return false +func allowedChannelTypesForRequest(c *gin.Context, modelName string) []int { + if strings.HasPrefix(c.Request.URL.Path, "/api/v3/contents/generations/tasks") { + return service.VideoAssetChannelTypesForFamily(service.VideoAssetFamilySeedance) } - return model.IsChannelEnabledForGroupModel(group, modelName, ch.Id) + return nil } -func requiredChannelTypeForRequest(c *gin.Context, modelName string) int { - if strings.HasPrefix(c.Request.URL.Path, "/api/v3/contents/generations/tasks") { - return service.DoubaoAssetChannelTypeForModel(modelName) +func channelTypeMatchesRequest(channelType int, requiredChannelType int, allowedChannelTypes []int) bool { + if requiredChannelType != 0 { + return channelType == requiredChannelType + } + if len(allowedChannelTypes) == 0 { + return true } + for _, allowedType := range allowedChannelTypes { + if channelType == allowedType { + return true + } + } + return false +} + +func shouldUseBoundSeedanceVideoChannel(_ *gin.Context, ch *model.Channel, group string, modelName string) bool { + return service.IsUsableVideoAssetChannelForFamily(ch, group, modelName, service.VideoAssetFamilySeedance) +} + +func requiredChannelTypeForRequest(c *gin.Context, modelName string) int { return 0 } -func shouldCheckDoubaoVideoBindingPreference(c *gin.Context, group string, modelName string) bool { +func shouldCheckSeedanceVideoBindingPreference(c *gin.Context, group string, modelName string) bool { group = strings.TrimSpace(group) modelName = strings.TrimSpace(modelName) if group == "" || group == "auto" || modelName == "" { @@ -364,7 +382,7 @@ func shouldCheckDoubaoVideoBindingPreference(c *gin.Context, group string, model if c.GetInt("relay_mode") != relayconstant.RelayModeVideoSubmit { return false } - return service.HasDoubaoVideoChannelForGroupModel(group, modelName) + return service.HasVideoAssetChannelForGroupModel(group, modelName, service.VideoAssetFamilySeedance) } func shouldPersistDoubaoVideoBinding(c *gin.Context) bool { @@ -376,13 +394,13 @@ func persistDoubaoVideoBindingIfNeeded(userId int, group string, ch *model.Chann if group == "" || group == "auto" { return nil } - if ch == nil || !service.IsDoubaoAssetChannelType(ch.Type) { + if ch == nil { return nil } - if !service.IsUsableDoubaoAssetChannel(ch, group) { + if !service.IsUsableVideoAssetChannelForFamily(ch, group, "", service.VideoAssetFamilySeedance) { return nil } - return model.BindUserAssetChannel(userId, ch.Type, group, ch.Id) + return service.BindVideoAssetChannel(userId, group, ch, service.VideoAssetFamilySeedance) } func SetupContextForSelectedChannel(c *gin.Context, channel *model.Channel, modelName string) *types.NewAPIError { diff --git a/middleware/doubao_asset_binding_test.go b/middleware/doubao_asset_binding_test.go index 39f3303..a495530 100644 --- a/middleware/doubao_asset_binding_test.go +++ b/middleware/doubao_asset_binding_test.go @@ -12,6 +12,7 @@ import ( "github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/model" relayconstant "github.com/QuantumNous/new-api/relay/constant" + "github.com/QuantumNous/new-api/service" "github.com/gin-gonic/gin" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -94,7 +95,7 @@ func tianyiyunVideoChannelForBindingTest(id int, group string) *model.Channel { return channel } -func TestPreferBoundDoubaoVideoChannel(t *testing.T) { +func TestPreferBoundSeedanceVideoChannel(t *testing.T) { db := setupDoubaoVideoBindingDB(t) channel := doubaoVideoChannelForBindingTest(77, "default") priority := int64(channel.Id) @@ -107,7 +108,7 @@ func TestPreferBoundDoubaoVideoChannel(t *testing.T) { Weight: 10, }).Error) - assert.True(t, shouldUseBoundDoubaoVideoChannel(nil, channel, "default", "seedance-2")) + assert.True(t, shouldUseBoundSeedanceVideoChannel(nil, channel, "default", "seedance-2")) } func TestPersistDoubaoVideoBindingSkipsNonDoubaoWithoutDB(t *testing.T) { @@ -163,14 +164,33 @@ func TestPersistDoubaoVideoBindingWritesTianyiYunBindingByType(t *testing.T) { assert.Equal(t, channel.Id, binding.ChannelId) } -func TestShouldCheckDoubaoVideoBindingPreferenceSkipsNonVideoRelayMode(t *testing.T) { +func TestPersistDoubaoVideoBindingReplacesOtherSeedanceFamilyBinding(t *testing.T) { + db := setupDoubaoVideoBindingDB(t) + require.NoError(t, db.Create(doubaoVideoChannelForBindingTest(7, "default")).Error) + require.NoError(t, db.Create(tianyiyunVideoChannelForBindingTest(16, "default")).Error) + require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", 7)) + ch, err := model.CacheGetChannel(16) + require.NoError(t, err) + + require.NoError(t, persistDoubaoVideoBindingIfNeeded(10, "default", ch)) + + oldBinding, err := model.GetUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default") + require.NoError(t, err) + assert.Nil(t, oldBinding) + newBinding, err := model.GetUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "default") + require.NoError(t, err) + require.NotNil(t, newBinding) + assert.Equal(t, 16, newBinding.ChannelId) +} + +func TestShouldCheckSeedanceVideoBindingPreferenceSkipsNonVideoRelayMode(t *testing.T) { c, _ := gin.CreateTestContext(nil) c.Set("relay_mode", relayconstant.RelayModeChatCompletions) - assert.False(t, shouldCheckDoubaoVideoBindingPreference(c, "default", "seedance-2")) + assert.False(t, shouldCheckSeedanceVideoBindingPreference(c, "default", "seedance-2")) } -func TestShouldCheckDoubaoVideoBindingPreferenceRequiresVideoSubmitRelayMode(t *testing.T) { +func TestShouldCheckSeedanceVideoBindingPreferenceRequiresVideoSubmitRelayMode(t *testing.T) { db := setupDoubaoVideoBindingDB(t) channel := doubaoVideoChannelForBindingTest(101, "default") require.NoError(t, db.Create(channel).Error) @@ -184,15 +204,18 @@ func TestShouldCheckDoubaoVideoBindingPreferenceRequiresVideoSubmitRelayMode(t * c, _ := gin.CreateTestContext(nil) c.Set("relay_mode", relayconstant.RelayModeVideoSubmit) - assert.True(t, shouldCheckDoubaoVideoBindingPreference(c, "default", "seedance-2")) + assert.True(t, shouldCheckSeedanceVideoBindingPreference(c, "default", "seedance-2")) } -func TestRequiredChannelTypeForSeedanceTasksUsesModelSpecificChannel(t *testing.T) { +func TestSeedanceTasksUseAllowedFamilyTypesInsteadOfModelPrefix(t *testing.T) { c, _ := gin.CreateTestContext(nil) c.Request = httptest.NewRequest(http.MethodPost, "/api/v3/contents/generations/tasks", strings.NewReader(`{}`)) - require.Equal(t, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, requiredChannelTypeForRequest(c, "cdance2.0-0611")) - require.Equal(t, constant.ChannelTypeDoubaoVideoCompatibleAiping, requiredChannelTypeForRequest(c, "doubao-seedance-2-0-260128")) + require.Equal(t, 0, requiredChannelTypeForRequest(c, "Doubao-Seedance-2.0")) + require.ElementsMatch(t, + service.VideoAssetChannelTypesForFamily(service.VideoAssetFamilySeedance), + allowedChannelTypesForRequest(c, "Doubao-Seedance-2.0"), + ) } func TestShouldPersistDoubaoVideoBindingRequiresVideoSubmitRelayMode(t *testing.T) { diff --git a/model/user_asset_channel.go b/model/user_asset_channel.go index 8d08093..7d374da 100644 --- a/model/user_asset_channel.go +++ b/model/user_asset_channel.go @@ -35,6 +35,10 @@ func GetUserAssetChannel(userId int, channelType int, group string) (*UserAssetC } func BindUserAssetChannel(userId int, channelType int, group string, channelId int) error { + return BindUserAssetChannelWithTx(DB, userId, channelType, group, channelId) +} + +func BindUserAssetChannelWithTx(tx *gorm.DB, userId int, channelType int, group string, channelId int) error { now := common.GetTimestamp() binding := UserAssetChannel{ UserId: userId, @@ -45,7 +49,7 @@ func BindUserAssetChannel(userId int, channelType int, group string, channelId i UpdatedAt: now, } - return DB.Clauses(clause.OnConflict{ + return tx.Session(&gorm.Session{SkipDefaultTransaction: true}).Clauses(clause.OnConflict{ Columns: []clause.Column{ {Name: "user_id"}, {Name: "channel_type"}, @@ -58,6 +62,25 @@ func BindUserAssetChannel(userId int, channelType int, group string, channelId i }).Create(&binding).Error } +func GetUserAssetChannelsByTypes(userId int, channelTypes []int, group string) ([]UserAssetChannel, error) { + if len(channelTypes) == 0 { + return nil, nil + } + var bindings []UserAssetChannel + err := DB.Where("user_id = ? AND channel_type IN ? AND "+commonGroupCol+" = ?", userId, channelTypes, group). + Order("updated_at DESC, id DESC"). + Find(&bindings).Error + return bindings, err +} + +func DeleteUserAssetChannelsByTypesWithTx(tx *gorm.DB, userId int, channelTypes []int, group string) error { + if len(channelTypes) == 0 { + return nil + } + return tx.Session(&gorm.Session{SkipDefaultTransaction: true}).Where("user_id = ? AND channel_type IN ? AND "+commonGroupCol+" = ?", userId, channelTypes, group). + Delete(&UserAssetChannel{}).Error +} + func UnbindUserAssetChannel(userId int, channelType int, group string) error { return DB.Where("user_id = ? AND channel_type = ? AND "+commonGroupCol+" = ?", userId, channelType, group). Delete(&UserAssetChannel{}).Error diff --git a/model/user_asset_channel_test.go b/model/user_asset_channel_test.go index cf5c150..a6caf07 100644 --- a/model/user_asset_channel_test.go +++ b/model/user_asset_channel_test.go @@ -71,6 +71,75 @@ func TestBindUserAssetChannel_CreateAndUpdate(t *testing.T) { assert.Equal(t, int64(1), count) } +func TestGetUserAssetChannelsByTypesSortsLatestFirst(t *testing.T) { + db := setupUserAssetChannelDB(t) + require.NoError(t, BindUserAssetChannel(1, 2, "default", 100)) + require.NoError(t, BindUserAssetChannel(1, 3, "default", 200)) + require.NoError(t, BindUserAssetChannel(1, 4, "default", 300)) + require.NoError(t, db.Model(&UserAssetChannel{}). + Where("user_id = ? AND channel_type = ?", 1, 2). + Update("updated_at", int64(100)).Error) + require.NoError(t, db.Model(&UserAssetChannel{}). + Where("user_id = ? AND channel_type = ?", 1, 3). + Update("updated_at", int64(200)).Error) + + bindings, err := GetUserAssetChannelsByTypes(1, []int{2, 3}, "default") + + require.NoError(t, err) + require.Len(t, bindings, 2) + assert.Equal(t, 3, bindings[0].ChannelType) + assert.Equal(t, 200, bindings[0].ChannelId) + assert.Equal(t, 2, bindings[1].ChannelType) +} + +func TestBindUserAssetChannelWithTxUpserts(t *testing.T) { + db := setupUserAssetChannelDB(t) + + require.NoError(t, BindUserAssetChannel(1, 2, "default", 100)) + require.NoError(t, BindUserAssetChannel(1, 2, "default", 200)) + + binding, err := GetUserAssetChannel(1, 2, "default") + require.NoError(t, err) + require.NotNil(t, binding) + assert.Equal(t, 200, binding.ChannelId) + + var count int64 + require.NoError(t, db.Model(&UserAssetChannel{}).Count(&count).Error) + assert.Equal(t, int64(1), count) +} + +func TestDeleteUserAssetChannelsByTypesWithTxDeletesOnlyRequestedRows(t *testing.T) { + db := setupUserAssetChannelDB(t) + require.NoError(t, BindUserAssetChannel(1, 2, "default", 100)) + require.NoError(t, BindUserAssetChannel(1, 3, "default", 200)) + require.NoError(t, BindUserAssetChannel(1, 4, "default", 300)) + require.NoError(t, BindUserAssetChannel(1, 2, "vip", 400)) + require.NoError(t, BindUserAssetChannel(2, 2, "default", 500)) + + require.NoError(t, db.Transaction(func(tx *gorm.DB) error { + return DeleteUserAssetChannelsByTypesWithTx(tx, 1, []int{2, 3}, "default") + })) + + binding, err := GetUserAssetChannel(1, 2, "default") + require.NoError(t, err) + assert.Nil(t, binding) + binding, err = GetUserAssetChannel(1, 3, "default") + require.NoError(t, err) + assert.Nil(t, binding) + binding, err = GetUserAssetChannel(1, 4, "default") + require.NoError(t, err) + require.NotNil(t, binding) + assert.Equal(t, 300, binding.ChannelId) + binding, err = GetUserAssetChannel(1, 2, "vip") + require.NoError(t, err) + require.NotNil(t, binding) + assert.Equal(t, 400, binding.ChannelId) + binding, err = GetUserAssetChannel(2, 2, "default") + require.NoError(t, err) + require.NotNil(t, binding) + assert.Equal(t, 500, binding.ChannelId) +} + func TestBindUserAssetChannel_IsolatesUserTypeGroup(t *testing.T) { setupUserAssetChannelDB(t) diff --git a/relay/channel/task/doubao_tianyiyun/adaptor_test.go b/relay/channel/task/doubao_tianyiyun/adaptor_test.go index 559c413..e671f44 100644 --- a/relay/channel/task/doubao_tianyiyun/adaptor_test.go +++ b/relay/channel/task/doubao_tianyiyun/adaptor_test.go @@ -84,6 +84,30 @@ func TestBuildRequestBodyPreservesTianyiYunNativeFields(t *testing.T) { require.NotContains(t, string(data), `"seconds"`) } +func TestBuildRequestBodyUsesMappedTianyiYunUpstreamModel(t *testing.T) { + adaptor := &TaskAdaptor{} + c := newTianyiYunTaskRequestContext(t, `{ + "model":"Doubao-Seedance-2.0", + "prompt":"write a clean product video" + }`) + info := &relaycommon.RelayInfo{ + OriginModelName: "Doubao-Seedance-2.0", + ChannelMeta: &relaycommon.ChannelMeta{ + UpstreamModelName: "cdance2.0-0611", + IsModelMapped: true, + }, + TaskRelayInfo: &relaycommon.TaskRelayInfo{}, + } + + body, err := adaptor.BuildRequestBody(c, info) + require.NoError(t, err) + data, err := io.ReadAll(body) + require.NoError(t, err) + + require.Contains(t, string(data), `"model":"cdance2.0-0611"`) + require.NotContains(t, string(data), `"model":"Doubao-Seedance-2.0"`) +} + func TestDoResponseRewritesIDToPublicTaskID(t *testing.T) { adaptor := &TaskAdaptor{} w := httptest.NewRecorder() diff --git a/relay/channel/task/kling/aiping/adaptor_test.go b/relay/channel/task/kling/aiping/adaptor_test.go index e3cca9a..8091a9a 100644 --- a/relay/channel/task/kling/aiping/adaptor_test.go +++ b/relay/channel/task/kling/aiping/adaptor_test.go @@ -162,6 +162,39 @@ func TestBuildRequestBodyMapsCanonicalModelNameForAiping(t *testing.T) { require.Equal(t, "Kling-V2.6", got["model_name"]) } +func TestBuildRequestBodyUsesMappedKlingUpstreamModel(t *testing.T) { + gin.SetMode(gin.TestMode) + c, _ := gin.CreateTestContext(nil) + info := &relaycommon.RelayInfo{ + OriginModelName: "public-kling-alias", + ChannelMeta: &relaycommon.ChannelMeta{ + UpstreamModelName: "kling-v2-6", + IsModelMapped: true, + }, + TaskRelayInfo: &relaycommon.TaskRelayInfo{ + Action: ActionText2Video, + }, + } + relaycommon.StoreTaskRequest(c, info, ActionText2Video, relaycommon.TaskSubmitReq{ + Model: "public-kling-alias", + Metadata: map[string]any{ + "model_name": "public-kling-alias", + "prompt": "prompt", + "duration": "5", + }, + }) + + body, err := (&TaskAdaptor{}).BuildRequestBody(c, info) + require.NoError(t, err) + data, err := io.ReadAll(body) + require.NoError(t, err) + + var got map[string]any + require.NoError(t, common.Unmarshal(data, &got)) + require.Equal(t, "Kling-V2.6", got["model_name"]) + require.NotContains(t, got, "model") +} + func TestValidateRequestRejectsOmniModelOnTextOrImageRoutes(t *testing.T) { gin.SetMode(gin.TestMode) c, _ := gin.CreateTestContext(nil) diff --git a/service/channel_select.go b/service/channel_select.go index 6f70973..3ff6b34 100644 --- a/service/channel_select.go +++ b/service/channel_select.go @@ -22,6 +22,7 @@ type RetryParam struct { Retry *int RequireMatrixUsageBilling bool RequiredChannelType int + AllowedChannelTypes []int resetNextTry bool } @@ -184,13 +185,16 @@ func CacheGetRandomSatisfiedChannel(param *RetryParam) (*model.Channel, string, } func matrixUsageBillingChannelFilter(param *RetryParam) model.ChannelFilter { - if !param.RequireMatrixUsageBilling && param.RequiredChannelType == 0 { + if !param.RequireMatrixUsageBilling && param.RequiredChannelType == 0 && len(param.AllowedChannelTypes) == 0 { return nil } return func(channel *model.Channel) (bool, error) { if param.RequiredChannelType != 0 && channel.Type != param.RequiredChannelType { return false, nil } + if !channelTypeAllowed(channel.Type, param.AllowedChannelTypes) { + return false, nil + } if param.RequireMatrixUsageBilling { mappedModel, err := helper.ResolvePricingModelMapping(param.ModelName, channel.GetModelMapping()) if err != nil { @@ -202,6 +206,18 @@ func matrixUsageBillingChannelFilter(param *RetryParam) model.ChannelFilter { } } +func channelTypeAllowed(channelType int, allowed []int) bool { + if len(allowed) == 0 { + return true + } + for _, allowedType := range allowed { + if channelType == allowedType { + return true + } + } + return false +} + func matrixUsageBillingChannelUnavailableError() *types.NewAPIError { return types.NewErrorWithStatusCode( fmt.Errorf("matrix_usage_billing_channel_unavailable"), diff --git a/service/channel_select_test.go b/service/channel_select_test.go index 2a302e1..30b32f2 100644 --- a/service/channel_select_test.go +++ b/service/channel_select_test.go @@ -3,12 +3,19 @@ package service import ( "net/http" "net/http/httptest" + "os" + "strings" "testing" + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/model" "github.com/QuantumNous/new-api/setting" "github.com/gin-gonic/gin" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "gorm.io/gorm" + "gorm.io/gorm/logger" ) func buildRetryContext(t *testing.T) *gin.Context { @@ -22,6 +29,81 @@ func buildRetryContext(t *testing.T) *gin.Context { return c } +func setupServiceChannelSelectDB(t *testing.T) *gorm.DB { + t.Helper() + + oldDB := model.DB + oldSQLitePath := common.SQLitePath + oldMemoryCacheEnabled := common.MemoryCacheEnabled + oldIsMasterNode := common.IsMasterNode + oldUsingSQLite := common.UsingSQLite + oldUsingMySQL := common.UsingMySQL + oldUsingPostgreSQL := common.UsingPostgreSQL + oldSQLDSN, hadSQLDSN := os.LookupEnv("SQL_DSN") + + common.SQLitePath = "file:" + strings.NewReplacer("/", "_", " ", "_").Replace(t.Name()) + "?mode=memory&cache=shared" + common.MemoryCacheEnabled = false + common.IsMasterNode = false + common.UsingSQLite = false + common.UsingMySQL = false + common.UsingPostgreSQL = false + require.NoError(t, os.Setenv("SQL_DSN", "local")) + require.NoError(t, model.InitDB()) + + model.DB = model.DB.Session(&gorm.Session{Logger: logger.Default.LogMode(logger.Silent)}) + db := model.DB + sqlDB, err := db.DB() + require.NoError(t, err) + sqlDB.SetMaxOpenConns(1) + require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.Ability{})) + + t.Cleanup(func() { + _ = sqlDB.Close() + model.DB = oldDB + common.SQLitePath = oldSQLitePath + common.MemoryCacheEnabled = oldMemoryCacheEnabled + common.IsMasterNode = oldIsMasterNode + common.UsingSQLite = oldUsingSQLite + common.UsingMySQL = oldUsingMySQL + common.UsingPostgreSQL = oldUsingPostgreSQL + if hadSQLDSN { + _ = os.Setenv("SQL_DSN", oldSQLDSN) + } else { + _ = os.Unsetenv("SQL_DSN") + } + }) + + return db +} + +func createServiceChannelSelectChannel(t *testing.T, db *gorm.DB, id int, channelType int, group string, modelName string) { + t.Helper() + + priority := int64(10) + weight := uint(10) + autoBan := 1 + require.NoError(t, db.Create(&model.Channel{ + Id: id, + Type: channelType, + Key: "key", + Status: common.ChannelStatusEnabled, + Name: "channel", + Group: group, + Models: modelName, + Priority: &priority, + Weight: &weight, + AutoBan: &autoBan, + }).Error) + require.NoError(t, db.Create(&model.Ability{ + Group: group, + Model: modelName, + ChannelId: id, + Enabled: true, + Priority: &priority, + Weight: 10, + }).Error) +} + // --- RetryParam tests --- func TestRetryParam_IncreaseRetry(t *testing.T) { @@ -77,6 +159,26 @@ func TestCacheGetRandomSatisfiedChannel_AutoGroup_NoAutoGroups(t *testing.T) { assert.Equal(t, "auto", group) } +func TestCacheGetRandomSatisfiedChannelFiltersAllowedChannelTypes(t *testing.T) { + db := setupServiceChannelSelectDB(t) + createServiceChannelSelectChannel(t, db, 101, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", "Doubao-Seedance-2.0") + createServiceChannelSelectChannel(t, db, 102, constant.ChannelTypeOpenAI, "default", "Doubao-Seedance-2.0") + + c := buildRetryContext(t) + channel, group, err := CacheGetRandomSatisfiedChannel(&RetryParam{ + Ctx: c, + TokenGroup: "default", + ModelName: "Doubao-Seedance-2.0", + Retry: common.GetPointer(0), + AllowedChannelTypes: []int{constant.ChannelTypeDoubaoVideoCompatibleAiping, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun}, + }) + + require.NoError(t, err) + require.NotNil(t, channel) + assert.Equal(t, "default", group) + assert.Equal(t, 101, channel.Id) +} + // --- GetUserAutoGroup tests --- func TestGetUserAutoGroup_Basic(t *testing.T) { diff --git a/service/doubao_asset_channel.go b/service/doubao_asset_channel.go index 8c40ac3..b74fb7f 100644 --- a/service/doubao_asset_channel.go +++ b/service/doubao_asset_channel.go @@ -1,163 +1,47 @@ package service import ( - "errors" "strings" - "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/model" - "gorm.io/gorm" ) const DoubaoAssetChannelPageSize = 50 const DoubaoAssetChannelType = constant.ChannelTypeDoubaoVideoCompatibleAiping -var DoubaoAssetChannelTypes = []int{ - constant.ChannelTypeDoubaoVideoCompatibleAiping, - constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, -} +var DoubaoAssetChannelTypes = VideoAssetChannelTypesForFamily(VideoAssetFamilySeedance) func ResolveDoubaoAssetChannel(userId int, tokenGroup string) (*model.Channel, error) { return ResolveDoubaoAssetChannelForModel(userId, tokenGroup, "") } func ResolveDoubaoAssetChannelForModel(userId int, tokenGroup string, modelName string) (*model.Channel, error) { - channelType := DoubaoAssetChannelTypeForModel(modelName) - channel, err := getBoundDoubaoAssetChannelByType(userId, tokenGroup, channelType) - if err != nil { - return nil, err - } - if channel != nil { - return channel, nil - } - - channel, err = autoMatchDoubaoAssetChannelByType(tokenGroup, channelType) - if err != nil { - return nil, err - } - if channel == nil { - return nil, nil - } - - if err = model.BindUserAssetChannel(userId, channel.Type, tokenGroup, channel.Id); err != nil { - return nil, err - } - return channel, nil + return ResolveVideoAssetChannelForModel(userId, tokenGroup, modelName, VideoAssetFamilySeedance) } func GetBoundDoubaoAssetChannel(userId int, tokenGroup string) (*model.Channel, error) { - return getBoundDoubaoAssetChannelByType(userId, tokenGroup, DoubaoAssetChannelType) + return GetBoundDoubaoAssetChannelForModel(userId, tokenGroup, "") } func GetBoundDoubaoAssetChannelForModel(userId int, tokenGroup string, modelName string) (*model.Channel, error) { - return getBoundDoubaoAssetChannelByType(userId, tokenGroup, DoubaoAssetChannelTypeForModel(modelName)) -} - -func getBoundDoubaoAssetChannelByType(userId int, tokenGroup string, channelType int) (*model.Channel, error) { - binding, err := model.GetUserAssetChannel(userId, channelType, tokenGroup) - if err != nil { - return nil, err - } - if binding == nil { - return nil, nil - } - - channel, err := model.CacheGetChannel(binding.ChannelId) - if err != nil { - if errors.Is(err, gorm.ErrRecordNotFound) || common.MemoryCacheEnabled { - return nil, nil - } - return nil, err - } - if !IsUsableDoubaoAssetChannel(channel, tokenGroup) { - return nil, nil - } - return channel, nil + return GetBoundVideoAssetChannelForModel(userId, tokenGroup, modelName, VideoAssetFamilySeedance) } func HasDoubaoVideoChannelForGroupModel(tokenGroup string, modelName string) bool { - if tokenGroup == "" || modelName == "" { - return false - } - - for _, channelType := range DoubaoAssetChannelTypes { - for startIdx := 0; ; startIdx += DoubaoAssetChannelPageSize { - candidates, err := model.GetChannelsByType(startIdx, DoubaoAssetChannelPageSize, true, channelType) - if err != nil { - return false - } - for _, candidate := range candidates { - if candidate == nil || !MatchDoubaoAssetGroup(candidate.GetGroups(), tokenGroup) { - continue - } - if model.IsChannelEnabledForGroupModel(tokenGroup, modelName, candidate.Id) { - return true - } - } - if len(candidates) < DoubaoAssetChannelPageSize { - break - } - } - } - return false + return HasVideoAssetChannelForGroupModel(tokenGroup, modelName, VideoAssetFamilySeedance) } func AutoMatchDoubaoAssetChannel(tokenGroup string) (*model.Channel, error) { - return autoMatchDoubaoAssetChannelByType(tokenGroup, DoubaoAssetChannelType) -} - -func autoMatchDoubaoAssetChannelByType(tokenGroup string, channelType int) (*model.Channel, error) { - for startIdx := 0; ; startIdx += DoubaoAssetChannelPageSize { - candidates, err := model.GetChannelsByType(startIdx, DoubaoAssetChannelPageSize, true, channelType) - if err != nil { - return nil, err - } - for _, candidate := range candidates { - if candidate == nil || !MatchDoubaoAssetGroup(candidate.GetGroups(), tokenGroup) { - continue - } - - channel, err := model.CacheGetChannel(candidate.Id) - if err != nil { - if errors.Is(err, gorm.ErrRecordNotFound) || common.MemoryCacheEnabled { - continue - } - return nil, err - } - if IsUsableDoubaoAssetChannel(channel, tokenGroup) { - return channel, nil - } - } - if len(candidates) < DoubaoAssetChannelPageSize { - return nil, nil - } - } + return autoMatchVideoAssetChannelForModel(tokenGroup, "", VideoAssetFamilySeedance) } func IsUsableDoubaoAssetChannel(channel *model.Channel, tokenGroup string) bool { - if channel == nil { - return false - } - if !IsDoubaoAssetChannelType(channel.Type) { - return false - } - if channel.Status != common.ChannelStatusEnabled { - return false - } - if strings.TrimSpace(channel.Key) == "" || len(channel.GetKeys()) == 0 { - return false - } - return MatchDoubaoAssetGroup(channel.GetGroups(), tokenGroup) + return IsUsableVideoAssetChannelForFamily(channel, tokenGroup, "", VideoAssetFamilySeedance) } func IsDoubaoAssetChannelType(channelType int) bool { - for _, t := range DoubaoAssetChannelTypes { - if channelType == t { - return true - } - } - return false + return channelTypeAllowed(channelType, DoubaoAssetChannelTypes) } func DoubaoAssetChannelTypeForModel(modelName string) int { diff --git a/service/doubao_asset_channel_test.go b/service/doubao_asset_channel_test.go index 11ee29d..0a14e2f 100644 --- a/service/doubao_asset_channel_test.go +++ b/service/doubao_asset_channel_test.go @@ -132,9 +132,9 @@ func TestResolveDoubaoAssetChannel_AutoMatchesAndPersistsWhenUnbound(t *testing. func TestResolveDoubaoAssetChannel_AutoMatchesTianyiYunAndPersistsByType(t *testing.T) { db := setupDoubaoAssetChannelDB(t) createDoubaoAssetChannelForTest(t, db, 1, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "default", "tianyiyun-key", common.ChannelStatusEnabled) - createDoubaoAssetAbilityForTest(t, db, "default", "cdance2.0-0611", 1, true) + createDoubaoAssetAbilityForTest(t, db, "default", "Doubao-Seedance-2.0", 1, true) - channel, err := ResolveDoubaoAssetChannelForModel(10, "default", "cdance2.0-0611") + channel, err := ResolveDoubaoAssetChannelForModel(10, "default", "Doubao-Seedance-2.0") require.NoError(t, err) require.NotNil(t, channel) @@ -147,6 +147,19 @@ func TestResolveDoubaoAssetChannel_AutoMatchesTianyiYunAndPersistsByType(t *test assert.Equal(t, 1, binding.ChannelId) } +func TestGetBoundDoubaoAssetChannelForModelFindsTianyiYunBindingForPublicModel(t *testing.T) { + db := setupDoubaoAssetChannelDB(t) + createDoubaoAssetChannelForTest(t, db, 16, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "default", "tianyiyun-key", common.ChannelStatusEnabled) + createDoubaoAssetAbilityForTest(t, db, "default", "Doubao-Seedance-2.0", 16, true) + require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "default", 16)) + + channel, err := GetBoundDoubaoAssetChannelForModel(10, "default", "Doubao-Seedance-2.0") + + require.NoError(t, err) + require.NotNil(t, channel) + assert.Equal(t, 16, channel.Id) +} + func TestResolveDoubaoAssetChannel_RebindsOpenAIMisbindToAiping(t *testing.T) { db := setupDoubaoAssetChannelDB(t) createDoubaoAssetChannelForTest(t, db, 1, constant.ChannelTypeOpenAI, "default", "openai-key", common.ChannelStatusEnabled) @@ -230,10 +243,10 @@ func TestHasDoubaoVideoChannelForGroupModel(t *testing.T) { createDoubaoAssetAbilityForTest(t, db, "vip", "seedance-2", 2, true) createDoubaoAssetAbilityForTest(t, db, "default", "disabled-model", 1, false) createDoubaoAssetAbilityForTest(t, db, "default", "seedance-2", 3, true) - createDoubaoAssetAbilityForTest(t, db, "default", "cdance2.0-0611", 4, true) + createDoubaoAssetAbilityForTest(t, db, "default", "Doubao-Seedance-2.0", 4, true) assert.True(t, HasDoubaoVideoChannelForGroupModel("default", "seedance-2")) - assert.True(t, HasDoubaoVideoChannelForGroupModel("default", "cdance2.0-0611")) + assert.True(t, HasDoubaoVideoChannelForGroupModel("default", "Doubao-Seedance-2.0")) assert.False(t, HasDoubaoVideoChannelForGroupModel("default", "disabled-model")) assert.False(t, HasDoubaoVideoChannelForGroupModel("missing", "seedance-2")) } diff --git a/service/kling_asset_channel.go b/service/kling_asset_channel.go index 4836202..6519f59 100644 --- a/service/kling_asset_channel.go +++ b/service/kling_asset_channel.go @@ -1,102 +1,32 @@ package service import ( - "errors" - "strings" - - "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/model" - "gorm.io/gorm" ) const KlingAssetChannelType = constant.ChannelTypeKlingAiping -func ResolveKlingAssetChannel(userId int, tokenGroup string) (*model.Channel, error) { - channel, err := GetBoundKlingAssetChannel(userId, tokenGroup) - if err != nil { - return nil, err - } - if channel != nil { - return channel, nil - } +func ResolveKlingAssetChannelForModel(userId int, tokenGroup string, modelName string) (*model.Channel, error) { + return ResolveVideoAssetChannelForModel(userId, tokenGroup, modelName, VideoAssetFamilyKling) +} - channel, err = AutoMatchKlingAssetChannel(tokenGroup) - if err != nil { - return nil, err - } - if channel == nil { - return nil, nil - } +func ResolveKlingAssetChannel(userId int, tokenGroup string) (*model.Channel, error) { + return ResolveKlingAssetChannelForModel(userId, tokenGroup, "") +} - if err = model.BindUserAssetChannel(userId, KlingAssetChannelType, tokenGroup, channel.Id); err != nil { - return nil, err - } - return channel, nil +func GetBoundKlingAssetChannelForModel(userId int, tokenGroup string, modelName string) (*model.Channel, error) { + return GetBoundVideoAssetChannelForModel(userId, tokenGroup, modelName, VideoAssetFamilyKling) } func GetBoundKlingAssetChannel(userId int, tokenGroup string) (*model.Channel, error) { - binding, err := model.GetUserAssetChannel(userId, KlingAssetChannelType, tokenGroup) - if err != nil { - return nil, err - } - if binding == nil { - return nil, nil - } - - channel, err := model.CacheGetChannel(binding.ChannelId) - if err != nil { - if errors.Is(err, gorm.ErrRecordNotFound) || common.MemoryCacheEnabled { - return nil, nil - } - return nil, err - } - if !IsUsableKlingAssetChannel(channel, tokenGroup) { - return nil, nil - } - return channel, nil + return GetBoundKlingAssetChannelForModel(userId, tokenGroup, "") } func AutoMatchKlingAssetChannel(tokenGroup string) (*model.Channel, error) { - for startIdx := 0; ; startIdx += DoubaoAssetChannelPageSize { - candidates, err := model.GetChannelsByType(startIdx, DoubaoAssetChannelPageSize, true, KlingAssetChannelType) - if err != nil { - return nil, err - } - for _, candidate := range candidates { - if candidate == nil || !MatchDoubaoAssetGroup(candidate.GetGroups(), tokenGroup) { - continue - } - - channel, err := model.CacheGetChannel(candidate.Id) - if err != nil { - if errors.Is(err, gorm.ErrRecordNotFound) || common.MemoryCacheEnabled { - continue - } - return nil, err - } - if IsUsableKlingAssetChannel(channel, tokenGroup) { - return channel, nil - } - } - if len(candidates) < DoubaoAssetChannelPageSize { - return nil, nil - } - } + return autoMatchVideoAssetChannelForModel(tokenGroup, "", VideoAssetFamilyKling) } func IsUsableKlingAssetChannel(channel *model.Channel, tokenGroup string) bool { - if channel == nil { - return false - } - if channel.Type != KlingAssetChannelType { - return false - } - if channel.Status != common.ChannelStatusEnabled { - return false - } - if strings.TrimSpace(channel.Key) == "" || len(channel.GetKeys()) == 0 { - return false - } - return MatchDoubaoAssetGroup(channel.GetGroups(), tokenGroup) + return IsUsableVideoAssetChannelForFamily(channel, tokenGroup, "", VideoAssetFamilyKling) } diff --git a/service/kling_asset_channel_test.go b/service/kling_asset_channel_test.go index 9eab53f..3d579b6 100644 --- a/service/kling_asset_channel_test.go +++ b/service/kling_asset_channel_test.go @@ -53,6 +53,30 @@ func TestResolveKlingAssetChannel_IgnoresDisabledChannel(t *testing.T) { assert.Equal(t, 2, channel.Id) } +func TestGetBoundKlingAssetChannelForModelRequiresAbility(t *testing.T) { + db := setupDoubaoAssetChannelDB(t) + createDoubaoAssetChannelForTest(t, db, 59, constant.ChannelTypeKlingAiping, "default", "key", common.ChannelStatusEnabled) + require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeKlingAiping, "default", 59)) + + channel, err := GetBoundKlingAssetChannelForModel(10, "default", "kling-v2-6") + + require.NoError(t, err) + assert.Nil(t, channel) +} + +func TestGetBoundKlingAssetChannelForModelFindsBoundChannel(t *testing.T) { + db := setupDoubaoAssetChannelDB(t) + createDoubaoAssetChannelForTest(t, db, 59, constant.ChannelTypeKlingAiping, "default", "key", common.ChannelStatusEnabled) + createDoubaoAssetAbilityForTest(t, db, "default", "kling-v2-6", 59, true) + require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeKlingAiping, "default", 59)) + + channel, err := GetBoundKlingAssetChannelForModel(10, "default", "kling-v2-6") + + require.NoError(t, err) + require.NotNil(t, channel) + assert.Equal(t, 59, channel.Id) +} + func TestIsUsableKlingAssetChannel_RejectsOtherChannelType(t *testing.T) { ch := &model.Channel{Type: constant.ChannelTypeDoubaoVideoCompatibleAiping, Key: "k", Status: common.ChannelStatusEnabled, Group: "default"} assert.False(t, IsUsableKlingAssetChannel(ch, "default")) @@ -61,4 +85,4 @@ func TestIsUsableKlingAssetChannel_RejectsOtherChannelType(t *testing.T) { func TestIsUsableKlingAssetChannel_RejectsGroupMismatch(t *testing.T) { ch := &model.Channel{Type: constant.ChannelTypeKlingAiping, Key: "k", Status: common.ChannelStatusEnabled, Group: "vip"} assert.False(t, IsUsableKlingAssetChannel(ch, "default")) -} \ No newline at end of file +} diff --git a/service/video_asset_channel.go b/service/video_asset_channel.go new file mode 100644 index 0000000..f4478b4 --- /dev/null +++ b/service/video_asset_channel.go @@ -0,0 +1,171 @@ +package service + +import ( + "errors" + "strings" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/model" + "gorm.io/gorm" +) + +type VideoAssetFamily string + +const ( + VideoAssetFamilySeedance VideoAssetFamily = "seedance" + VideoAssetFamilyKling VideoAssetFamily = "kling" +) + +func VideoAssetChannelTypesForFamily(family VideoAssetFamily) []int { + switch family { + case VideoAssetFamilySeedance: + return []int{ + constant.ChannelTypeDoubaoVideoCompatibleAiping, + constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, + } + case VideoAssetFamilyKling: + return []int{ + constant.ChannelTypeKlingAiping, + } + default: + return nil + } +} + +func VideoAssetFamilyForChannelType(channelType int) (VideoAssetFamily, bool) { + for _, family := range []VideoAssetFamily{VideoAssetFamilySeedance, VideoAssetFamilyKling} { + if channelTypeAllowed(channelType, VideoAssetChannelTypesForFamily(family)) { + return family, true + } + } + return "", false +} + +func GetBoundVideoAssetChannelForModel(userId int, tokenGroup string, modelName string, family VideoAssetFamily) (*model.Channel, error) { + bindings, err := model.GetUserAssetChannelsByTypes(userId, VideoAssetChannelTypesForFamily(family), tokenGroup) + if err != nil { + return nil, err + } + for _, binding := range bindings { + channel, err := model.CacheGetChannel(binding.ChannelId) + if err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) || common.MemoryCacheEnabled { + continue + } + return nil, err + } + if IsUsableVideoAssetChannelForFamily(channel, tokenGroup, modelName, family) { + return channel, nil + } + } + return nil, nil +} + +func ResolveVideoAssetChannelForModel(userId int, tokenGroup string, modelName string, family VideoAssetFamily) (*model.Channel, error) { + channel, err := GetBoundVideoAssetChannelForModel(userId, tokenGroup, modelName, family) + if err != nil { + return nil, err + } + if channel != nil { + return channel, nil + } + + channel, err = autoMatchVideoAssetChannelForModel(tokenGroup, modelName, family) + if err != nil { + return nil, err + } + if channel == nil { + return nil, nil + } + if err = BindVideoAssetChannel(userId, tokenGroup, channel, family); err != nil { + return nil, err + } + return channel, nil +} + +func BindVideoAssetChannel(userId int, tokenGroup string, channel *model.Channel, family VideoAssetFamily) error { + if channel == nil { + return nil + } + channelTypes := VideoAssetChannelTypesForFamily(family) + if len(channelTypes) == 0 || !channelTypeAllowed(channel.Type, channelTypes) { + return nil + } + otherTypes := make([]int, 0, len(channelTypes)) + for _, channelType := range channelTypes { + if channelType != channel.Type { + otherTypes = append(otherTypes, channelType) + } + } + return model.DB.Transaction(func(tx *gorm.DB) error { + if err := model.DeleteUserAssetChannelsByTypesWithTx(tx, userId, otherTypes, tokenGroup); err != nil { + return err + } + return model.BindUserAssetChannelWithTx(tx, userId, channel.Type, tokenGroup, channel.Id) + }) +} + +func HasVideoAssetChannelForGroupModel(tokenGroup string, modelName string, family VideoAssetFamily) bool { + tokenGroup = strings.TrimSpace(tokenGroup) + modelName = strings.TrimSpace(modelName) + if tokenGroup == "" || modelName == "" { + return false + } + channel, err := autoMatchVideoAssetChannelForModel(tokenGroup, modelName, family) + return err == nil && channel != nil +} + +func autoMatchVideoAssetChannelForModel(tokenGroup string, modelName string, family VideoAssetFamily) (*model.Channel, error) { + for _, channelType := range VideoAssetChannelTypesForFamily(family) { + for startIdx := 0; ; startIdx += DoubaoAssetChannelPageSize { + candidates, err := model.GetChannelsByType(startIdx, DoubaoAssetChannelPageSize, true, channelType) + if err != nil { + return nil, err + } + for _, candidate := range candidates { + if candidate == nil || !MatchDoubaoAssetGroup(candidate.GetGroups(), tokenGroup) { + continue + } + channel, err := model.CacheGetChannel(candidate.Id) + if err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) || common.MemoryCacheEnabled { + continue + } + return nil, err + } + if IsUsableVideoAssetChannelForFamily(channel, tokenGroup, modelName, family) { + return channel, nil + } + } + if len(candidates) < DoubaoAssetChannelPageSize { + break + } + } + } + return nil, nil +} + +func IsUsableVideoAssetChannelForFamily(channel *model.Channel, tokenGroup string, modelName string, family VideoAssetFamily) bool { + if channel == nil { + return false + } + channelTypes := VideoAssetChannelTypesForFamily(family) + if len(channelTypes) == 0 || !channelTypeAllowed(channel.Type, channelTypes) { + return false + } + if channel.Status != common.ChannelStatusEnabled { + return false + } + if strings.TrimSpace(channel.Key) == "" || len(channel.GetKeys()) == 0 { + return false + } + if !MatchDoubaoAssetGroup(channel.GetGroups(), tokenGroup) { + return false + } + modelName = strings.TrimSpace(modelName) + if modelName == "" { + return true + } + return model.IsChannelEnabledForGroupModel(tokenGroup, modelName, channel.Id) +} diff --git a/service/video_asset_channel_test.go b/service/video_asset_channel_test.go new file mode 100644 index 0000000..1ea2109 --- /dev/null +++ b/service/video_asset_channel_test.go @@ -0,0 +1,107 @@ +package service + +import ( + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/model" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestGetBoundVideoAssetChannelForModelFindsSeedanceTianyiYunBinding(t *testing.T) { + db := setupDoubaoAssetChannelDB(t) + createDoubaoAssetChannelForTest(t, db, 16, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "default", "key", common.ChannelStatusEnabled) + createDoubaoAssetAbilityForTest(t, db, "default", "Doubao-Seedance-2.0", 16, true) + require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "default", 16)) + + ch, err := GetBoundVideoAssetChannelForModel(10, "default", "Doubao-Seedance-2.0", VideoAssetFamilySeedance) + + require.NoError(t, err) + require.NotNil(t, ch) + assert.Equal(t, 16, ch.Id) +} + +func TestGetBoundVideoAssetChannelForModelFindsKlingBindingByFamily(t *testing.T) { + db := setupDoubaoAssetChannelDB(t) + createDoubaoAssetChannelForTest(t, db, 59, constant.ChannelTypeKlingAiping, "default", "key", common.ChannelStatusEnabled) + createDoubaoAssetAbilityForTest(t, db, "default", "kling-v2-6", 59, true) + require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeKlingAiping, "default", 59)) + + ch, err := GetBoundVideoAssetChannelForModel(10, "default", "kling-v2-6", VideoAssetFamilyKling) + + require.NoError(t, err) + require.NotNil(t, ch) + assert.Equal(t, 59, ch.Id) +} + +func TestGetBoundVideoAssetChannelForModelRejectsBindingWithoutAbility(t *testing.T) { + db := setupDoubaoAssetChannelDB(t) + createDoubaoAssetChannelForTest(t, db, 16, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "default", "key", common.ChannelStatusEnabled) + require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "default", 16)) + + ch, err := GetBoundVideoAssetChannelForModel(10, "default", "Doubao-Seedance-2.0", VideoAssetFamilySeedance) + + require.NoError(t, err) + assert.Nil(t, ch) +} + +func TestBindVideoAssetChannelReplacesOtherFamilyBindings(t *testing.T) { + db := setupDoubaoAssetChannelDB(t) + createDoubaoAssetChannelForTest(t, db, 7, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", "aiping-key", common.ChannelStatusEnabled) + createDoubaoAssetChannelForTest(t, db, 16, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "default", "tianyiyun-key", common.ChannelStatusEnabled) + createDoubaoAssetAbilityForTest(t, db, "default", "Doubao-Seedance-2.0", 7, true) + createDoubaoAssetAbilityForTest(t, db, "default", "Doubao-Seedance-2.0", 16, true) + require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", 7)) + selected, err := model.CacheGetChannel(16) + require.NoError(t, err) + + require.NoError(t, BindVideoAssetChannel(10, "default", selected, VideoAssetFamilySeedance)) + + oldBinding, err := model.GetUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default") + require.NoError(t, err) + assert.Nil(t, oldBinding) + newBinding, err := model.GetUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "default") + require.NoError(t, err) + require.NotNil(t, newBinding) + assert.Equal(t, 16, newBinding.ChannelId) +} + +func TestGetBoundVideoAssetChannelForModelPrefersLatestFamilyBindingForLegacyRows(t *testing.T) { + db := setupDoubaoAssetChannelDB(t) + createDoubaoAssetChannelForTest(t, db, 7, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", "aiping-key", common.ChannelStatusEnabled) + createDoubaoAssetChannelForTest(t, db, 16, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "default", "tianyiyun-key", common.ChannelStatusEnabled) + createDoubaoAssetAbilityForTest(t, db, "default", "Doubao-Seedance-2.0", 7, true) + createDoubaoAssetAbilityForTest(t, db, "default", "Doubao-Seedance-2.0", 16, true) + require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", 7)) + require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "default", 16)) + require.NoError(t, db.Model(&model.UserAssetChannel{}). + Where("user_id = ? AND channel_type = ?", 10, constant.ChannelTypeDoubaoVideoCompatibleAiping). + Update("updated_at", int64(100)).Error) + require.NoError(t, db.Model(&model.UserAssetChannel{}). + Where("user_id = ? AND channel_type = ?", 10, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun). + Update("updated_at", int64(200)).Error) + + ch, err := GetBoundVideoAssetChannelForModel(10, "default", "Doubao-Seedance-2.0", VideoAssetFamilySeedance) + + require.NoError(t, err) + require.NotNil(t, ch) + assert.Equal(t, 16, ch.Id) +} + +func TestResolveVideoAssetChannelForModelAutoSelectsAndPersistsFamilyBinding(t *testing.T) { + db := setupDoubaoAssetChannelDB(t) + createDoubaoAssetChannelForTest(t, db, 16, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "default", "tianyiyun-key", common.ChannelStatusEnabled) + createDoubaoAssetAbilityForTest(t, db, "default", "Doubao-Seedance-2.0", 16, true) + + ch, err := ResolveVideoAssetChannelForModel(10, "default", "Doubao-Seedance-2.0", VideoAssetFamilySeedance) + + require.NoError(t, err) + require.NotNil(t, ch) + assert.Equal(t, 16, ch.Id) + binding, err := model.GetUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "default") + require.NoError(t, err) + require.NotNil(t, binding) + assert.Equal(t, 16, binding.ChannelId) +}