diff --git a/common/api_type.go b/common/api_type.go index 8a8811f..dd96cfd 100644 --- a/common/api_type.go +++ b/common/api_type.go @@ -55,6 +55,8 @@ func ChannelType2APIType(channelType int) (int, bool) { apiType = constant.APITypeVolcEngine case constant.ChannelTypeDoubaoVideoCompatibleAiping: apiType = constant.APITypeVolcEngine + case constant.ChannelTypeDoubaoVideoCompatibleTianyiYun: + apiType = constant.APITypeVolcEngine case constant.ChannelTypeBaiduV2: apiType = constant.APITypeBaiduV2 case constant.ChannelTypeOpenRouter: diff --git a/common/endpoint_type.go b/common/endpoint_type.go index fb7a693..b7eb845 100644 --- a/common/endpoint_type.go +++ b/common/endpoint_type.go @@ -30,7 +30,7 @@ func GetEndpointTypesByChannelType(channelType int, modelName string) []constant endpointTypes = []constant.EndpointType{constant.EndpointTypeOpenAI, constant.EndpointTypeOpenAIResponse} case constant.ChannelTypeSora: endpointTypes = []constant.EndpointType{constant.EndpointTypeOpenAIVideo} - case constant.ChannelTypeDoubaoVideo, constant.ChannelTypeDoubaoVideoCompatibleAiping: + case constant.ChannelTypeDoubaoVideo, constant.ChannelTypeDoubaoVideoCompatibleAiping, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun: endpointTypes = []constant.EndpointType{constant.EndpointTypeDoubaoVideo} default: if IsOpenAIResponseOnlyModel(modelName) { diff --git a/common/tianyiyun_channel_test.go b/common/tianyiyun_channel_test.go new file mode 100644 index 0000000..8140273 --- /dev/null +++ b/common/tianyiyun_channel_test.go @@ -0,0 +1,21 @@ +package common + +import ( + "testing" + + "github.com/QuantumNous/new-api/constant" + "github.com/stretchr/testify/require" +) + +func TestTianyiYunChannelUsesVolcEngineAPIType(t *testing.T) { + apiType, ok := ChannelType2APIType(constant.ChannelTypeDoubaoVideoCompatibleTianyiYun) + + require.True(t, ok) + require.Equal(t, constant.APITypeVolcEngine, apiType) +} + +func TestTianyiYunChannelUsesDoubaoVideoEndpointType(t *testing.T) { + got := GetEndpointTypesByChannelType(constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "cdance2.0-0611") + + require.Equal(t, []constant.EndpointType{constant.EndpointTypeDoubaoVideo}, got) +} diff --git a/constant/channel.go b/constant/channel.go index d54666a..03409d8 100644 --- a/constant/channel.go +++ b/constant/channel.go @@ -1,63 +1,64 @@ package constant const ( - ChannelTypeUnknown = 0 - ChannelTypeOpenAI = 1 - ChannelTypeMidjourney = 2 - ChannelTypeAzure = 3 - ChannelTypeOllama = 4 - ChannelTypeMidjourneyPlus = 5 - ChannelTypeOpenAIMax = 6 - ChannelTypeOhMyGPT = 7 - ChannelTypeCustom = 8 - ChannelTypeAILS = 9 - ChannelTypeAIProxy = 10 - ChannelTypePaLM = 11 - ChannelTypeAPI2GPT = 12 - ChannelTypeAIGC2D = 13 - ChannelTypeAnthropic = 14 - ChannelTypeBaidu = 15 - ChannelTypeZhipu = 16 - ChannelTypeAli = 17 - ChannelTypeXunfei = 18 - ChannelType360 = 19 - ChannelTypeOpenRouter = 20 - ChannelTypeAIProxyLibrary = 21 - ChannelTypeFastGPT = 22 - ChannelTypeTencent = 23 - ChannelTypeGemini = 24 - ChannelTypeMoonshot = 25 - ChannelTypeZhipu_v4 = 26 - ChannelTypePerplexity = 27 - ChannelTypeLingYiWanWu = 31 - ChannelTypeAws = 33 - ChannelTypeCohere = 34 - ChannelTypeMiniMax = 35 - ChannelTypeSunoAPI = 36 - ChannelTypeDify = 37 - ChannelTypeJina = 38 - ChannelCloudflare = 39 - ChannelTypeSiliconFlow = 40 - ChannelTypeVertexAi = 41 - ChannelTypeMistral = 42 - ChannelTypeDeepSeek = 43 - ChannelTypeMokaAI = 44 - ChannelTypeVolcEngine = 45 - ChannelTypeBaiduV2 = 46 - ChannelTypeXinference = 47 - ChannelTypeXai = 48 - ChannelTypeCoze = 49 - ChannelTypeKling = 50 - ChannelTypeJimeng = 51 - ChannelTypeVidu = 52 - ChannelTypeSubmodel = 53 - ChannelTypeDoubaoVideo = 54 - ChannelTypeSora = 55 - ChannelTypeReplicate = 56 - ChannelTypeCodex = 57 - ChannelTypeDoubaoVideoCompatibleAiping = 58 - ChannelTypeKlingAiping = 59 - ChannelTypeDummy // this one is only for count, do not add any channel after this + ChannelTypeUnknown = 0 + ChannelTypeOpenAI = 1 + ChannelTypeMidjourney = 2 + ChannelTypeAzure = 3 + ChannelTypeOllama = 4 + ChannelTypeMidjourneyPlus = 5 + ChannelTypeOpenAIMax = 6 + ChannelTypeOhMyGPT = 7 + ChannelTypeCustom = 8 + ChannelTypeAILS = 9 + ChannelTypeAIProxy = 10 + ChannelTypePaLM = 11 + ChannelTypeAPI2GPT = 12 + ChannelTypeAIGC2D = 13 + ChannelTypeAnthropic = 14 + ChannelTypeBaidu = 15 + ChannelTypeZhipu = 16 + ChannelTypeAli = 17 + ChannelTypeXunfei = 18 + ChannelType360 = 19 + ChannelTypeOpenRouter = 20 + ChannelTypeAIProxyLibrary = 21 + ChannelTypeFastGPT = 22 + ChannelTypeTencent = 23 + ChannelTypeGemini = 24 + ChannelTypeMoonshot = 25 + ChannelTypeZhipu_v4 = 26 + ChannelTypePerplexity = 27 + ChannelTypeLingYiWanWu = 31 + ChannelTypeAws = 33 + ChannelTypeCohere = 34 + ChannelTypeMiniMax = 35 + ChannelTypeSunoAPI = 36 + ChannelTypeDify = 37 + ChannelTypeJina = 38 + ChannelCloudflare = 39 + ChannelTypeSiliconFlow = 40 + ChannelTypeVertexAi = 41 + ChannelTypeMistral = 42 + ChannelTypeDeepSeek = 43 + ChannelTypeMokaAI = 44 + ChannelTypeVolcEngine = 45 + ChannelTypeBaiduV2 = 46 + ChannelTypeXinference = 47 + ChannelTypeXai = 48 + ChannelTypeCoze = 49 + ChannelTypeKling = 50 + ChannelTypeJimeng = 51 + ChannelTypeVidu = 52 + ChannelTypeSubmodel = 53 + ChannelTypeDoubaoVideo = 54 + ChannelTypeSora = 55 + ChannelTypeReplicate = 56 + ChannelTypeCodex = 57 + ChannelTypeDoubaoVideoCompatibleAiping = 58 + ChannelTypeKlingAiping = 59 + ChannelTypeDoubaoVideoCompatibleTianyiYun = 60 + ChannelTypeDummy // this one is only for count, do not add any channel after this ) @@ -122,65 +123,67 @@ var ChannelBaseURLs = []string{ "https://chatgpt.com", //57 "", //58 "https://aiping.cn/api", //59 + "https://ai.ctaigw.cn", //60 } var ChannelTypeNames = map[int]string{ - ChannelTypeUnknown: "Unknown", - ChannelTypeOpenAI: "OpenAI", - ChannelTypeMidjourney: "Midjourney", - ChannelTypeAzure: "Azure", - ChannelTypeOllama: "Ollama", - ChannelTypeMidjourneyPlus: "MidjourneyPlus", - ChannelTypeOpenAIMax: "OpenAIMax", - ChannelTypeOhMyGPT: "OhMyGPT", - ChannelTypeCustom: "Custom", - ChannelTypeAILS: "AILS", - ChannelTypeAIProxy: "AIProxy", - ChannelTypePaLM: "PaLM", - ChannelTypeAPI2GPT: "API2GPT", - ChannelTypeAIGC2D: "AIGC2D", - ChannelTypeAnthropic: "Anthropic", - ChannelTypeBaidu: "Baidu", - ChannelTypeZhipu: "Zhipu", - ChannelTypeAli: "Ali", - ChannelTypeXunfei: "Xunfei", - ChannelType360: "360", - ChannelTypeOpenRouter: "OpenRouter", - ChannelTypeAIProxyLibrary: "AIProxyLibrary", - ChannelTypeFastGPT: "FastGPT", - ChannelTypeTencent: "Tencent", - ChannelTypeGemini: "Gemini", - ChannelTypeMoonshot: "Moonshot", - ChannelTypeZhipu_v4: "ZhipuV4", - ChannelTypePerplexity: "Perplexity", - ChannelTypeLingYiWanWu: "LingYiWanWu", - ChannelTypeAws: "AWS", - ChannelTypeCohere: "Cohere", - ChannelTypeMiniMax: "MiniMax", - ChannelTypeSunoAPI: "SunoAPI", - ChannelTypeDify: "Dify", - ChannelTypeJina: "Jina", - ChannelCloudflare: "Cloudflare", - ChannelTypeSiliconFlow: "SiliconFlow", - ChannelTypeVertexAi: "VertexAI", - ChannelTypeMistral: "Mistral", - ChannelTypeDeepSeek: "DeepSeek", - ChannelTypeMokaAI: "MokaAI", - ChannelTypeVolcEngine: "VolcEngine", - ChannelTypeBaiduV2: "BaiduV2", - ChannelTypeXinference: "Xinference", - ChannelTypeXai: "xAI", - ChannelTypeCoze: "Coze", - ChannelTypeKling: "Kling", - ChannelTypeJimeng: "Jimeng", - ChannelTypeVidu: "Vidu", - ChannelTypeSubmodel: "Submodel", - ChannelTypeDoubaoVideo: "DoubaoVideo", - ChannelTypeSora: "Sora", - ChannelTypeReplicate: "Replicate", - ChannelTypeCodex: "Codex", - ChannelTypeDoubaoVideoCompatibleAiping: "DoubaoVideoCompatibleAiping", - ChannelTypeKlingAiping: "KlingAiping", + ChannelTypeUnknown: "Unknown", + ChannelTypeOpenAI: "OpenAI", + ChannelTypeMidjourney: "Midjourney", + ChannelTypeAzure: "Azure", + ChannelTypeOllama: "Ollama", + ChannelTypeMidjourneyPlus: "MidjourneyPlus", + ChannelTypeOpenAIMax: "OpenAIMax", + ChannelTypeOhMyGPT: "OhMyGPT", + ChannelTypeCustom: "Custom", + ChannelTypeAILS: "AILS", + ChannelTypeAIProxy: "AIProxy", + ChannelTypePaLM: "PaLM", + ChannelTypeAPI2GPT: "API2GPT", + ChannelTypeAIGC2D: "AIGC2D", + ChannelTypeAnthropic: "Anthropic", + ChannelTypeBaidu: "Baidu", + ChannelTypeZhipu: "Zhipu", + ChannelTypeAli: "Ali", + ChannelTypeXunfei: "Xunfei", + ChannelType360: "360", + ChannelTypeOpenRouter: "OpenRouter", + ChannelTypeAIProxyLibrary: "AIProxyLibrary", + ChannelTypeFastGPT: "FastGPT", + ChannelTypeTencent: "Tencent", + ChannelTypeGemini: "Gemini", + ChannelTypeMoonshot: "Moonshot", + ChannelTypeZhipu_v4: "ZhipuV4", + ChannelTypePerplexity: "Perplexity", + ChannelTypeLingYiWanWu: "LingYiWanWu", + ChannelTypeAws: "AWS", + ChannelTypeCohere: "Cohere", + ChannelTypeMiniMax: "MiniMax", + ChannelTypeSunoAPI: "SunoAPI", + ChannelTypeDify: "Dify", + ChannelTypeJina: "Jina", + ChannelCloudflare: "Cloudflare", + ChannelTypeSiliconFlow: "SiliconFlow", + ChannelTypeVertexAi: "VertexAI", + ChannelTypeMistral: "Mistral", + ChannelTypeDeepSeek: "DeepSeek", + ChannelTypeMokaAI: "MokaAI", + ChannelTypeVolcEngine: "VolcEngine", + ChannelTypeBaiduV2: "BaiduV2", + ChannelTypeXinference: "Xinference", + ChannelTypeXai: "xAI", + ChannelTypeCoze: "Coze", + ChannelTypeKling: "Kling", + ChannelTypeJimeng: "Jimeng", + ChannelTypeVidu: "Vidu", + ChannelTypeSubmodel: "Submodel", + ChannelTypeDoubaoVideo: "DoubaoVideo", + ChannelTypeSora: "Sora", + ChannelTypeReplicate: "Replicate", + ChannelTypeCodex: "Codex", + ChannelTypeDoubaoVideoCompatibleAiping: "DoubaoVideoCompatibleAiping", + ChannelTypeKlingAiping: "KlingAiping", + ChannelTypeDoubaoVideoCompatibleTianyiYun: "DoubaoVideoCompatibleTianyiYun", } func GetChannelTypeName(channelType int) string { diff --git a/controller/channel-test.go b/controller/channel-test.go index ab12132..2b7aa5d 100644 --- a/controller/channel-test.go +++ b/controller/channel-test.go @@ -65,6 +65,8 @@ func testChannel(channel *model.Channel, testModel string, endpointType string, constant.ChannelTypeKling, constant.ChannelTypeJimeng, constant.ChannelTypeDoubaoVideo, + constant.ChannelTypeDoubaoVideoCompatibleAiping, + constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, constant.ChannelTypeVidu, } if lo.Contains(unsupportedTestChannelTypes, channel.Type) { diff --git a/controller/channel_test_tianyiyun_test.go b/controller/channel_test_tianyiyun_test.go new file mode 100644 index 0000000..31b9b9d --- /dev/null +++ b/controller/channel_test_tianyiyun_test.go @@ -0,0 +1,31 @@ +package controller + +import ( + "strings" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/model" + "github.com/stretchr/testify/require" +) + +func TestTestChannelRejectsAsyncVideoChannels(t *testing.T) { + channels := []*model.Channel{ + {Type: constant.ChannelTypeDoubaoVideoCompatibleAiping, Models: "doubao-seedance-2-0-260128", Status: common.ChannelStatusEnabled}, + {Type: constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, Models: "cdance2.0-0611", Status: common.ChannelStatusEnabled}, + } + + for _, channel := range channels { + result := testChannel(channel, "", "", false) + + require.Error(t, result.localErr) + require.True(t, strings.Contains(result.localErr.Error(), "channel test is not supported")) + } +} + +func TestRequiredTaskChannelTypeForTianyiYunSeedanceModel(t *testing.T) { + c := newControllerJSONContext(t, "/api/v3/contents/generations/tasks", `{"model":"cdance2.0-0611"}`) + + require.Equal(t, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, requiredTaskChannelTypeForRequest(c)) +} diff --git a/controller/doubao_aiping_video.go b/controller/doubao_aiping_video.go index 7d3587a..515d238 100644 --- a/controller/doubao_aiping_video.go +++ b/controller/doubao_aiping_video.go @@ -10,6 +10,7 @@ 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" ) @@ -82,11 +83,9 @@ func AipingNativeVideoSubmit(c *gin.Context) { }) return } - c.Set("channel_type", constant.ChannelTypeDoubaoVideoCompatibleAiping) if relayInfo.ChannelMeta == nil { relayInfo.ChannelMeta = &relaycommon.ChannelMeta{} } - relayInfo.ChannelType = constant.ChannelTypeDoubaoVideoCompatibleAiping if err := prepareAipingNativeVideoSubmit(c, relayInfo); err != nil { c.JSON(http.StatusBadRequest, &dto.TaskError{ Code: "invalid_request", @@ -95,6 +94,9 @@ func AipingNativeVideoSubmit(c *gin.Context) { }) return } + channelType := service.DoubaoAssetChannelTypeForModel(relayInfo.OriginModelName) + c.Set("channel_type", channelType) + relayInfo.ChannelType = channelType relayTaskWithInfo(c, relayInfo) } diff --git a/controller/relay.go b/controller/relay.go index 40bc7e6..bbe50bb 100644 --- a/controller/relay.go +++ b/controller/relay.go @@ -317,11 +317,11 @@ func persistTaskVideoBindingIfNeeded(userId int, group string, channel *model.Ch return nil } switch channel.Type { - case service.DoubaoAssetChannelType: + case service.DoubaoAssetChannelType, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun: if !service.IsUsableDoubaoAssetChannel(channel, group) { return nil } - return model.BindUserAssetChannel(userId, service.DoubaoAssetChannelType, group, channel.Id) + return model.BindUserAssetChannel(userId, channel.Type, group, channel.Id) case constant.ChannelTypeKlingAiping: if !service.IsUsableKlingAssetChannel(channel, group) { return nil @@ -368,7 +368,7 @@ 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 constant.ChannelTypeDoubaoVideoCompatibleAiping + return service.DoubaoAssetChannelTypeForModel(taskModelNameForRequiredChannelType(c)) } if isKlingAipingNativePath(c.Request.URL.Path) { return constant.ChannelTypeKlingAiping @@ -376,6 +376,27 @@ func requiredTaskChannelTypeForRequest(c *gin.Context) int { 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"` + } + if err := common.Unmarshal(body, &raw); err != nil { + return "" + } + return raw.Model +} + func shouldRetry(c *gin.Context, openaiErr *types.NewAPIError, retryTimes int) bool { if openaiErr == nil { return false diff --git a/middleware/distributor.go b/middleware/distributor.go index 64b2c8d..a4516ee 100644 --- a/middleware/distributor.go +++ b/middleware/distributor.go @@ -65,7 +65,7 @@ func Distribute() func(c *gin.Context) { } var selectGroup string usingGroup := common.GetContextKeyString(c, constant.ContextKeyUsingGroup) - requiredChannelType := requiredChannelTypeForRequest(c) + requiredChannelType := requiredChannelTypeForRequest(c, modelRequest.Model) // check path is /pg/chat/completions if strings.HasPrefix(c.Request.URL.Path, "/pg/chat/completions") { playgroundRequest := &dto.PlayGroundRequest{} @@ -110,7 +110,7 @@ func Distribute() func(c *gin.Context) { } if channel == nil && shouldCheckDoubaoVideoBindingPreference(c, usingGroup, modelRequest.Model) { - bound, bindErr := service.GetBoundDoubaoAssetChannel(c.GetInt("id"), usingGroup) + 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) { @@ -348,9 +348,9 @@ func shouldUseBoundDoubaoVideoChannel(_ *gin.Context, ch *model.Channel, group s return model.IsChannelEnabledForGroupModel(group, modelName, ch.Id) } -func requiredChannelTypeForRequest(c *gin.Context) int { +func requiredChannelTypeForRequest(c *gin.Context, modelName string) int { if strings.HasPrefix(c.Request.URL.Path, "/api/v3/contents/generations/tasks") { - return constant.ChannelTypeDoubaoVideoCompatibleAiping + return service.DoubaoAssetChannelTypeForModel(modelName) } return 0 } @@ -376,13 +376,13 @@ func persistDoubaoVideoBindingIfNeeded(userId int, group string, ch *model.Chann if group == "" || group == "auto" { return nil } - if ch == nil || ch.Type != service.DoubaoAssetChannelType { + if ch == nil || !service.IsDoubaoAssetChannelType(ch.Type) { return nil } if !service.IsUsableDoubaoAssetChannel(ch, group) { return nil } - return model.BindUserAssetChannel(userId, service.DoubaoAssetChannelType, group, ch.Id) + return model.BindUserAssetChannel(userId, ch.Type, group, ch.Id) } func SetupContextForSelectedChannel(c *gin.Context, channel *model.Channel, modelName string) *types.NewAPIError { diff --git a/middleware/doubao_asset_binding_test.go b/middleware/doubao_asset_binding_test.go index d2fe6dd..39f3303 100644 --- a/middleware/doubao_asset_binding_test.go +++ b/middleware/doubao_asset_binding_test.go @@ -2,6 +2,8 @@ package middleware import ( "fmt" + "net/http" + "net/http/httptest" "os" "strings" "testing" @@ -85,6 +87,13 @@ func doubaoVideoChannelForBindingTest(id int, group string) *model.Channel { } } +func tianyiyunVideoChannelForBindingTest(id int, group string) *model.Channel { + channel := doubaoVideoChannelForBindingTest(id, group) + channel.Type = constant.ChannelTypeDoubaoVideoCompatibleTianyiYun + channel.Models = "cdance2.0-0611" + return channel +} + func TestPreferBoundDoubaoVideoChannel(t *testing.T) { db := setupDoubaoVideoBindingDB(t) channel := doubaoVideoChannelForBindingTest(77, "default") @@ -142,6 +151,18 @@ func TestPersistDoubaoVideoBindingWritesBinding(t *testing.T) { assert.Equal(t, channel.Id, binding.ChannelId) } +func TestPersistDoubaoVideoBindingWritesTianyiYunBindingByType(t *testing.T) { + setupDoubaoVideoBindingDB(t) + channel := tianyiyunVideoChannelForBindingTest(199, "vip") + + require.NoError(t, persistDoubaoVideoBindingIfNeeded(10, "vip", channel)) + + binding, err := model.GetUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "vip") + require.NoError(t, err) + require.NotNil(t, binding) + assert.Equal(t, channel.Id, binding.ChannelId) +} + func TestShouldCheckDoubaoVideoBindingPreferenceSkipsNonVideoRelayMode(t *testing.T) { c, _ := gin.CreateTestContext(nil) c.Set("relay_mode", relayconstant.RelayModeChatCompletions) @@ -166,6 +187,14 @@ func TestShouldCheckDoubaoVideoBindingPreferenceRequiresVideoSubmitRelayMode(t * assert.True(t, shouldCheckDoubaoVideoBindingPreference(c, "default", "seedance-2")) } +func TestRequiredChannelTypeForSeedanceTasksUsesModelSpecificChannel(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")) +} + func TestShouldPersistDoubaoVideoBindingRequiresVideoSubmitRelayMode(t *testing.T) { c, _ := gin.CreateTestContext(nil) c.Set("relay_mode", relayconstant.RelayModeChatCompletions) diff --git a/relay/channel/task/doubao_tianyiyun/adaptor.go b/relay/channel/task/doubao_tianyiyun/adaptor.go new file mode 100644 index 0000000..81c17fb --- /dev/null +++ b/relay/channel/task/doubao_tianyiyun/adaptor.go @@ -0,0 +1,308 @@ +package doubao_tianyiyun + +import ( + "bytes" + "fmt" + "io" + "net/http" + "strconv" + "strings" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/dto" + "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/relay/channel" + taskcommon "github.com/QuantumNous/new-api/relay/channel/task/taskcommon" + relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/QuantumNous/new-api/service" + relaytypes "github.com/QuantumNous/new-api/types" + "github.com/gin-gonic/gin" + "github.com/pkg/errors" + "github.com/samber/lo" +) + +type ContentItem struct { + Type string `json:"type,omitempty"` + Text string `json:"text,omitempty"` + ImageURL *MediaURL `json:"image_url,omitempty"` + VideoURL *MediaURL `json:"video_url,omitempty"` + AudioURL *MediaURL `json:"audio_url,omitempty"` + DraftTask *DraftTask `json:"draft_task,omitempty"` + Role string `json:"role,omitempty"` +} + +type MediaURL struct { + URL string `json:"url,omitempty"` +} + +type DraftTask struct { + ID string `json:"id,omitempty"` +} + +type requestPayload struct { + Model string `json:"model"` + Content []ContentItem `json:"content,omitempty"` + CallbackURL string `json:"callback_url,omitempty"` + ReturnLastFrame *dto.BoolValue `json:"return_last_frame,omitempty"` + ServiceTier string `json:"service_tier,omitempty"` + ExecutionExpiresAfter *dto.IntValue `json:"execution_expires_after,omitempty"` + GenerateAudio *dto.BoolValue `json:"generate_audio,omitempty"` + Draft *dto.BoolValue `json:"draft,omitempty"` + Tools []struct { + Type string `json:"type,omitempty"` + } `json:"tools,omitempty"` + Resolution string `json:"resolution,omitempty"` + Ratio string `json:"ratio,omitempty"` + Duration *dto.IntValue `json:"duration,omitempty"` + Frames *dto.IntValue `json:"frames,omitempty"` + Seed *dto.IntValue `json:"seed,omitempty"` + CameraFixed *dto.BoolValue `json:"camera_fixed,omitempty"` + Watermark *dto.BoolValue `json:"watermark,omitempty"` + SafetyIdentifier string `json:"safety_identifier,omitempty"` + Priority *dto.IntValue `json:"priority,omitempty"` +} + +type responsePayload struct { + ID string `json:"id"` +} + +type responseTask struct { + ID string `json:"id"` + Model string `json:"model"` + Status string `json:"status"` + Content struct { + VideoURL string `json:"video_url"` + } `json:"content"` + Usage struct { + CompletionTokens int `json:"completion_tokens"` + TotalTokens int `json:"total_tokens"` + } `json:"usage"` + Error struct { + Code string `json:"code"` + Message string `json:"message"` + } `json:"error"` + CreatedAt int64 `json:"created_at"` + UpdatedAt int64 `json:"updated_at"` +} + +type TaskAdaptor struct { + taskcommon.BaseBilling + ChannelType int + apiKey string + baseURL string +} + +func (a *TaskAdaptor) Init(info *relaycommon.RelayInfo) { + a.ChannelType = info.ChannelType + a.baseURL = strings.TrimRight(info.ChannelBaseUrl, "/") + a.apiKey = info.ApiKey +} + +func (a *TaskAdaptor) ValidateRequestAndSetAction(c *gin.Context, info *relaycommon.RelayInfo) (taskErr *dto.TaskError) { + if _, err := relaycommon.GetTaskRequest(c); err == nil { + info.Action = constant.TaskActionGenerate + return nil + } + return relaycommon.ValidateBasicTaskRequest(c, info, constant.TaskActionGenerate) +} + +func (a *TaskAdaptor) BuildRequestURL(info *relaycommon.RelayInfo) (string, error) { + baseURL := a.baseURL + if strings.TrimSpace(baseURL) == "" && info != nil { + baseURL = strings.TrimRight(info.ChannelBaseUrl, "/") + } + return fmt.Sprintf("%s/v1/contents/generations/tasks", baseURL), nil +} + +func (a *TaskAdaptor) BuildRequestHeader(c *gin.Context, req *http.Request, info *relaycommon.RelayInfo) error { + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "application/json") + req.Header.Set("Authorization", "Bearer "+a.apiKey) + return nil +} + +func (a *TaskAdaptor) BuildRequestBody(c *gin.Context, info *relaycommon.RelayInfo) (io.Reader, error) { + req, err := relaycommon.GetTaskRequest(c) + if err != nil { + return nil, err + } + body, err := a.convertToRequestPayload(&req) + if err != nil { + return nil, errors.Wrap(err, "convert request payload failed") + } + if info.IsModelMapped { + body.Model = info.UpstreamModelName + } else { + info.UpstreamModelName = body.Model + } + data, err := common.Marshal(body) + if err != nil { + return nil, err + } + return bytes.NewReader(data), nil +} + +func (a *TaskAdaptor) DoRequest(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (*http.Response, error) { + return channel.DoTaskApiRequest(a, c, info, requestBody) +} + +func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (string, []byte, *dto.TaskError) { + responseBody, err := io.ReadAll(resp.Body) + if err != nil { + return "", nil, service.TaskErrorWrapper(err, "read_response_body_failed", http.StatusInternalServerError) + } + _ = resp.Body.Close() + + var dResp responsePayload + if err := common.Unmarshal(responseBody, &dResp); err != nil { + return "", nil, service.TaskErrorWrapper(errors.Wrapf(err, "body: %s", responseBody), "unmarshal_response_body_failed", http.StatusInternalServerError) + } + if strings.TrimSpace(dResp.ID) == "" { + return "", nil, service.TaskErrorWrapper(fmt.Errorf("task_id is empty"), "invalid_response", http.StatusInternalServerError) + } + + clientPayload := map[string]any{} + if err := common.Unmarshal(responseBody, &clientPayload); err != nil { + return "", nil, service.TaskErrorWrapper(err, "unmarshal_response_body_failed", http.StatusInternalServerError) + } + clientPayload = relaytypes.CloneMapAny(clientPayload) + clientPayload["id"] = info.PublicTaskID + if _, ok := clientPayload["created_at"]; !ok { + clientPayload["created_at"] = time.Now().Unix() + } + if _, ok := clientPayload["model"]; !ok { + clientPayload["model"] = info.OriginModelName + } + c.JSON(http.StatusOK, clientPayload) + return dResp.ID, responseBody, nil +} + +func (a *TaskAdaptor) FetchTask(baseUrl, key string, body map[string]any, proxy string) (*http.Response, error) { + req, err := a.buildFetchRequest(baseUrl, key, body) + if err != nil { + return nil, err + } + client, err := service.GetHttpClientWithProxy(proxy) + if err != nil { + return nil, fmt.Errorf("new proxy http client failed: %w", err) + } + return client.Do(req) +} + +func (a *TaskAdaptor) buildFetchRequest(baseUrl, key string, body map[string]any) (*http.Request, error) { + taskID, ok := body["task_id"].(string) + if !ok || strings.TrimSpace(taskID) == "" { + return nil, fmt.Errorf("invalid task_id") + } + uri := fmt.Sprintf("%s/v1/contents/generations/tasks/%s", strings.TrimRight(baseUrl, "/"), taskID) + req, err := http.NewRequest(http.MethodGet, uri, nil) + if err != nil { + return nil, err + } + req.Header.Set("Accept", "application/json") + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+key) + return req, nil +} + +func (a *TaskAdaptor) GetModelList() []string { + return ModelList +} + +func (a *TaskAdaptor) GetChannelName() string { + return ChannelName +} + +func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (*requestPayload, error) { + r := requestPayload{ + Model: req.Model, + Content: []ContentItem{}, + } + if req.HasImage() { + for _, imgURL := range req.Images { + r.Content = append(r.Content, ContentItem{ + Type: "image_url", + ImageURL: &MediaURL{URL: imgURL}, + }) + } + } + if err := taskcommon.UnmarshalMetadata(req.Metadata, &r); err != nil { + return nil, errors.Wrap(err, "unmarshal metadata failed") + } + if sec, _ := strconv.Atoi(req.Seconds); sec > 0 && r.Duration == nil { + r.Duration = lo.ToPtr(dto.IntValue(sec)) + } + hasTextContent := lo.SomeBy(r.Content, func(c ContentItem) bool { + return c.Type == "text" && strings.TrimSpace(c.Text) != "" + }) + if !hasTextContent && strings.TrimSpace(req.Prompt) != "" { + r.Content = append(r.Content, ContentItem{Type: "text", Text: req.Prompt}) + } + return &r, nil +} + +func (a *TaskAdaptor) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, error) { + resTask := responseTask{} + if err := common.Unmarshal(respBody, &resTask); err != nil { + return nil, errors.Wrap(err, "unmarshal task result failed") + } + + taskResult := relaycommon.TaskInfo{Code: 0} + switch resTask.Status { + case "pending", "queued": + taskResult.Status = model.TaskStatusQueued + taskResult.Progress = "10%" + case "processing", "running": + taskResult.Status = model.TaskStatusInProgress + taskResult.Progress = "50%" + case "succeeded", "success": + taskResult.Status = model.TaskStatusSuccess + taskResult.Progress = "100%" + taskResult.Url = resTask.Content.VideoURL + taskResult.CompletionTokens = resTask.Usage.CompletionTokens + taskResult.TotalTokens = resTask.Usage.TotalTokens + case "failed", "expired", "cancelled": + taskResult.Status = model.TaskStatusFailure + taskResult.Progress = "100%" + taskResult.Reason = resTask.Error.Message + if taskResult.Reason == "" { + taskResult.Reason = "task " + resTask.Status + } + default: + taskResult.Status = model.TaskStatusInProgress + taskResult.Progress = "30%" + } + return &taskResult, nil +} + +func (a *TaskAdaptor) ConvertToOpenAIVideo(originTask *model.Task) ([]byte, error) { + var dResp responseTask + if err := common.Unmarshal(originTask.Data, &dResp); err != nil { + return nil, errors.Wrap(err, "unmarshal tianyiyun doubao task data failed") + } + + openAIVideo := dto.NewOpenAIVideo() + openAIVideo.ID = originTask.TaskID + openAIVideo.TaskID = originTask.TaskID + openAIVideo.Status = originTask.Status.ToVideoStatus() + openAIVideo.SetProgressStr(originTask.Progress) + openAIVideo.SetMetadata("url", dResp.Content.VideoURL) + openAIVideo.CreatedAt = originTask.CreatedAt + openAIVideo.CompletedAt = originTask.UpdatedAt + openAIVideo.Model = originTask.Properties.OriginModelName + + if dResp.Status == "failed" { + message := dResp.Error.Message + if message == "" { + message = "task failed" + } + code := dResp.Error.Code + if code == "" { + code = "failed" + } + openAIVideo.Error = &dto.OpenAIVideoError{Message: message, Code: code} + } + return common.Marshal(openAIVideo) +} diff --git a/relay/channel/task/doubao_tianyiyun/adaptor_test.go b/relay/channel/task/doubao_tianyiyun/adaptor_test.go new file mode 100644 index 0000000..559c413 --- /dev/null +++ b/relay/channel/task/doubao_tianyiyun/adaptor_test.go @@ -0,0 +1,182 @@ +package doubao_tianyiyun + +import ( + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/model" + relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func TestBuildRequestURLUsesTianyiYunEndpoint(t *testing.T) { + adaptor := &TaskAdaptor{} + adaptor.Init(&relaycommon.RelayInfo{ + ChannelMeta: &relaycommon.ChannelMeta{ + ChannelType: constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, + ChannelBaseUrl: "https://ai.ctaigw.cn", + ApiKey: "sk-test", + }, + }) + + got, err := adaptor.BuildRequestURL(&relaycommon.RelayInfo{ + ChannelMeta: &relaycommon.ChannelMeta{ChannelBaseUrl: "https://ai.ctaigw.cn"}, + }) + + require.NoError(t, err) + require.Equal(t, "https://ai.ctaigw.cn/v1/contents/generations/tasks", got) +} + +func TestFetchTaskUsesTianyiYunEndpoint(t *testing.T) { + adaptor := &TaskAdaptor{} + req, err := adaptor.buildFetchRequest("https://ai.ctaigw.cn", "sk-test", map[string]any{"task_id": "task-123"}) + + require.NoError(t, err) + require.Equal(t, http.MethodGet, req.Method) + require.Equal(t, "https://ai.ctaigw.cn/v1/contents/generations/tasks/task-123", req.URL.String()) + require.Equal(t, "Bearer sk-test", req.Header.Get("Authorization")) + require.Equal(t, "application/json", req.Header.Get("Accept")) +} + +func TestGetModelListIncludesTianyiYunSeedanceModels(t *testing.T) { + models := (&TaskAdaptor{}).GetModelList() + + require.Contains(t, models, "cdance2.0-0611") + require.Contains(t, models, "cdance2.0-fast-0611") + require.Equal(t, "DoubaoVideoCompatibleTianyiYun", (&TaskAdaptor{}).GetChannelName()) +} + +func TestBuildRequestBodyPreservesTianyiYunNativeFields(t *testing.T) { + adaptor := &TaskAdaptor{} + c := newTianyiYunTaskRequestContext(t, `{ + "model":"cdance2.0-0611", + "prompt":"write a clean product video", + "metadata":{ + "content":[ + {"type":"image_url","image_url":{"url":"https://example.test/input.png"},"role":"reference_image"} + ], + "ratio":"16:9", + "duration":5, + "watermark":false + } + }`) + info := &relaycommon.RelayInfo{ + ChannelMeta: &relaycommon.ChannelMeta{UpstreamModelName: "cdance2.0-0611"}, + } + + 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.Contains(t, string(data), `"image_url":{"url":"https://example.test/input.png"}`) + require.Contains(t, string(data), `"text":"write a clean product video"`) + require.Contains(t, string(data), `"ratio":"16:9"`) + require.Contains(t, string(data), `"duration":5`) + require.Contains(t, string(data), `"watermark":false`) + require.NotContains(t, string(data), `"seconds"`) +} + +func TestDoResponseRewritesIDToPublicTaskID(t *testing.T) { + adaptor := &TaskAdaptor{} + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + info := &relaycommon.RelayInfo{ + OriginModelName: "cdance2.0-0611", + TaskRelayInfo: &relaycommon.TaskRelayInfo{PublicTaskID: "task_public"}, + } + resp := &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader(`{ + "id":"task-upstream", + "status":"queued", + "created_at":1781496040, + "model":"cdance2.0-0611" + }`)), + } + + upstreamID, data, taskErr := adaptor.DoResponse(c, resp, info) + + require.Nil(t, taskErr) + require.Equal(t, "task-upstream", upstreamID) + require.Contains(t, string(data), `"id":"task-upstream"`) + require.Equal(t, http.StatusOK, w.Code) + require.Contains(t, w.Body.String(), `"id":"task_public"`) + require.Contains(t, w.Body.String(), `"status":"queued"`) + require.Contains(t, w.Body.String(), `"model":"cdance2.0-0611"`) +} + +func TestParseTaskResultMapsTianyiYunStatuses(t *testing.T) { + adaptor := &TaskAdaptor{} + cases := []struct { + name string + body string + wantStatus model.TaskStatus + wantURL string + wantReason string + wantUsage bool + }{ + { + name: "success", + body: `{"id":"task-ok","status":"success","content":{"video_url":"https://example.test/out.mp4"},"usage":{"completion_tokens":7,"total_tokens":11}}`, + wantStatus: model.TaskStatusSuccess, + wantURL: "https://example.test/out.mp4", + wantUsage: true, + }, + { + name: "succeeded", + body: `{"id":"task-ok","status":"succeeded","content":{"video_url":"https://example.test/out.mp4"}}`, + wantStatus: model.TaskStatusSuccess, + wantURL: "https://example.test/out.mp4", + }, + { + name: "running", + body: `{"id":"task-run","status":"running"}`, + wantStatus: model.TaskStatusInProgress, + }, + { + name: "failed", + body: `{"id":"task-fail","status":"failed","error":{"code":"BadRequest","message":"bad prompt"}}`, + wantStatus: model.TaskStatusFailure, + wantReason: "bad prompt", + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + got, err := adaptor.ParseTaskResult([]byte(tc.body)) + + require.NoError(t, err) + require.Equal(t, string(tc.wantStatus), got.Status) + if tc.wantURL != "" { + require.Equal(t, tc.wantURL, got.Url) + } + if tc.wantUsage { + require.Equal(t, 7, got.CompletionTokens) + require.Equal(t, 11, got.TotalTokens) + } + if tc.wantReason != "" { + require.Equal(t, tc.wantReason, got.Reason) + } + }) + } +} + +func newTianyiYunTaskRequestContext(t *testing.T, body string) *gin.Context { + t.Helper() + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + c.Request = httptest.NewRequest(http.MethodPost, "/api/v3/contents/generations/tasks", strings.NewReader(body)) + c.Request.Header.Set("Content-Type", "application/json") + var req relaycommon.TaskSubmitReq + require.NoError(t, common.Unmarshal([]byte(body), &req)) + relaycommon.StoreTaskRequest(c, &relaycommon.RelayInfo{}, constant.TaskActionGenerate, req) + return c +} diff --git a/relay/channel/task/doubao_tianyiyun/constants.go b/relay/channel/task/doubao_tianyiyun/constants.go new file mode 100644 index 0000000..e09f5d8 --- /dev/null +++ b/relay/channel/task/doubao_tianyiyun/constants.go @@ -0,0 +1,8 @@ +package doubao_tianyiyun + +var ModelList = []string{ + "cdance2.0-0611", + "cdance2.0-fast-0611", +} + +var ChannelName = "DoubaoVideoCompatibleTianyiYun" diff --git a/relay/helper/matrix_usage_capability.go b/relay/helper/matrix_usage_capability.go index 7317cdb..9fcd1f3 100644 --- a/relay/helper/matrix_usage_capability.go +++ b/relay/helper/matrix_usage_capability.go @@ -21,6 +21,8 @@ var matrixUsageCapabilities = []MatrixUsageCapability{ {BillingModelName: "doubao-seedance-2-0-fast-260128", UpstreamModelName: "doubao-seedance-2-0-fast-260128", ChannelType: constant.ChannelTypeDoubaoVideo}, {BillingModelName: "doubao-seedance-2-0-260128", UpstreamModelName: "doubao-seedance-2-0-260128", ChannelType: constant.ChannelTypeDoubaoVideoCompatibleAiping}, {BillingModelName: "doubao-seedance-2-0-fast-260128", UpstreamModelName: "doubao-seedance-2-0-fast-260128", ChannelType: constant.ChannelTypeDoubaoVideoCompatibleAiping}, + {BillingModelName: "cdance2.0-0611", UpstreamModelName: "cdance2.0-0611", ChannelType: constant.ChannelTypeDoubaoVideoCompatibleTianyiYun}, + {BillingModelName: "cdance2.0-fast-0611", UpstreamModelName: "cdance2.0-fast-0611", ChannelType: constant.ChannelTypeDoubaoVideoCompatibleTianyiYun}, } func SupportsMatrixUsageBilling(modelName string, channelType int, mappedModel string) bool { diff --git a/relay/helper/matrix_usage_capability_test.go b/relay/helper/matrix_usage_capability_test.go index c711479..93ed3fb 100644 --- a/relay/helper/matrix_usage_capability_test.go +++ b/relay/helper/matrix_usage_capability_test.go @@ -23,6 +23,12 @@ func TestSupportsMatrixUsageBilling_DoubaoAipingChannel(t *testing.T) { require.True(t, SupportsMatrixUsageBilling("doubao-seedance-2-0-fast-260128", constant.ChannelTypeDoubaoVideoCompatibleAiping, "doubao-seedance-2-0-fast-260128")) } +func TestSupportsMatrixUsageBilling_TianyiYunSeedanceChannel(t *testing.T) { + require.True(t, SupportsMatrixUsageBilling("cdance2.0-0611", constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "cdance2.0-0611")) + require.True(t, SupportsMatrixUsageBilling("cdance2.0-fast-0611", constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "cdance2.0-fast-0611")) + require.False(t, SupportsMatrixUsageBilling("cdance2.0-0611", constant.ChannelTypeDoubaoVideoCompatibleAiping, "cdance2.0-0611")) +} + func TestResolvePricingModelMapping(t *testing.T) { mapped, err := ResolvePricingModelMapping("seedance-2", `{"seedance-2":"doubao-video-seedance"}`) require.NoError(t, err) diff --git a/relay/relay_adaptor.go b/relay/relay_adaptor.go index adeb06f..d0b5f47 100644 --- a/relay/relay_adaptor.go +++ b/relay/relay_adaptor.go @@ -33,6 +33,7 @@ import ( taskali "github.com/QuantumNous/new-api/relay/channel/task/ali" taskdoubao "github.com/QuantumNous/new-api/relay/channel/task/doubao" taskdoubaoaiping "github.com/QuantumNous/new-api/relay/channel/task/doubao_aiping" + taskdoubaotianyiyun "github.com/QuantumNous/new-api/relay/channel/task/doubao_tianyiyun" taskGemini "github.com/QuantumNous/new-api/relay/channel/task/gemini" "github.com/QuantumNous/new-api/relay/channel/task/hailuo" taskjimeng "github.com/QuantumNous/new-api/relay/channel/task/jimeng" @@ -157,6 +158,8 @@ func GetTaskAdaptor(platform constant.TaskPlatform) channel.TaskAdaptor { return &taskdoubao.TaskAdaptor{} case constant.ChannelTypeDoubaoVideoCompatibleAiping: return &taskdoubaoaiping.TaskAdaptor{} + case constant.ChannelTypeDoubaoVideoCompatibleTianyiYun: + return &taskdoubaotianyiyun.TaskAdaptor{} case constant.ChannelTypeKlingAiping: return &klingaiping.TaskAdaptor{} case constant.ChannelTypeSora, constant.ChannelTypeOpenAI: diff --git a/relay/relay_adaptor_tianyiyun_test.go b/relay/relay_adaptor_tianyiyun_test.go new file mode 100644 index 0000000..c3be900 --- /dev/null +++ b/relay/relay_adaptor_tianyiyun_test.go @@ -0,0 +1,16 @@ +package relay + +import ( + "strconv" + "testing" + + "github.com/QuantumNous/new-api/constant" + tianyiyun "github.com/QuantumNous/new-api/relay/channel/task/doubao_tianyiyun" + "github.com/stretchr/testify/require" +) + +func TestGetTaskAdaptorReturnsTianyiYunSeedanceAdaptor(t *testing.T) { + adaptor := GetTaskAdaptor(constant.TaskPlatform(strconv.Itoa(constant.ChannelTypeDoubaoVideoCompatibleTianyiYun))) + + require.IsType(t, &tianyiyun.TaskAdaptor{}, adaptor) +} diff --git a/router/tianyiyun_seedance_e2e_test.go b/router/tianyiyun_seedance_e2e_test.go new file mode 100644 index 0000000..9bea8fa --- /dev/null +++ b/router/tianyiyun_seedance_e2e_test.go @@ -0,0 +1,188 @@ +package router + +import ( + "net/http" + "net/http/httptest" + "os" + "strings" + "sync/atomic" + "testing" + + "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/QuantumNous/new-api/setting/ratio_setting" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" + "gorm.io/gorm" + "gorm.io/gorm/logger" +) + +func setupTianyiYunSeedanceE2EDB(t *testing.T) *gorm.DB { + t.Helper() + + oldDB := model.DB + oldLOGDB := model.LOG_DB + oldSQLitePath := common.SQLitePath + oldMemoryCacheEnabled := common.MemoryCacheEnabled + oldRedisEnabled := common.RedisEnabled + oldIsMasterNode := common.IsMasterNode + oldUsingSQLite := common.UsingSQLite + oldUsingMySQL := common.UsingMySQL + oldUsingPostgreSQL := common.UsingPostgreSQL + oldSQLDSN, hadSQLDSN := os.LookupEnv("SQL_DSN") + oldModelRatio := ratio_setting.ModelRatio2JSONString() + oldGroupRatio := ratio_setting.GroupRatio2JSONString() + + common.SQLitePath = "file:" + strings.NewReplacer("/", "_", " ", "_").Replace(t.Name()) + "?mode=memory&cache=shared" + common.MemoryCacheEnabled = false + common.RedisEnabled = 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()) + require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(`{"cdance2.0-0611":0}`)) + require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1}`)) + service.InitHttpClient() + + model.DB = model.DB.Session(&gorm.Session{Logger: logger.Default.LogMode(logger.Silent)}) + model.LOG_DB = model.DB + db := model.DB + + sqlDB, err := db.DB() + require.NoError(t, err) + sqlDB.SetMaxOpenConns(1) + require.NoError(t, db.AutoMigrate( + &model.User{}, + &model.Token{}, + &model.Channel{}, + &model.Ability{}, + &model.UserAssetChannel{}, + &model.Task{}, + &model.Log{}, + &model.SubscriptionPlan{}, + &model.SubscriptionOrder{}, + &model.UserSubscription{}, + &model.SubscriptionPreConsumeRecord{}, + )) + + t.Cleanup(func() { + _ = sqlDB.Close() + model.DB = oldDB + model.LOG_DB = oldLOGDB + common.SQLitePath = oldSQLitePath + common.MemoryCacheEnabled = oldMemoryCacheEnabled + common.RedisEnabled = oldRedisEnabled + common.IsMasterNode = oldIsMasterNode + common.UsingSQLite = oldUsingSQLite + common.UsingMySQL = oldUsingMySQL + common.UsingPostgreSQL = oldUsingPostgreSQL + _ = ratio_setting.UpdateModelRatioByJSONString(oldModelRatio) + _ = ratio_setting.UpdateGroupRatioByJSONString(oldGroupRatio) + if hadSQLDSN { + _ = os.Setenv("SQL_DSN", oldSQLDSN) + } else { + _ = os.Unsetenv("SQL_DSN") + } + }) + + return db +} + +func TestTianyiYunSeedanceSubmitE2EUsesBoundUserChannel(t *testing.T) { + gin.SetMode(gin.TestMode) + db := setupTianyiYunSeedanceE2EDB(t) + + var upstreamCalls int32 + var gotPath string + var gotAuth string + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + atomic.AddInt32(&upstreamCalls, 1) + gotPath = r.URL.Path + gotAuth = r.Header.Get("Authorization") + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"id":"task-upstream","status":"queued","model":"cdance2.0-0611"}`)) + })) + t.Cleanup(upstream.Close) + + require.NoError(t, db.Create(&model.User{ + Id: 10, + Username: "e2e-user", + Status: common.UserStatusEnabled, + Group: "default", + Quota: 1000000, + }).Error) + require.NoError(t, db.Create(&model.Token{ + Id: 20, + UserId: 10, + Key: "e2etokenkey", + Status: common.TokenStatusEnabled, + Name: "e2e-token", + ExpiredTime: -1, + UnlimitedQuota: true, + Group: "default", + }).Error) + priority := int64(1) + weight := uint(10) + autoBan := 1 + baseURL := upstream.URL + require.NoError(t, db.Create(&model.Channel{ + Id: 30, + Type: constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, + Key: "upstream-secret", + Status: common.ChannelStatusEnabled, + Name: "tianyiyun-e2e", + Group: "default", + Models: "cdance2.0-0611", + BaseURL: &baseURL, + Priority: &priority, + Weight: &weight, + AutoBan: &autoBan, + CreatedTime: 30, + }).Error) + require.NoError(t, db.Create(&model.Ability{ + Group: "default", + Model: "cdance2.0-0611", + ChannelId: 30, + Enabled: true, + Priority: &priority, + Weight: 10, + }).Error) + require.NoError(t, model.BindUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "default", 30)) + + r := gin.New() + SetVideoRouter(r) + body := `{ + "model":"cdance2.0-0611", + "content":[{"type":"text","text":"e2e prompt"}], + "ratio":"16:9", + "duration":5, + "watermark":false + }` + req := httptest.NewRequest(http.MethodPost, "/api/v3/contents/generations/tasks", strings.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer sk-e2etokenkey") + w := httptest.NewRecorder() + + r.ServeHTTP(w, req) + + require.Equal(t, http.StatusOK, w.Code, w.Body.String()) + require.Equal(t, int32(1), atomic.LoadInt32(&upstreamCalls)) + require.Equal(t, "/v1/contents/generations/tasks", gotPath) + require.Equal(t, "Bearer upstream-secret", gotAuth) + require.Contains(t, w.Body.String(), `"id":"task_`) + + binding, err := model.GetUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "default") + require.NoError(t, err) + require.NotNil(t, binding) + require.Equal(t, 30, binding.ChannelId) + + var task model.Task + require.NoError(t, db.Where("user_id = ? AND channel_id = ?", 10, 30).First(&task).Error) + require.Equal(t, constant.TaskPlatform("60"), task.Platform) + require.Equal(t, "task-upstream", task.PrivateData.UpstreamTaskID) + require.Equal(t, "cdance2.0-0611", task.Properties.OriginModelName) +} diff --git a/service/doubao_asset_channel.go b/service/doubao_asset_channel.go index ee3626d..8c40ac3 100644 --- a/service/doubao_asset_channel.go +++ b/service/doubao_asset_channel.go @@ -13,8 +13,18 @@ import ( const DoubaoAssetChannelPageSize = 50 const DoubaoAssetChannelType = constant.ChannelTypeDoubaoVideoCompatibleAiping +var DoubaoAssetChannelTypes = []int{ + constant.ChannelTypeDoubaoVideoCompatibleAiping, + constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, +} + func ResolveDoubaoAssetChannel(userId int, tokenGroup string) (*model.Channel, error) { - channel, err := GetBoundDoubaoAssetChannel(userId, tokenGroup) + 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 } @@ -22,7 +32,7 @@ func ResolveDoubaoAssetChannel(userId int, tokenGroup string) (*model.Channel, e return channel, nil } - channel, err = AutoMatchDoubaoAssetChannel(tokenGroup) + channel, err = autoMatchDoubaoAssetChannelByType(tokenGroup, channelType) if err != nil { return nil, err } @@ -30,14 +40,22 @@ func ResolveDoubaoAssetChannel(userId int, tokenGroup string) (*model.Channel, e return nil, nil } - if err = model.BindUserAssetChannel(userId, DoubaoAssetChannelType, tokenGroup, channel.Id); err != nil { + if err = model.BindUserAssetChannel(userId, channel.Type, tokenGroup, channel.Id); err != nil { return nil, err } return channel, nil } func GetBoundDoubaoAssetChannel(userId int, tokenGroup string) (*model.Channel, error) { - binding, err := model.GetUserAssetChannel(userId, DoubaoAssetChannelType, tokenGroup) + return getBoundDoubaoAssetChannelByType(userId, tokenGroup, DoubaoAssetChannelType) +} + +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 } @@ -63,28 +81,35 @@ func HasDoubaoVideoChannelForGroupModel(tokenGroup string, modelName string) boo return false } - for startIdx := 0; ; startIdx += DoubaoAssetChannelPageSize { - candidates, err := model.GetChannelsByType(startIdx, DoubaoAssetChannelPageSize, true, DoubaoAssetChannelType) - if err != nil { - return false - } - for _, candidate := range candidates { - if candidate == nil || !MatchDoubaoAssetGroup(candidate.GetGroups(), tokenGroup) { - continue + for _, channelType := range DoubaoAssetChannelTypes { + for startIdx := 0; ; startIdx += DoubaoAssetChannelPageSize { + candidates, err := model.GetChannelsByType(startIdx, DoubaoAssetChannelPageSize, true, channelType) + if err != nil { + return false } - if model.IsChannelEnabledForGroupModel(tokenGroup, modelName, candidate.Id) { - return true + 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 } - } - if len(candidates) < DoubaoAssetChannelPageSize { - return false } } + return false } 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, DoubaoAssetChannelType) + candidates, err := model.GetChannelsByType(startIdx, DoubaoAssetChannelPageSize, true, channelType) if err != nil { return nil, err } @@ -114,7 +139,7 @@ func IsUsableDoubaoAssetChannel(channel *model.Channel, tokenGroup string) bool if channel == nil { return false } - if channel.Type != DoubaoAssetChannelType { + if !IsDoubaoAssetChannelType(channel.Type) { return false } if channel.Status != common.ChannelStatusEnabled { @@ -126,6 +151,22 @@ func IsUsableDoubaoAssetChannel(channel *model.Channel, tokenGroup string) bool return MatchDoubaoAssetGroup(channel.GetGroups(), tokenGroup) } +func IsDoubaoAssetChannelType(channelType int) bool { + for _, t := range DoubaoAssetChannelTypes { + if channelType == t { + return true + } + } + return false +} + +func DoubaoAssetChannelTypeForModel(modelName string) int { + if strings.HasPrefix(strings.ToLower(strings.TrimSpace(modelName)), "cdance") { + return constant.ChannelTypeDoubaoVideoCompatibleTianyiYun + } + return DoubaoAssetChannelType +} + func MatchDoubaoAssetGroup(channelGroups []string, tokenGroup string) bool { tokenGroup = strings.TrimSpace(tokenGroup) if tokenGroup == "" { diff --git a/service/doubao_asset_channel_test.go b/service/doubao_asset_channel_test.go index f7964eb..11ee29d 100644 --- a/service/doubao_asset_channel_test.go +++ b/service/doubao_asset_channel_test.go @@ -129,6 +129,24 @@ func TestResolveDoubaoAssetChannel_AutoMatchesAndPersistsWhenUnbound(t *testing. assert.Equal(t, 1, binding.ChannelId) } +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) + + channel, err := ResolveDoubaoAssetChannelForModel(10, "default", "cdance2.0-0611") + + require.NoError(t, err) + require.NotNil(t, channel) + assert.Equal(t, 1, channel.Id) + assert.Equal(t, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, channel.Type) + + binding, err := model.GetUserAssetChannel(10, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "default") + require.NoError(t, err) + require.NotNil(t, binding) + assert.Equal(t, 1, binding.ChannelId) +} + func TestResolveDoubaoAssetChannel_RebindsOpenAIMisbindToAiping(t *testing.T) { db := setupDoubaoAssetChannelDB(t) createDoubaoAssetChannelForTest(t, db, 1, constant.ChannelTypeOpenAI, "default", "openai-key", common.ChannelStatusEnabled) @@ -207,12 +225,15 @@ func TestHasDoubaoVideoChannelForGroupModel(t *testing.T) { createDoubaoAssetChannelForTest(t, db, 1, constant.ChannelTypeDoubaoVideoCompatibleAiping, "default", "default-key", common.ChannelStatusEnabled) createDoubaoAssetChannelForTest(t, db, 2, constant.ChannelTypeDoubaoVideoCompatibleAiping, "vip", "vip-key", common.ChannelStatusEnabled) createDoubaoAssetChannelForTest(t, db, 3, constant.ChannelTypeOpenAI, "default", "openai-key", common.ChannelStatusEnabled) + createDoubaoAssetChannelForTest(t, db, 4, constant.ChannelTypeDoubaoVideoCompatibleTianyiYun, "default", "tianyiyun-key", common.ChannelStatusEnabled) createDoubaoAssetAbilityForTest(t, db, "default", "seedance-2", 1, true) 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) assert.True(t, HasDoubaoVideoChannelForGroupModel("default", "seedance-2")) + assert.True(t, HasDoubaoVideoChannelForGroupModel("default", "cdance2.0-0611")) assert.False(t, HasDoubaoVideoChannelForGroupModel("default", "disabled-model")) assert.False(t, HasDoubaoVideoChannelForGroupModel("missing", "seedance-2")) } diff --git a/web/src/constants/channel.constants.js b/web/src/constants/channel.constants.js index 065392c..bde772f 100644 --- a/web/src/constants/channel.constants.js +++ b/web/src/constants/channel.constants.js @@ -199,6 +199,11 @@ export const CHANNEL_OPTIONS = [ color: 'green', label: 'KlingAiping', }, + { + value: 60, + color: 'green', + label: 'DoubaoVideoCompatibleTianyiYun', + }, ]; export const MODEL_TABLE_PAGE_SIZE = 10; diff --git a/web/src/helpers/render.jsx b/web/src/helpers/render.jsx index bdfee36..66d78d2 100644 --- a/web/src/helpers/render.jsx +++ b/web/src/helpers/render.jsx @@ -95,12 +95,12 @@ import { SiOkta, SiOpenid, SiReddit, - SiSlack, SiTelegram, SiTwitch, SiWechat, SiX, } from 'react-icons/si'; +import { FaSlack } from 'react-icons/fa'; // 获取侧边栏Lucide图标组件 export function getLucideIcon(key, selected = false) { @@ -509,7 +509,7 @@ const oauthProviderIconMap = { facebook: SiFacebook, x: SiX, twitter: SiX, - slack: SiSlack, + slack: FaSlack, telegram: SiTelegram, wechat: SiWechat, keycloak: SiKeycloak,