From 28fdb8af37f0025f705ef54fd1779fa3be38f601 Mon Sep 17 00:00:00 2001
From: feitianbubu <feitianbubu@qq.com>
Date: Sun, 10 Aug 2025 15:55:28 +0800
Subject: [PATCH] feat: add jimeng video official api

---
 middleware/distributor.go    |  4 ++-
 middleware/jimeng_adapter.go | 66 ++++++++++++++++++++++++++++++++++++
 relay/relay_task.go          |  3 ++
 router/video-router.go       |  8 +++++
 4 files changed, 80 insertions(+), 1 deletion(-)
 create mode 100644 middleware/jimeng_adapter.go

diff --git a/middleware/distributor.go b/middleware/distributor.go
index dea30abf..b7e5ef3c 100644
--- a/middleware/distributor.go
+++ b/middleware/distributor.go
@@ -174,7 +174,9 @@ func getModelRequest(c *gin.Context) (*ModelRequest, bool, error) {
 			relayMode = relayconstant.RelayModeVideoFetchByID
 			shouldSelectChannel = false
 		}
-		c.Set("relay_mode", relayMode)
+		if _, ok := c.Get("relay_mode"); !ok {
+			c.Set("relay_mode", relayMode)
+		}
 	} else if strings.HasPrefix(c.Request.URL.Path, "/v1beta/models/") || strings.HasPrefix(c.Request.URL.Path, "/v1/models/") {
 		// Gemini API 路径处理: /v1beta/models/gemini-2.0-flash:generateContent
 		relayMode := relayconstant.RelayModeGemini
diff --git a/middleware/jimeng_adapter.go b/middleware/jimeng_adapter.go
new file mode 100644
index 00000000..ce5e1467
--- /dev/null
+++ b/middleware/jimeng_adapter.go
@@ -0,0 +1,66 @@
+package middleware
+
+import (
+	"bytes"
+	"encoding/json"
+	"github.com/gin-gonic/gin"
+	"io"
+	"net/http"
+	"one-api/common"
+	"one-api/constant"
+	relayconstant "one-api/relay/constant"
+)
+
+func JimengRequestConvert() func(c *gin.Context) {
+	return func(c *gin.Context) {
+		action := c.Query("Action")
+		if action == "" {
+			abortWithOpenAiMessage(c, http.StatusBadRequest, "Action query parameter is required")
+			return
+		}
+
+		// Handle Jimeng official API request
+		var originalReq map[string]interface{}
+		if err := common.UnmarshalBodyReusable(c, &originalReq); err != nil {
+			abortWithOpenAiMessage(c, http.StatusBadRequest, "Invalid request body")
+			return
+		}
+		model, _ := originalReq["req_key"].(string)
+		prompt, _ := originalReq["prompt"].(string)
+
+		unifiedReq := map[string]interface{}{
+			"model":    model,
+			"prompt":   prompt,
+			"metadata": originalReq,
+		}
+
+		jsonData, err := json.Marshal(unifiedReq)
+		if err != nil {
+			abortWithOpenAiMessage(c, http.StatusInternalServerError, "Failed to marshal request body")
+			return
+		}
+
+		// Update request body
+		c.Request.Body = io.NopCloser(bytes.NewBuffer(jsonData))
+		c.Set(common.KeyRequestBody, jsonData)
+
+		if image, ok := originalReq["image"]; !ok || image == "" {
+			c.Set("action", constant.TaskActionTextGenerate)
+		}
+
+		c.Request.URL.Path = "/v1/video/generations"
+
+		if action == "CVSync2AsyncGetResult" {
+			taskId, ok := originalReq["task_id"].(string)
+			if !ok || taskId == "" {
+				abortWithOpenAiMessage(c, http.StatusBadRequest, "task_id is required for CVSync2AsyncGetResult")
+				return
+			}
+			c.Request.URL.Path = "/v1/video/generations/" + taskId
+			c.Request.Method = http.MethodGet
+			c.Set("task_id", taskId)
+			c.Set("relay_mode", relayconstant.RelayModeVideoFetchByID)
+		}
+		c.Next()
+	}
+}
diff --git a/relay/relay_task.go b/relay/relay_task.go
index ce00527b..0ccc3b33 100644
--- a/relay/relay_task.go
+++ b/relay/relay_task.go
@@ -258,6 +258,9 @@ func sunoFetchByIDRespBodyBuilder(c *gin.Context) (respBody []byte, taskResp *dt
 
 func videoFetchByIDRespBodyBuilder(c *gin.Context) (respBody []byte, taskResp *dto.TaskError) {
 	taskId := c.Param("task_id")
+	if taskId == "" {
+		taskId = c.GetString("task_id")
+	}
 	userId := c.GetInt("id")
 
 	originTask, exist, err := model.GetByTaskId(userId, taskId)
diff --git a/router/video-router.go b/router/video-router.go
index 0bd8cd83..bcc05eae 100644
--- a/router/video-router.go
+++ b/router/video-router.go
@@ -23,4 +23,12 @@ func SetVideoRouter(router *gin.Engine) {
 		klingV1Router.GET("/videos/text2video/:task_id", controller.RelayTask)
 		klingV1Router.GET("/videos/image2video/:task_id", controller.RelayTask)
 	}
+
+	// Jimeng official API routes - direct mapping to official API format
+	jimengOfficialGroup := router.Group("jimeng")
+	jimengOfficialGroup.Use(middleware.JimengRequestConvert(), middleware.TokenAuth(), middleware.Distribute())
+	{
+		// Maps to: /?Action=CVSync2AsyncSubmitTask&Version=2022-08-31 and /?Action=CVSync2AsyncGetResult&Version=2022-08-31
+		jimengOfficialGroup.POST("/", controller.RelayTask)
+	}
 }