feat: suno api 支持

controller/relay.go

@@ -190,3 +190,94 @@ func RelayNotFound(c *gin.Context) {
 		"error": err,
 	})
 }
+
+func RelayTask(c *gin.Context) {
+	retryTimes := common.RetryTimes
+	channelId := c.GetInt("channel_id")
+	relayMode := c.GetInt("relay_mode")
+	group := c.GetString("group")
+	originalModel := c.GetString("original_model")
+	c.Set("use_channel", []string{fmt.Sprintf("%d", channelId)})
+	taskErr := taskRelayHandler(c, relayMode)
+	if taskErr == nil {
+		retryTimes = 0
+	}
+	for i := 0; shouldRetryTaskRelay(c, channelId, taskErr, retryTimes) && i < retryTimes; i++ {
+		channel, err := model.CacheGetRandomSatisfiedChannel(group, originalModel, i)
+		if err != nil {
+			common.LogError(c.Request.Context(), fmt.Sprintf("CacheGetRandomSatisfiedChannel failed: %s", err.Error()))
+			break
+		}
+		channelId = channel.Id
+		useChannel := c.GetStringSlice("use_channel")
+		useChannel = append(useChannel, fmt.Sprintf("%d", channelId))
+		c.Set("use_channel", useChannel)
+		common.LogInfo(c.Request.Context(), fmt.Sprintf("using channel #%d to retry (remain times %d)", channel.Id, i))
+		middleware.SetupContextForSelectedChannel(c, channel, originalModel)
+
+		requestBody, err := common.GetRequestBody(c)
+		c.Request.Body = io.NopCloser(bytes.NewBuffer(requestBody))
+		taskErr = taskRelayHandler(c, relayMode)
+	}
+	useChannel := c.GetStringSlice("use_channel")
+	if len(useChannel) > 1 {
+		retryLogStr := fmt.Sprintf("重试：%s", strings.Trim(strings.Join(strings.Fields(fmt.Sprint(useChannel)), "->"), "[]"))
+		common.LogInfo(c.Request.Context(), retryLogStr)
+	}
+	if taskErr != nil {
+		if taskErr.StatusCode == http.StatusTooManyRequests {
+			taskErr.Message = "当前分组上游负载已饱和，请稍后再试"
+		}
+		c.JSON(taskErr.StatusCode, taskErr)
+	}
+}
+
+func taskRelayHandler(c *gin.Context, relayMode int) *dto.TaskError {
+	var err *dto.TaskError
+	switch relayMode {
+	case relayconstant.RelayModeSunoFetch, relayconstant.RelayModeSunoFetchByID:
+		err = relay.RelayTaskFetch(c, relayMode)
+	default:
+		err = relay.RelayTaskSubmit(c, relayMode)
+	}
+	return err
+}
+
+func shouldRetryTaskRelay(c *gin.Context, channelId int, taskErr *dto.TaskError, retryTimes int) bool {
+	if taskErr == nil {
+		return false
+	}
+	if retryTimes <= 0 {
+		return false
+	}
+	if _, ok := c.Get("specific_channel_id"); ok {
+		return false
+	}
+	if taskErr.StatusCode == http.StatusTooManyRequests {
+		return true
+	}
+	if taskErr.StatusCode == 307 {
+		return true
+	}
+	if taskErr.StatusCode/100 == 5 {
+		// 超时不重试
+		if taskErr.StatusCode == 504 || taskErr.StatusCode == 524 {
+			return false
+		}
+		return true
+	}
+	if taskErr.StatusCode == http.StatusBadRequest {
+		return false
+	}
+	if taskErr.StatusCode == 408 {
+		// azure处理超时不重试
+		return false
+	}
+	if taskErr.LocalError {
+		return false
+	}
+	if taskErr.StatusCode/100 == 2 {
+		return false
+	}
+	return true
+}

controller/task.go

@@ -0,0 +1,92 @@
+package controller
+
+import (
+	"github.com/gin-gonic/gin"
+	"log"
+	"one-api/common"
+	"one-api/constant"
+	"one-api/model"
+	"strconv"
+	"time"
+)
+
+func UpdateTaskBulk() {
+	//revocer
+	//imageModel := "midjourney"
+	for {
+		time.Sleep(time.Duration(15) * time.Second)
+		common.SysLog("任务进度轮询开始")
+		allTasks := model.GetAllUnFinishSyncTasks(500)
+		platformTask := make(map[constant.TaskPlatform][]*model.Task)
+		for _, t := range allTasks {
+			platformTask[t.Platform] = append(platformTask[t.Platform], t)
+		}
+		for platform, tasks := range platformTask {
+			UpdateTaskByPlatform(platform, tasks)
+		}
+		common.SysLog("任务进度轮询完成")
+	}
+}
+
+func GetAllMidjourney(c *gin.Context) {
+	p, _ := strconv.Atoi(c.Query("p"))
+	if p < 0 {
+		p = 0
+	}
+
+	// 解析其他查询参数
+	queryParams := model.TaskQueryParams{
+		ChannelID:      c.Query("channel_id"),
+		MjID:           c.Query("mj_id"),
+		StartTimestamp: c.Query("start_timestamp"),
+		EndTimestamp:   c.Query("end_timestamp"),
+	}
+
+	logs := model.GetAllTasks(p*common.ItemsPerPage, common.ItemsPerPage, queryParams)
+	if logs == nil {
+		logs = make([]*model.Midjourney, 0)
+	}
+	if constant.MjForwardUrlEnabled {
+		for i, midjourney := range logs {
+			midjourney.ImageUrl = constant.ServerAddress + "/mj/image/" + midjourney.MjId
+			logs[i] = midjourney
+		}
+	}
+	c.JSON(200, gin.H{
+		"success": true,
+		"message": "",
+		"data":    logs,
+	})
+}
+
+func GetUserMidjourney(c *gin.Context) {
+	p, _ := strconv.Atoi(c.Query("p"))
+	if p < 0 {
+		p = 0
+	}
+
+	userId := c.GetInt("id")
+	log.Printf("userId = %d \n", userId)
+
+	queryParams := model.TaskQueryParams{
+		MjID:           c.Query("mj_id"),
+		StartTimestamp: c.Query("start_timestamp"),
+		EndTimestamp:   c.Query("end_timestamp"),
+	}
+
+	logs := model.GetAllUserTask(userId, p*common.ItemsPerPage, common.ItemsPerPage, queryParams)
+	if logs == nil {
+		logs = make([]*model.Midjourney, 0)
+	}
+	if constant.MjForwardUrlEnabled {
+		for i, midjourney := range logs {
+			midjourney.ImageUrl = constant.ServerAddress + "/mj/image/" + midjourney.MjId
+			logs[i] = midjourney
+		}
+	}
+	c.JSON(200, gin.H{
+		"success": true,
+		"message": "",
+		"data":    logs,
+	})
+}