new-api/controller/doubao_aiping_video.go

package controller

import (
	"fmt"
	"net/http"
	"strings"

	"github.com/QuantumNous/new-api/common"
	"github.com/QuantumNous/new-api/constant"
	"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/types"
	"github.com/gin-gonic/gin"
)

func convertAipingNativeVideoRequest(body []byte) (relaycommon.TaskSubmitReq, error) {
	var raw map[string]any
	if err := common.Unmarshal(body, &raw); err != nil {
		return relaycommon.TaskSubmitReq{}, err
	}
	modelName, _ := raw["model"].(string)
	if strings.TrimSpace(modelName) == "" {
		return relaycommon.TaskSubmitReq{}, fmt.Errorf("model field is required")
	}

	metadata := map[string]any{}
	for key, value := range raw {
		if key == "model" || key == "content" {
			continue
		}
		metadata[key] = value
	}

	var passthroughContent []any
	if rawContent, ok := raw["content"].([]any); ok {
		for _, item := range rawContent {
			itemMap, ok := item.(map[string]any)
			if !ok {
				passthroughContent = append(passthroughContent, item)
				continue
			}
			passthroughContent = append(passthroughContent, itemMap)
		}
	}
	if len(passthroughContent) > 0 {
		metadata["content"] = passthroughContent
	}

	return relaycommon.TaskSubmitReq{
		Model:    strings.TrimSpace(modelName),
		Metadata: metadata,
	}, nil
}

func prepareAipingNativeVideoSubmit(c *gin.Context, info *relaycommon.RelayInfo) error {
	bodyStorage, err := common.GetBodyStorage(c)
	if err != nil {
		return err
	}
	body, err := bodyStorage.Bytes()
	if err != nil {
		return err
	}
	req, err := convertAipingNativeVideoRequest(body)
	if err != nil {
		return err
	}
	info.OriginModelName = req.Model
	info.Action = constant.TaskActionGenerate
	relaycommon.StoreTaskRequest(c, info, constant.TaskActionGenerate, req)
	return nil
}

func AipingNativeVideoSubmit(c *gin.Context) {
	relayInfo, err := relaycommon.GenRelayInfo(c, types.RelayFormatTask, nil, nil)
	if err != nil {
		c.JSON(http.StatusInternalServerError, &dto.TaskError{
			Code:       "gen_relay_info_failed",
			Message:    err.Error(),
			StatusCode: http.StatusInternalServerError,
		})
		return
	}
	if relayInfo.ChannelMeta == nil {
		relayInfo.ChannelMeta = &relaycommon.ChannelMeta{}
	}
	if err := prepareAipingNativeVideoSubmit(c, relayInfo); err != nil {
		c.JSON(http.StatusBadRequest, &dto.TaskError{
			Code:       "invalid_request",
			Message:    err.Error(),
			StatusCode: http.StatusBadRequest,
		})
		return
	}
	relayTaskWithInfo(c, relayInfo)
}

func buildAipingNativeFetchResponse(task *model.Task) ([]byte, error) {
	payload := map[string]any{}
	if len(task.Data) > 0 {
		if err := common.Unmarshal(task.Data, &payload); err != nil {
			return nil, err
		}
	}
	delete(payload, "aiping_id")
	if _, ok := payload["error"]; !ok {
		payload["error"] = nil
	}
	payload["id"] = task.TaskID
	if _, ok := payload["model"]; !ok {
		payload["model"] = task.Properties.OriginModelName
	}
	if _, ok := payload["status"]; !ok {
		payload["status"] = mapAipingNativeStatus(task.Status)
	}
	if _, ok := payload["created_at"]; !ok {
		payload["created_at"] = task.CreatedAt
	}
	if _, ok := payload["updated_at"]; !ok {
		payload["updated_at"] = task.UpdatedAt
	}
	return common.Marshal(payload)
}

func mapAipingNativeStatus(status model.TaskStatus) string {
	switch status {
	case model.TaskStatusQueued, model.TaskStatusSubmitted:
		return "queued"
	case model.TaskStatusInProgress:
		return "running"
	case model.TaskStatusSuccess:
		return "succeeded"
	case model.TaskStatusFailure:
		return "failed"
	default:
		return "running"
	}
}

func AipingNativeVideoFetch(c *gin.Context) {
	taskID := c.Param("task_id")
	task, exist, err := model.GetByTaskId(c.GetInt("id"), taskID)
	if err != nil {
		c.JSON(http.StatusInternalServerError, &dto.TaskError{Code: "get_task_failed", Message: err.Error(), StatusCode: http.StatusInternalServerError})
		return
	}
	if !exist {
		c.JSON(http.StatusNotFound, gin.H{"error": gin.H{"code": "not_found", "message": "task not found"}})
		return
	}
	data, err := buildAipingNativeFetchResponse(task)
	if err != nil {
		c.JSON(http.StatusInternalServerError, &dto.TaskError{Code: "build_response_failed", Message: err.Error(), StatusCode: http.StatusInternalServerError})
		return
	}
	c.Data(http.StatusOK, "application/json", data)
}