new-api/relay/channel/task/kling/aiping/adaptor.go

package aiping

import (
	"bytes"
	"fmt"
	"io"
	"net/http"
	"strings"
	"time"

	"github.com/QuantumNous/new-api/common"
	"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"
	"github.com/gin-gonic/gin"
	"github.com/pkg/errors"
	"github.com/shopspring/decimal"
)

type responsePayload struct {
	Code      int    `json:"code"`
	Message   string `json:"message"`
	RequestID string `json:"request_id"`
	AipingID  string `json:"aiping_id"`
	Data      struct {
		TaskID        string `json:"task_id"`
		TaskStatus    string `json:"task_status"`
		TaskStatusMsg string `json:"task_status_msg"`
		TaskInfo      any    `json:"task_info,omitempty"`
		TaskResult    struct {
			Videos []struct {
				ID           string `json:"id"`
				URL          string `json:"url"`
				WatermarkURL string `json:"watermark_url,omitempty"`
				Duration     string `json:"duration"`
			} `json:"videos"`
		} `json:"task_result"`
		FinalUnitDeduction any   `json:"final_unit_deduction"`
		CreatedAt          int64 `json:"created_at"`
		UpdatedAt          int64 `json:"updated_at"`
	} `json:"data"`
}

type TaskAdaptor struct {
	taskcommon.BaseBilling
	apiKey  string
	baseURL string
}

func (a *TaskAdaptor) Init(info *relaycommon.RelayInfo) {
	a.apiKey = info.ApiKey
	a.baseURL = strings.TrimRight(info.ChannelBaseUrl, "/")
}

func (a *TaskAdaptor) ValidateRequestAndSetAction(c *gin.Context, info *relaycommon.RelayInfo) *dto.TaskError {
	req, err := relaycommon.GetTaskRequest(c)
	if err != nil {
		return service.TaskErrorWrapperLocal(fmt.Errorf("task request not prepared"), "invalid_request", http.StatusBadRequest)
	}
	modelName := resolveRequestModelName(req, info)
	if isOmniModel(modelName) && (info.Action == ActionText2Video || info.Action == ActionImage2Video) {
		return service.TaskErrorWrapperLocal(
			fmt.Errorf("model %s must use /v1/videos/omni-video", modelName),
			"invalid_model_route",
			http.StatusUnprocessableEntity,
		)
	}
	return nil
}

func (a *TaskAdaptor) BuildRequestURL(info *relaycommon.RelayInfo) (string, error) {
	route, ok := FindRouteByAction(info.Action)
	if !ok {
		return "", fmt.Errorf("unsupported kling aiping action: %s", info.Action)
	}
	return a.baseURL + route.UpstreamPath, 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 := cloneMap(req.Metadata)
	normalizeRequestBody(body, info)
	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 parsed responsePayload
	if err := common.Unmarshal(responseBody, &parsed); err != nil {
		return "", nil, service.TaskErrorWrapper(errors.Wrapf(err, "body: %s", responseBody), "unmarshal_response_body_failed", http.StatusInternalServerError)
	}
	if parsed.Code != 0 {
		msg := parsed.Message
		if msg == "" {
			msg = "upstream task failed"
		}
		return "", nil, service.TaskErrorWrapper(fmt.Errorf("%s", msg), "task_failed", http.StatusBadRequest)
	}
	upstreamTaskID := parsed.Data.TaskID
	if strings.TrimSpace(upstreamTaskID) == "" {
		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)
	}
	sanitizeNativeTaskPayload(clientPayload, info.PublicTaskID)
	c.JSON(http.StatusOK, clientPayload)
	return upstreamTaskID, 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) == "" {
		taskID, _ = body["upstream_task_id"].(string)
	}
	action, _ := body["action"].(string)
	if strings.TrimSpace(taskID) == "" {
		return nil, fmt.Errorf("invalid task_id")
	}
	route, ok := FindFetchRouteByAction(action)
	if !ok {
		return nil, fmt.Errorf("unsupported kling aiping action: %s", action)
	}
	url := strings.TrimRight(baseURL, "/") + strings.Replace(route.UpstreamPath, ":task_id", taskID, 1)
	req, err := http.NewRequest(http.MethodGet, url, nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("Accept", "application/json")
	req.Header.Set("Authorization", "Bearer "+key)
	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) ParseTaskResult(respBody []byte) (*relaycommon.TaskInfo, error) {
	var parsed responsePayload
	if err := common.Unmarshal(respBody, &parsed); err != nil {
		return nil, errors.Wrap(err, "unmarshal kling aiping task data failed")
	}
	info := &relaycommon.TaskInfo{
		Code:   parsed.Code,
		TaskID: parsed.Data.TaskID,
		Reason: parsed.Data.TaskStatusMsg,
	}
	switch parsed.Data.TaskStatus {
	case "submitted":
		info.Status = model.TaskStatusSubmitted
		info.Progress = taskcommon.ProgressSubmitted
	case "processing":
		info.Status = model.TaskStatusInProgress
		info.Progress = taskcommon.ProgressInProgress
	case "succeed":
		info.Status = model.TaskStatusSuccess
		info.Progress = taskcommon.ProgressComplete
		if tokens := finalUnitDeductionTokens(parsed.Data.FinalUnitDeduction); tokens > 0 {
			info.CompletionTokens = tokens
			info.TotalTokens = tokens
		}
	case "failed":
		info.Status = model.TaskStatusFailure
		info.Progress = taskcommon.ProgressComplete
	default:
		info.Status = model.TaskStatusInProgress
		info.Progress = taskcommon.ProgressInProgress
	}
	if len(parsed.Data.TaskResult.Videos) > 0 {
		info.Url = parsed.Data.TaskResult.Videos[0].URL
	}
	return info, nil
}

