Prechádzať zdrojové kódy

feat: add jimeng video official api

feitianbubu 4 mesiacov pred
rodič
commit
28fdb8af37

+ 3 - 1
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

+ 66 - 0
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()
+	}
+}

+ 3 - 0
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)

+ 8 - 0
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)
+	}
 }