Co-Authored-By: Codex <noreply@anthropic.com>master^2
| @@ -7,6 +7,7 @@ import ( | |||||
| "github.com/QuantumNous/new-api/common" | "github.com/QuantumNous/new-api/common" | ||||
| "github.com/QuantumNous/new-api/constant" | "github.com/QuantumNous/new-api/constant" | ||||
| "github.com/QuantumNous/new-api/model" | "github.com/QuantumNous/new-api/model" | ||||
| "github.com/QuantumNous/new-api/service" | |||||
| "github.com/stretchr/testify/require" | "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/dto" | ||||
| "github.com/QuantumNous/new-api/model" | "github.com/QuantumNous/new-api/model" | ||||
| relaycommon "github.com/QuantumNous/new-api/relay/common" | relaycommon "github.com/QuantumNous/new-api/relay/common" | ||||
| "github.com/QuantumNous/new-api/service" | |||||
| "github.com/QuantumNous/new-api/types" | "github.com/QuantumNous/new-api/types" | ||||
| "github.com/gin-gonic/gin" | "github.com/gin-gonic/gin" | ||||
| ) | ) | ||||
| @@ -94,9 +93,6 @@ func AipingNativeVideoSubmit(c *gin.Context) { | |||||
| }) | }) | ||||
| return | return | ||||
| } | } | ||||
| channelType := service.DoubaoAssetChannelTypeForModel(relayInfo.OriginModelName) | |||||
| c.Set("channel_type", channelType) | |||||
| relayInfo.ChannelType = channelType | |||||
| relayTaskWithInfo(c, relayInfo) | relayTaskWithInfo(c, relayInfo) | ||||
| } | } | ||||
| @@ -54,7 +54,7 @@ func resolveKlingAipingBoundChannel(c *gin.Context, relayInfo *relaycommon.Relay | |||||
| if group == "" { | if group == "" { | ||||
| return nil | return nil | ||||
| } | } | ||||
| channel, err := service.ResolveKlingAssetChannel(c.GetInt("id"), group) | |||||
| channel, err := service.GetBoundKlingAssetChannelForModel(c.GetInt("id"), group, relayInfo.OriginModelName) | |||||
| if err != nil { | if err != nil { | ||||
| logger.LogError(c, fmt.Sprintf("resolve kling aiping bound channel failed: %v", err)) | logger.LogError(c, fmt.Sprintf("resolve kling aiping bound channel failed: %v", err)) | ||||
| return nil | return nil | ||||
| @@ -141,7 +141,7 @@ func KlingAipingNativeProxy(c *gin.Context) { | |||||
| TokenGroup: group, | TokenGroup: group, | ||||
| ModelName: route.BillingModel, | ModelName: route.BillingModel, | ||||
| Retry: common.GetPointer(0), | Retry: common.GetPointer(0), | ||||
| RequiredChannelType: constant.ChannelTypeKlingAiping, | |||||
| AllowedChannelTypes: service.VideoAssetChannelTypesForFamily(service.VideoAssetFamilyKling), | |||||
| }) | }) | ||||
| if err != nil { | if err != nil { | ||||
| c.JSON(http.StatusServiceUnavailable, gin.H{"code": 503, "message": err.Error()}) | 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) == "" { | if group == "" || group == "auto" || strings.TrimSpace(billingModel) == "" { | ||||
| return nil | return nil | ||||
| } | } | ||||
| channel, err := service.ResolveKlingAssetChannel(c.GetInt("id"), group) | |||||
| channel, err := service.GetBoundKlingAssetChannelForModel(c.GetInt("id"), group, billingModel) | |||||
| if err != nil { | if err != nil { | ||||
| logger.LogError(c, fmt.Sprintf("resolve kling aiping bound channel (proxy) failed: %v", err)) | logger.LogError(c, fmt.Sprintf("resolve kling aiping bound channel (proxy) failed: %v", err)) | ||||
| return nil | return nil | ||||
| @@ -86,15 +86,24 @@ func TestConfigureKlingAipingTaskRelayInfoForcesChannelSelection(t *testing.T) { | |||||
| func TestRequiredTaskChannelTypeOnlyMatchesKlingAipingNativePaths(t *testing.T) { | func TestRequiredTaskChannelTypeOnlyMatchesKlingAipingNativePaths(t *testing.T) { | ||||
| c := newControllerJSONContext(t, "/v1/videos/text2video", `{}`) | c := newControllerJSONContext(t, "/v1/videos/text2video", `{}`) | ||||
| c.Request.URL.Path = "/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 = newControllerJSONContext(t, "/v1/videos/video-extend", `{}`) | ||||
| c.Request.URL.Path = "/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 = newControllerJSONContext(t, "/v1/videos/video_123/remix", `{}`) | ||||
| c.Request.URL.Path = "/v1/videos/video_123/remix" | c.Request.URL.Path = "/v1/videos/video_123/remix" | ||||
| require.Zero(t, requiredTaskChannelTypeForRequest(c)) | require.Zero(t, requiredTaskChannelTypeForRequest(c)) | ||||
| require.Nil(t, allowedTaskChannelTypesForRequest(c)) | |||||
| } | } | ||||
| func TestKlingAipingTaskDataObjectUsesPublicTaskIDAndWatermarkURL(t *testing.T) { | func TestKlingAipingTaskDataObjectUsesPublicTaskIDAndWatermarkURL(t *testing.T) { | ||||
| @@ -316,19 +316,14 @@ func persistTaskVideoBindingIfNeeded(userId int, group string, channel *model.Ch | |||||
| if group == "" || group == "auto" || channel == nil { | if group == "" || group == "auto" || channel == nil { | ||||
| return 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) { | 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 { | 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 | 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 { | 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, | ModelName: relayInfo.OriginModelName, | ||||
| Retry: common.GetPointer(0), | Retry: common.GetPointer(0), | ||||
| RequiredChannelType: requiredTaskChannelTypeForRequest(c), | RequiredChannelType: requiredTaskChannelTypeForRequest(c), | ||||
| AllowedChannelTypes: allowedTaskChannelTypesForRequest(c), | |||||
| } | } | ||||
| for ; retryParam.GetRetry() <= common.RetryTimes; retryParam.IncreaseRetry() { | for ; retryParam.GetRetry() <= common.RetryTimes; retryParam.IncreaseRetry() { | ||||
| @@ -65,6 +65,7 @@ func Distribute() func(c *gin.Context) { | |||||
| } | } | ||||
| var selectGroup string | var selectGroup string | ||||
| usingGroup := common.GetContextKeyString(c, constant.ContextKeyUsingGroup) | usingGroup := common.GetContextKeyString(c, constant.ContextKeyUsingGroup) | ||||
| allowedChannelTypes := allowedChannelTypesForRequest(c, modelRequest.Model) | |||||
| requiredChannelType := requiredChannelTypeForRequest(c, modelRequest.Model) | requiredChannelType := requiredChannelTypeForRequest(c, modelRequest.Model) | ||||
| // check path is /pg/chat/completions | // check path is /pg/chat/completions | ||||
| if strings.HasPrefix(c.Request.URL.Path, "/pg/chat/completions") { | if strings.HasPrefix(c.Request.URL.Path, "/pg/chat/completions") { | ||||
| @@ -88,7 +89,7 @@ func Distribute() func(c *gin.Context) { | |||||
| if preferredChannelID, found := service.GetPreferredChannelByAffinity(c, modelRequest.Model, usingGroup); found { | if preferredChannelID, found := service.GetPreferredChannelByAffinity(c, modelRequest.Model, usingGroup); found { | ||||
| preferred, err := model.CacheGetChannel(preferredChannelID) | preferred, err := model.CacheGetChannel(preferredChannelID) | ||||
| if err == nil && preferred != nil && preferred.Status == common.ChannelStatusEnabled && | if err == nil && preferred != nil && preferred.Status == common.ChannelStatusEnabled && | ||||
| (requiredChannelType == 0 || preferred.Type == requiredChannelType) { | |||||
| channelTypeMatchesRequest(preferred.Type, requiredChannelType, allowedChannelTypes) { | |||||
| if usingGroup == "auto" { | if usingGroup == "auto" { | ||||
| userGroup := common.GetContextKeyString(c, constant.ContextKeyUserGroup) | userGroup := common.GetContextKeyString(c, constant.ContextKeyUserGroup) | ||||
| autoGroups := service.GetUserAutoGroup(userGroup) | 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) | bound, bindErr := service.GetBoundDoubaoAssetChannelForModel(c.GetInt("id"), usingGroup, modelRequest.Model) | ||||
| if bindErr != nil { | if bindErr != nil { | ||||
| logger.LogError(c, fmt.Sprintf("get bound doubao video channel failed: %v", bindErr)) | 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 | channel = bound | ||||
| selectGroup = usingGroup | selectGroup = usingGroup | ||||
| } | } | ||||
| @@ -126,6 +127,7 @@ func Distribute() func(c *gin.Context) { | |||||
| TokenGroup: usingGroup, | TokenGroup: usingGroup, | ||||
| Retry: common.GetPointer(0), | Retry: common.GetPointer(0), | ||||
| RequiredChannelType: requiredChannelType, | RequiredChannelType: requiredChannelType, | ||||
| AllowedChannelTypes: allowedChannelTypes, | |||||
| }) | }) | ||||
| if err != nil { | if err != nil { | ||||
| showGroup := usingGroup | showGroup := usingGroup | ||||
| @@ -341,21 +343,37 @@ func getModelRequest(c *gin.Context) (*ModelRequest, bool, error) { | |||||
| return &modelRequest, shouldSelectChannel, nil | 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 | 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) | group = strings.TrimSpace(group) | ||||
| modelName = strings.TrimSpace(modelName) | modelName = strings.TrimSpace(modelName) | ||||
| if group == "" || group == "auto" || 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 { | if c.GetInt("relay_mode") != relayconstant.RelayModeVideoSubmit { | ||||
| return false | return false | ||||
| } | } | ||||
| return service.HasDoubaoVideoChannelForGroupModel(group, modelName) | |||||
| return service.HasVideoAssetChannelForGroupModel(group, modelName, service.VideoAssetFamilySeedance) | |||||
| } | } | ||||
| func shouldPersistDoubaoVideoBinding(c *gin.Context) bool { | func shouldPersistDoubaoVideoBinding(c *gin.Context) bool { | ||||
| @@ -376,13 +394,13 @@ func persistDoubaoVideoBindingIfNeeded(userId int, group string, ch *model.Chann | |||||
| if group == "" || group == "auto" { | if group == "" || group == "auto" { | ||||
| return nil | return nil | ||||
| } | } | ||||
| if ch == nil || !service.IsDoubaoAssetChannelType(ch.Type) { | |||||
| if ch == nil { | |||||
| return nil | return nil | ||||
| } | } | ||||
| if !service.IsUsableDoubaoAssetChannel(ch, group) { | |||||
| if !service.IsUsableVideoAssetChannelForFamily(ch, group, "", service.VideoAssetFamilySeedance) { | |||||
| return nil | 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 { | 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/constant" | ||||
| "github.com/QuantumNous/new-api/model" | "github.com/QuantumNous/new-api/model" | ||||
| relayconstant "github.com/QuantumNous/new-api/relay/constant" | relayconstant "github.com/QuantumNous/new-api/relay/constant" | ||||
| "github.com/QuantumNous/new-api/service" | |||||
| "github.com/gin-gonic/gin" | "github.com/gin-gonic/gin" | ||||
| "github.com/stretchr/testify/assert" | "github.com/stretchr/testify/assert" | ||||
| "github.com/stretchr/testify/require" | "github.com/stretchr/testify/require" | ||||
| @@ -94,7 +95,7 @@ func tianyiyunVideoChannelForBindingTest(id int, group string) *model.Channel { | |||||
| return channel | return channel | ||||
| } | } | ||||
| func TestPreferBoundDoubaoVideoChannel(t *testing.T) { | |||||
| func TestPreferBoundSeedanceVideoChannel(t *testing.T) { | |||||
| db := setupDoubaoVideoBindingDB(t) | db := setupDoubaoVideoBindingDB(t) | ||||
| channel := doubaoVideoChannelForBindingTest(77, "default") | channel := doubaoVideoChannelForBindingTest(77, "default") | ||||
| priority := int64(channel.Id) | priority := int64(channel.Id) | ||||
| @@ -107,7 +108,7 @@ func TestPreferBoundDoubaoVideoChannel(t *testing.T) { | |||||
| Weight: 10, | Weight: 10, | ||||
| }).Error) | }).Error) | ||||
| assert.True(t, shouldUseBoundDoubaoVideoChannel(nil, channel, "default", "seedance-2")) | |||||
| assert.True(t, shouldUseBoundSeedanceVideoChannel(nil, channel, "default", "seedance-2")) | |||||
| } | } | ||||
| func TestPersistDoubaoVideoBindingSkipsNonDoubaoWithoutDB(t *testing.T) { | func TestPersistDoubaoVideoBindingSkipsNonDoubaoWithoutDB(t *testing.T) { | ||||
| @@ -163,14 +164,33 @@ func TestPersistDoubaoVideoBindingWritesTianyiYunBindingByType(t *testing.T) { | |||||
| assert.Equal(t, channel.Id, binding.ChannelId) | 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, _ := gin.CreateTestContext(nil) | ||||
| c.Set("relay_mode", relayconstant.RelayModeChatCompletions) | 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) | db := setupDoubaoVideoBindingDB(t) | ||||
| channel := doubaoVideoChannelForBindingTest(101, "default") | channel := doubaoVideoChannelForBindingTest(101, "default") | ||||
| require.NoError(t, db.Create(channel).Error) | require.NoError(t, db.Create(channel).Error) | ||||
| @@ -184,15 +204,18 @@ func TestShouldCheckDoubaoVideoBindingPreferenceRequiresVideoSubmitRelayMode(t * | |||||
| c, _ := gin.CreateTestContext(nil) | c, _ := gin.CreateTestContext(nil) | ||||
| c.Set("relay_mode", relayconstant.RelayModeVideoSubmit) | 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, _ := gin.CreateTestContext(nil) | ||||
| c.Request = httptest.NewRequest(http.MethodPost, "/api/v3/contents/generations/tasks", strings.NewReader(`{}`)) | 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) { | 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 { | 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() | now := common.GetTimestamp() | ||||
| binding := UserAssetChannel{ | binding := UserAssetChannel{ | ||||
| UserId: userId, | UserId: userId, | ||||
| @@ -45,7 +49,7 @@ func BindUserAssetChannel(userId int, channelType int, group string, channelId i | |||||
| UpdatedAt: now, | UpdatedAt: now, | ||||
| } | } | ||||
| return DB.Clauses(clause.OnConflict{ | |||||
| return tx.Session(&gorm.Session{SkipDefaultTransaction: true}).Clauses(clause.OnConflict{ | |||||
| Columns: []clause.Column{ | Columns: []clause.Column{ | ||||
| {Name: "user_id"}, | {Name: "user_id"}, | ||||
| {Name: "channel_type"}, | {Name: "channel_type"}, | ||||
| @@ -58,6 +62,25 @@ func BindUserAssetChannel(userId int, channelType int, group string, channelId i | |||||
| }).Create(&binding).Error | }).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 { | func UnbindUserAssetChannel(userId int, channelType int, group string) error { | ||||
| return DB.Where("user_id = ? AND channel_type = ? AND "+commonGroupCol+" = ?", userId, channelType, group). | return DB.Where("user_id = ? AND channel_type = ? AND "+commonGroupCol+" = ?", userId, channelType, group). | ||||
| Delete(&UserAssetChannel{}).Error | Delete(&UserAssetChannel{}).Error | ||||
| @@ -71,6 +71,75 @@ func TestBindUserAssetChannel_CreateAndUpdate(t *testing.T) { | |||||
| assert.Equal(t, int64(1), count) | 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) { | func TestBindUserAssetChannel_IsolatesUserTypeGroup(t *testing.T) { | ||||
| setupUserAssetChannelDB(t) | setupUserAssetChannelDB(t) | ||||
| @@ -84,6 +84,30 @@ func TestBuildRequestBodyPreservesTianyiYunNativeFields(t *testing.T) { | |||||
| require.NotContains(t, string(data), `"seconds"`) | 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) { | func TestDoResponseRewritesIDToPublicTaskID(t *testing.T) { | ||||
| adaptor := &TaskAdaptor{} | adaptor := &TaskAdaptor{} | ||||
| w := httptest.NewRecorder() | w := httptest.NewRecorder() | ||||
| @@ -162,6 +162,39 @@ func TestBuildRequestBodyMapsCanonicalModelNameForAiping(t *testing.T) { | |||||
| require.Equal(t, "Kling-V2.6", got["model_name"]) | 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) { | func TestValidateRequestRejectsOmniModelOnTextOrImageRoutes(t *testing.T) { | ||||
| gin.SetMode(gin.TestMode) | gin.SetMode(gin.TestMode) | ||||
| c, _ := gin.CreateTestContext(nil) | c, _ := gin.CreateTestContext(nil) | ||||
| @@ -22,6 +22,7 @@ type RetryParam struct { | |||||
| Retry *int | Retry *int | ||||
| RequireMatrixUsageBilling bool | RequireMatrixUsageBilling bool | ||||
| RequiredChannelType int | RequiredChannelType int | ||||
| AllowedChannelTypes []int | |||||
| resetNextTry bool | resetNextTry bool | ||||
| } | } | ||||
| @@ -184,13 +185,16 @@ func CacheGetRandomSatisfiedChannel(param *RetryParam) (*model.Channel, string, | |||||
| } | } | ||||
| func matrixUsageBillingChannelFilter(param *RetryParam) model.ChannelFilter { | 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 nil | ||||
| } | } | ||||
| return func(channel *model.Channel) (bool, error) { | return func(channel *model.Channel) (bool, error) { | ||||
| if param.RequiredChannelType != 0 && channel.Type != param.RequiredChannelType { | if param.RequiredChannelType != 0 && channel.Type != param.RequiredChannelType { | ||||
| return false, nil | return false, nil | ||||
| } | } | ||||
| if !channelTypeAllowed(channel.Type, param.AllowedChannelTypes) { | |||||
| return false, nil | |||||
| } | |||||
| if param.RequireMatrixUsageBilling { | if param.RequireMatrixUsageBilling { | ||||
| mappedModel, err := helper.ResolvePricingModelMapping(param.ModelName, channel.GetModelMapping()) | mappedModel, err := helper.ResolvePricingModelMapping(param.ModelName, channel.GetModelMapping()) | ||||
| if err != nil { | 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 { | func matrixUsageBillingChannelUnavailableError() *types.NewAPIError { | ||||
| return types.NewErrorWithStatusCode( | return types.NewErrorWithStatusCode( | ||||
| fmt.Errorf("matrix_usage_billing_channel_unavailable"), | fmt.Errorf("matrix_usage_billing_channel_unavailable"), | ||||
| @@ -3,12 +3,19 @@ package service | |||||
| import ( | import ( | ||||
| "net/http" | "net/http" | ||||
| "net/http/httptest" | "net/http/httptest" | ||||
| "os" | |||||
| "strings" | |||||
| "testing" | "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/QuantumNous/new-api/setting" | ||||
| "github.com/gin-gonic/gin" | "github.com/gin-gonic/gin" | ||||
| "github.com/stretchr/testify/assert" | "github.com/stretchr/testify/assert" | ||||
| "github.com/stretchr/testify/require" | "github.com/stretchr/testify/require" | ||||
| "gorm.io/gorm" | |||||
| "gorm.io/gorm/logger" | |||||
| ) | ) | ||||
| func buildRetryContext(t *testing.T) *gin.Context { | func buildRetryContext(t *testing.T) *gin.Context { | ||||
| @@ -22,6 +29,81 @@ func buildRetryContext(t *testing.T) *gin.Context { | |||||
| return c | 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 --- | // --- RetryParam tests --- | ||||
| func TestRetryParam_IncreaseRetry(t *testing.T) { | func TestRetryParam_IncreaseRetry(t *testing.T) { | ||||
| @@ -77,6 +159,26 @@ func TestCacheGetRandomSatisfiedChannel_AutoGroup_NoAutoGroups(t *testing.T) { | |||||
| assert.Equal(t, "auto", group) | 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 --- | // --- GetUserAutoGroup tests --- | ||||
| func TestGetUserAutoGroup_Basic(t *testing.T) { | func TestGetUserAutoGroup_Basic(t *testing.T) { | ||||
| @@ -1,163 +1,47 @@ | |||||
| package service | package service | ||||
| import ( | import ( | ||||
| "errors" | |||||
| "strings" | "strings" | ||||
| "github.com/QuantumNous/new-api/common" | |||||
| "github.com/QuantumNous/new-api/constant" | "github.com/QuantumNous/new-api/constant" | ||||
| "github.com/QuantumNous/new-api/model" | "github.com/QuantumNous/new-api/model" | ||||
| "gorm.io/gorm" | |||||
| ) | ) | ||||
| const DoubaoAssetChannelPageSize = 50 | const DoubaoAssetChannelPageSize = 50 | ||||
| const DoubaoAssetChannelType = constant.ChannelTypeDoubaoVideoCompatibleAiping | const DoubaoAssetChannelType = constant.ChannelTypeDoubaoVideoCompatibleAiping | ||||
| var DoubaoAssetChannelTypes = []int{ | |||||
| constant.ChannelTypeDoubaoVideoCompatibleAiping, | |||||
| constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, | |||||
| } | |||||
| var DoubaoAssetChannelTypes = VideoAssetChannelTypesForFamily(VideoAssetFamilySeedance) | |||||
| func ResolveDoubaoAssetChannel(userId int, tokenGroup string) (*model.Channel, error) { | func ResolveDoubaoAssetChannel(userId int, tokenGroup string) (*model.Channel, error) { | ||||
| return ResolveDoubaoAssetChannelForModel(userId, tokenGroup, "") | return ResolveDoubaoAssetChannelForModel(userId, tokenGroup, "") | ||||
| } | } | ||||
| func ResolveDoubaoAssetChannelForModel(userId int, tokenGroup string, modelName string) (*model.Channel, error) { | 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) { | 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) { | 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 { | 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) { | 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 { | 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 { | 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 { | func DoubaoAssetChannelTypeForModel(modelName string) int { | ||||
| @@ -132,9 +132,9 @@ func TestResolveDoubaoAssetChannel_AutoMatchesAndPersistsWhenUnbound(t *testing. | |||||
| func TestResolveDoubaoAssetChannel_AutoMatchesTianyiYunAndPersistsByType(t *testing.T) { | func TestResolveDoubaoAssetChannel_AutoMatchesTianyiYunAndPersistsByType(t *testing.T) { | ||||
| db := setupDoubaoAssetChannelDB(t) | db := setupDoubaoAssetChannelDB(t) | ||||
| createDoubaoAssetChannelForTest(t, db, 1, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "default", "tianyiyun-key", common.ChannelStatusEnabled) | 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.NoError(t, err) | ||||
| require.NotNil(t, channel) | require.NotNil(t, channel) | ||||
| @@ -147,6 +147,19 @@ func TestResolveDoubaoAssetChannel_AutoMatchesTianyiYunAndPersistsByType(t *test | |||||
| assert.Equal(t, 1, binding.ChannelId) | 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) { | func TestResolveDoubaoAssetChannel_RebindsOpenAIMisbindToAiping(t *testing.T) { | ||||
| db := setupDoubaoAssetChannelDB(t) | db := setupDoubaoAssetChannelDB(t) | ||||
| createDoubaoAssetChannelForTest(t, db, 1, constant.ChannelTypeOpenAI, "default", "openai-key", common.ChannelStatusEnabled) | 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, "vip", "seedance-2", 2, true) | ||||
| createDoubaoAssetAbilityForTest(t, db, "default", "disabled-model", 1, false) | createDoubaoAssetAbilityForTest(t, db, "default", "disabled-model", 1, false) | ||||
| createDoubaoAssetAbilityForTest(t, db, "default", "seedance-2", 3, true) | 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", "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("default", "disabled-model")) | ||||
| assert.False(t, HasDoubaoVideoChannelForGroupModel("missing", "seedance-2")) | assert.False(t, HasDoubaoVideoChannelForGroupModel("missing", "seedance-2")) | ||||
| } | } | ||||
| @@ -1,102 +1,32 @@ | |||||
| package service | package service | ||||
| import ( | import ( | ||||
| "errors" | |||||
| "strings" | |||||
| "github.com/QuantumNous/new-api/common" | |||||
| "github.com/QuantumNous/new-api/constant" | "github.com/QuantumNous/new-api/constant" | ||||
| "github.com/QuantumNous/new-api/model" | "github.com/QuantumNous/new-api/model" | ||||
| "gorm.io/gorm" | |||||
| ) | ) | ||||
| const KlingAssetChannelType = constant.ChannelTypeKlingAiping | 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) { | 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) { | 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 { | 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) | 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) { | func TestIsUsableKlingAssetChannel_RejectsOtherChannelType(t *testing.T) { | ||||
| ch := &model.Channel{Type: constant.ChannelTypeDoubaoVideoCompatibleAiping, Key: "k", Status: common.ChannelStatusEnabled, Group: "default"} | ch := &model.Channel{Type: constant.ChannelTypeDoubaoVideoCompatibleAiping, Key: "k", Status: common.ChannelStatusEnabled, Group: "default"} | ||||
| assert.False(t, IsUsableKlingAssetChannel(ch, "default")) | assert.False(t, IsUsableKlingAssetChannel(ch, "default")) | ||||
| @@ -61,4 +85,4 @@ func TestIsUsableKlingAssetChannel_RejectsOtherChannelType(t *testing.T) { | |||||
| func TestIsUsableKlingAssetChannel_RejectsGroupMismatch(t *testing.T) { | func TestIsUsableKlingAssetChannel_RejectsGroupMismatch(t *testing.T) { | ||||
| ch := &model.Channel{Type: constant.ChannelTypeKlingAiping, Key: "k", Status: common.ChannelStatusEnabled, Group: "vip"} | ch := &model.Channel{Type: constant.ChannelTypeKlingAiping, Key: "k", Status: common.ChannelStatusEnabled, Group: "vip"} | ||||
| assert.False(t, IsUsableKlingAssetChannel(ch, "default")) | 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) | |||||
| } | |||||