func finalUnitDeductionTokens(value any) int {
	var d decimal.Decimal
	var err error
	switch v := value.(type) {
	case string:
		if strings.TrimSpace(v) == "" {
			return 0
		}
		d, err = decimal.NewFromString(strings.TrimSpace(v))
	case float64:
		d = decimal.NewFromFloat(v)
	case int:
		d = decimal.NewFromInt(int64(v))
	case int64:
		d = decimal.NewFromInt(v)
	default:
		return 0
	}
	if err != nil || !d.IsPositive() {
		return 0
	}
	return int(d.Mul(decimal.NewFromInt(1000000)).IntPart())
}

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

func (a *TaskAdaptor) GetChannelName() string {
	return "kling-aiping"
}

func normalizeRequestBody(body map[string]any, info *relaycommon.RelayInfo) {
	delete(body, "action_control")
	delete(body, "uid")
	delete(body, "create_at")
	delete(body, "_standard_model")

	modelName := stringValue(body["model_name"])
	if modelName == "" {
		modelName = stringValue(body["model"])
	}
	if info != nil && info.IsModelMapped && info.UpstreamModelName != "" {
		modelName = info.UpstreamModelName
	}
	if modelName != "" {
		body["model_name"] = canonicalAipingModelName(modelName)
	}
	delete(body, "model")

	duration := body["duration"]
	if duration == nil {
		duration = body["seconds"]
	}
	if duration != nil {
		body["duration"] = fmt.Sprint(duration)
	}
	delete(body, "seconds")

	if imageList, ok := body["reference_images"]; ok {
		if _, exists := body["image_list"]; !exists {
			body["image_list"] = imageList
		}
		delete(body, "reference_images")
	}
	if imageList, ok := normalizeImageList(body["image_list"]); ok {
		body["image_list"] = imageList
	}

	if body["watermark_info"] == nil {
		return
	}
}

func normalizeImageList(v any) ([]map[string]any, bool) {
	items, ok := v.([]any)
	if !ok {
		return nil, false
	}
	out := make([]map[string]any, 0, len(items))
	for _, item := range items {
		switch val := item.(type) {
		case string:
			out = append(out, map[string]any{"image": val})
		case map[string]any:
			image := firstString(val, "image", "image_url", "url", "base64")
			if image == "" {
				out = append(out, val)
				continue
			}
			out = append(out, map[string]any{"image": image})
		default:
			out = append(out, map[string]any{"image": fmt.Sprint(val)})
		}
	}
	return out, true
}

func sanitizeNativeTaskPayload(payload map[string]any, publicTaskID string) {
	delete(payload, "aiping_id")
	data, _ := payload["data"].(map[string]any)
	if data == nil {
		return
	}
	data["task_id"] = publicTaskID
	ensureWatermarkURL(data)
	if _, ok := data["created_at"]; !ok {
		data["created_at"] = time.Now().UnixMilli()
	}
	if _, ok := data["updated_at"]; !ok {
		data["updated_at"] = data["created_at"]
	}
}

func ensureWatermarkURL(data map[string]any) {
	taskResult, _ := data["task_result"].(map[string]any)
	if taskResult == nil {
		return
	}
	videos, _ := taskResult["videos"].([]any)
	for _, videoAny := range videos {
		video, _ := videoAny.(map[string]any)
		if video == nil {
			continue
		}
		if _, ok := video["watermark_url"]; !ok {
			video["watermark_url"] = ""
		}
	}
}

func cloneMap(in map[string]any) map[string]any {
	out := make(map[string]any, len(in))
	for k, v := range in {
		out[k] = v
	}
	return out
}

func stringValue(v any) string {
	if s, ok := v.(string); ok {
		return strings.TrimSpace(s)
	}
	return ""
}

func canonicalAipingModelName(modelName string) string {
	trimmed := strings.TrimSpace(modelName)
	switch strings.ToLower(trimmed) {
	case "kling-v1":
		return "Kling-V1"
	case "kling-v1-6":
		return "Kling-V1.6"
	case "kling-v2-6":
		return "Kling-V2.6"
	case "kling-v3":
		return "Kling-V3"
	case "kling-video-o1":
		return "Kling-Video-O1"
	case "kling-v3-omni":
		return "Kling-V3-Omni"
	default:
		return trimmed
	}
}

func firstString(m map[string]any, keys ...string) string {
	for _, key := range keys {
		if s := stringValue(m[key]); s != "" {
			return s
		}
	}
	return ""
}

func resolveRequestModelName(req relaycommon.TaskSubmitReq, info *relaycommon.RelayInfo) string {
	if info != nil && info.ChannelMeta != nil && info.IsModelMapped && info.UpstreamModelName != "" {
		return strings.TrimSpace(info.UpstreamModelName)
	}
	if modelName := firstString(req.Metadata, "model_name", "model"); modelName != "" {
		return modelName
	}
	return strings.TrimSpace(req.Model)
}

func isOmniModel(modelName string) bool {
	normalized := strings.ToLower(strings.TrimSpace(modelName))
	return normalized == "kling-video-o1" || normalized == "kling-v3-omni"
}