new-api/relay/channel/task/chinamobile_seedance/adaptor.go

package chinamobile_seedance

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"
	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"
)

type responsePayload struct {
	ID string `json:"id"`
}

type upstreamError struct {
	Code    string `json:"code"`
	Message string `json:"message"`
}

func (e *upstreamError) UnmarshalJSON(data []byte) error {
	if len(bytes.TrimSpace(data)) == 0 || string(bytes.TrimSpace(data)) == "null" {
		return nil
	}
	var message string
	if err := common.Unmarshal(data, &message); err == nil {
		e.Message = message
		return nil
	}
	type alias upstreamError
	var parsed alias
	if err := common.Unmarshal(data, &parsed); err != nil {
		return err
	}
	*e = upstreamError(parsed)
	return nil
}

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     upstreamError `json:"error"`
	Message   string        `json:"message"`
	CreatedAt int64         `json:"created_at"`
	UpdatedAt int64         `json:"updated_at"`
}

type TaskAdaptor struct {
	taskcommon.BaseBilling
	ChannelType   int
	apiKey        string
	baseURL       string
	newSDK        seedanceSDKFactory
	submitTimeout time.Duration
	queryTimeout  time.Duration
}

func (a *TaskAdaptor) Init(info *relaycommon.RelayInfo) {
	if info != nil && info.ChannelMeta != nil {
		a.ChannelType = info.ChannelType
		a.baseURL = strings.TrimRight(info.ChannelBaseUrl, "/")
		a.apiKey = info.ApiKey
	}
	if a.newSDK == nil {
		a.newSDK = newRealSDK
	}
	if a.submitTimeout == 0 {
		a.submitTimeout = defaultSubmitTimeout
	}
	if a.queryTimeout == 0 {
		a.queryTimeout = defaultQueryTimeout
	}
}

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, "/")
	}
	if strings.TrimSpace(baseURL) == "" {
		baseURL = defaultBaseURL
	}
	return fmt.Sprintf("%s/contents/generations/tasks", strings.TrimRight(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")
	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 != nil {
		if info.IsModelMapped {
			body["model"] = info.UpstreamModelName
		} else if modelName, _ := body["model"].(string); modelName != "" {
			info.UpstreamModelName = modelName
		}
	}
	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) {
	data, err := io.ReadAll(requestBody)
	if err != nil {
		return nil, err
	}
	payload := map[string]interface{}{}
	if err := common.Unmarshal(data, &payload); err != nil {
		return nil, err
	}
	modelName := chinaMobileSeedanceModel(modelFromPayloadOrInfo(payload, info))
	payload["model"] = modelName
	if info != nil {
		info.UpstreamModelName = modelName
	}
	client, err := a.createSDK(info, modelName)
	if err != nil {
		return nil, err
	}
	submitTimeout := a.submitTimeout
	if submitTimeout == 0 {
		submitTimeout = defaultSubmitTimeout
	}
	taskID, err := runSDKCall[string](submitTimeout, func() (string, error) {
		return client.CreateVideoGenerationTask(payload)
	})
	if err != nil {
		return chinaMobileSeedanceErrorResponse(err), nil
	}
	respBody, err := common.Marshal(map[string]any{"id": taskID})
	if err != nil {
		return nil, err
	}
	return &http.Response{
		StatusCode: http.StatusOK,
		Header:     make(http.Header),
		Body:       io.NopCloser(bytes.NewReader(respBody)),
	}, nil
}

func (a *TaskAdaptor) DoResponse(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (taskID string, taskData []byte, taskErr *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)
	if info.PublicTaskID != "" {
		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) {
	taskID, _ := body["task_id"].(string)
	if strings.TrimSpace(taskID) == "" {
		return nil, fmt.Errorf("invalid task_id")
	}
	modelName, _ := body["model"].(string)
	client, err := a.createSDKFromValues(baseUrl, key, modelName)
	if err != nil {
		return nil, err
	}
	queryTimeout := a.queryTimeout
	if queryTimeout == 0 {
		queryTimeout = defaultQueryTimeout
	}
	result, err := runSDKCall[map[string]interface{}](queryTimeout, func() (map[string]interface{}, error) {
		return client.QueryVideoGenerationTask(taskID)
	})
	if err != nil {
		return nil, err
	}
	respBody, err := common.Marshal(result)
	if err != nil {
		return nil, err
	}
	return &http.Response{
		StatusCode: http.StatusOK,
		Header:     make(http.Header),
		Body:       io.NopCloser(bytes.NewReader(respBody)),
	}, 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 strings.ToLower(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 = upstreamReason(resTask)
	default:
		if resTask.Error.Message != "" || resTask.Message != "" {
			taskResult.Status = model.TaskStatusFailure
			taskResult.Progress = "100%"
			taskResult.Reason = upstreamReason(resTask)
		} else {
			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 chinamobile seedance 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 originTask.Status == model.TaskStatusFailure || dResp.Status == "failed" {
		message := upstreamReason(dResp)
		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)
}

func (a *TaskAdaptor) GetModelList() []string {
	return ModelList
}

func (a *TaskAdaptor) GetChannelName() string {
	return ChannelName
}

func (a *TaskAdaptor) createSDK(info *relaycommon.RelayInfo, modelName string) (seedanceSDK, error) {
	baseURL := a.baseURL
	apiKey := a.apiKey
	if info != nil && info.ChannelMeta != nil {
		baseURL = info.ChannelBaseUrl
		apiKey = info.ApiKey
	}
	return a.createSDKFromValues(baseURL, apiKey, modelName)
}

func (a *TaskAdaptor) createSDKFromValues(baseURL, apiKey, modelName string) (seedanceSDK, error) {
	factory := a.newSDK
	if factory == nil {
		factory = newRealSDK
	}
	return factory(strings.TrimRight(baseURL, "/"), apiKey, modelName)
}

func (a *TaskAdaptor) convertToRequestPayload(req *relaycommon.TaskSubmitReq) (map[string]interface{}, error) {
	payload := map[string]interface{}{
		"model": req.Model,
	}
	for k, v := range req.Metadata {
		payload[k] = v
	}
	if sec, _ := strconv.Atoi(req.Seconds); sec > 0 {
		if _, ok := payload["duration"]; !ok {
			payload["duration"] = sec
		}
	}
	content, err := normalizeContent(payload["content"])
	if err != nil {
		return nil, err
	}
	if len(content) == 0 && req.HasImage() {
		for _, imgURL := range req.Images {
			content = append(content, map[string]interface{}{
				"type":      "image_url",
				"image_url": map[string]interface{}{"url": imgURL},
			})
		}
	}
	if strings.TrimSpace(req.Prompt) != "" && !hasTextContent(content) {
		content = append(content, map[string]interface{}{
			"type": "text",
			"text": req.Prompt,
		})
	}
	if len(content) > 0 {
		payload["content"] = content
	}
	return payload, nil
}

func normalizeContent(input any) ([]interface{}, error) {
	if input == nil {
		return nil, nil
	}
	data, err := common.Marshal(input)
	if err != nil {
		return nil, err
	}
	var content []interface{}
	if err := common.Unmarshal(data, &content); err != nil {
		return nil, err
	}
	return content, nil
}

func hasTextContent(content []interface{}) bool {
	for _, item := range content {
		m, ok := item.(map[string]interface{})
		if !ok {
			continue
		}
		if m["type"] == "text" && strings.TrimSpace(fmt.Sprint(m["text"])) != "" {
			return true
		}
	}
	return false
}

func modelFromPayloadOrInfo(payload map[string]interface{}, info *relaycommon.RelayInfo) string {
	if modelName, _ := payload["model"].(string); strings.TrimSpace(modelName) != "" {
		return modelName
	}
	if info != nil {
		if strings.TrimSpace(info.UpstreamModelName) != "" {
			return info.UpstreamModelName
		}
		return info.OriginModelName
	}
	return defaultModel
}

func chinaMobileSeedanceModel(modelName string) string {
	switch strings.TrimSpace(modelName) {
	case "", "doubao-seedance-2-0-260128", "doubao-seedance-2-0-fast-260128":
		return defaultModel
	default:
		return modelName
	}
}

type chinaMobileSeedanceUpstreamError struct {
	ErrorCode    string `json:"ErrorCode"`
	ErrorMessage string `json:"ErrorMessage"`
}

func chinaMobileSeedanceErrorResponse(err error) *http.Response {
	statusCode := http.StatusBadGateway
	body := map[string]any{
		"code":    "upstream_error",
		"message": err.Error(),
	}
	if upstreamErr, ok := parseChinaMobileSeedanceError(err); ok {
		statusCode = chinaMobileSeedanceStatusCode(upstreamErr.ErrorCode)
		body["code"] = upstreamErr.ErrorCode
		body["message"] = upstreamErr.ErrorMessage
	}
	respBody, marshalErr := common.Marshal(body)
	if marshalErr != nil {
		respBody = []byte(`{"code":"upstream_error","message":"upstream error"}`)
	}
	return &http.Response{
		StatusCode: statusCode,
		Header:     make(http.Header),
		Body:       io.NopCloser(bytes.NewReader(respBody)),
	}
}

func parseChinaMobileSeedanceError(err error) (chinaMobileSeedanceUpstreamError, bool) {
	if err == nil {
		return chinaMobileSeedanceUpstreamError{}, false
	}
	text := err.Error()
	start := strings.Index(text, "{")
	end := strings.LastIndex(text, "}")
	if start < 0 || end < start {
		return chinaMobileSeedanceUpstreamError{}, false
	}
	var upstreamErr chinaMobileSeedanceUpstreamError
	if err := common.Unmarshal([]byte(text[start:end+1]), &upstreamErr); err != nil {
		return chinaMobileSeedanceUpstreamError{}, false
	}
	return upstreamErr, upstreamErr.ErrorCode != "" || upstreamErr.ErrorMessage != ""
}

func chinaMobileSeedanceStatusCode(code string) int {
	upperCode := strings.ToUpper(strings.TrimSpace(code))
	switch {
	case strings.Contains(upperCode, "PERMISSION") || strings.Contains(upperCode, "UNAUTHORIZED"):
		return http.StatusForbidden
	case strings.Contains(upperCode, "SENSITIVE") || strings.Contains(upperCode, "INVALID") || strings.Contains(upperCode, "BAD"):
		return http.StatusBadRequest
	case strings.Contains(upperCode, "RATE") || strings.Contains(upperCode, "LIMIT"):
		return http.StatusTooManyRequests
	default:
		return http.StatusBadGateway
	}
}

func upstreamReason(resTask responseTask) string {
	if resTask.Error.Message != "" {
		return resTask.Error.Message
	}
	if resTask.Message != "" {
		return resTask.Message
	}
	if resTask.Status != "" {
		return "task " + resTask.Status
	}
	return "upstream error"
}