| @@ -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), | |||
| ) | |||
| } | |||
| @@ -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) | |||
| } | |||
| @@ -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 | |||
| @@ -167,7 +167,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()}) | |||
| @@ -193,7 +193,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 | |||
| @@ -102,15 +102,24 @@ func TestConfigureKlingAipingTaskRelayInfoStoresDuration(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) { | |||
| @@ -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() { | |||
| @@ -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 { | |||
| @@ -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) { | |||
| @@ -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 | |||
| @@ -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) | |||
| @@ -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() | |||
| @@ -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) | |||
| @@ -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"), | |||
| @@ -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) { | |||
| @@ -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 { | |||
| @@ -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")) | |||
| } | |||
| @@ -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) | |||
| } | |||
| @@ -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")) | |||
| } | |||
| } | |||
| @@ -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) | |||
| } | |||
| @@ -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) | |||
| } | |||