feat: midjourney image variation function is ready

api/core/types/function.go

@@ -83,15 +83,15 @@ var InnerFunctions = []Function{
 			Properties: map[string]Property{
 				"prompt": {
 					Type:        "string",
-					Description: "绘画内容描述，提示词，此参数需要翻译成英文",
+					Description: "绘画内容描述，提示词，如果该参数中有中文的话，则需要翻译成英文",
 				},
 				"ar": {
 					Type:        "string",
-					Description: "图片长宽比，如 16:9, --ar 3:2",
+					Description: "图片长宽比，如 --ar 4:3",
 				},
 				"niji": {
 					Type:        "string",
-					Description: "动漫模型版本，如 --niji 5",
+					Description: "动漫模型版本，例如 --niji 5",
 				},
 			},
 			Required: []string{},

api/handler/mj_handler.go

@@ -8,6 +8,7 @@ import (
 	"chatplus/store/model"
 	"chatplus/utils"
 	"chatplus/utils/resp"
+	"fmt"
 	"github.com/gin-gonic/gin"
 	"gorm.io/gorm"
 	"time"
@@ -83,7 +84,7 @@ func (h *MidJourneyHandler) Notify(c *gin.Context) {
 		err := h.leveldb.Get(types.TaskStorePrefix+data.Key, &task)
 		if err != nil {
 			logger.Error("error with get MidJourney task: ", err)
-			resp.ERROR(c, err.Error())
+			resp.SUCCESS(c)
 			return
 		}
 
@@ -138,16 +139,18 @@ func (h *MidJourneyHandler) Notify(c *gin.Context) {
 	resp.SUCCESS(c, "SUCCESS")
 }
 
+type reqVo struct {
+	Index       int32  `json:"index"`
+	MessageId   string `json:"message_id"`
+	MessageHash string `json:"message_hash"`
+	SessionId   string `json:"session_id"`
+	Key         string `json:"key"`
+	Prompt      string `json:"prompt"`
+}
+
 // Upscale send upscale command to MidJourney Bot
 func (h *MidJourneyHandler) Upscale(c *gin.Context) {
-	var data struct {
-		Index       int32  `json:"index"`
-		MessageId   string `json:"message_id"`
-		MessageHash string `json:"message_hash"`
-		SessionId   string `json:"session_id"`
-		Key         string `json:"key"`
-	}
-
+	var data reqVo
 	if err := c.ShouldBindJSON(&data); err != nil ||
 		data.SessionId == "" ||
 		data.Key == "" {
@@ -170,7 +173,39 @@ func (h *MidJourneyHandler) Upscale(c *gin.Context) {
 		return
 	}
 
-	utils.ReplyMessage(wsClient, "已推送放大图片任务到 MidJourney 机器人，请耐心等待任务执行...")
+	content := fmt.Sprintf("**%s** 已推送 Upscale 任务到 MidJourney 机器人，请耐心等待任务执行...", data.Prompt)
+	utils.ReplyMessage(wsClient, content)
+	if h.App.MjTaskClients.Get(data.Key) == nil {
+		h.App.MjTaskClients.Put(data.Key, wsClient)
+	}
+	resp.SUCCESS(c)
+}
+
+func (h *MidJourneyHandler) Variation(c *gin.Context) {
+	var data reqVo
+	if err := c.ShouldBindJSON(&data); err != nil ||
+		data.SessionId == "" ||
+		data.Key == "" {
+		resp.ERROR(c, types.InvalidArgs)
+		return
+	}
+	wsClient := h.App.ChatClients.Get(data.SessionId)
+	if wsClient == nil {
+		resp.ERROR(c, "No Websocket client online")
+		return
+	}
+
+	err := h.mjFunc.Variation(function.MjVariationReq{
+		Index:       data.Index,
+		MessageId:   data.MessageId,
+		MessageHash: data.MessageHash,
+	})
+	if err != nil {
+		resp.ERROR(c, err.Error())
+		return
+	}
+	content := fmt.Sprintf("**%s** 已推送 Variation 任务到 MidJourney 机器人，请耐心等待任务执行...", data.Prompt)
+	utils.ReplyMessage(wsClient, content)
 	if h.App.MjTaskClients.Get(data.Key) == nil {
 		h.App.MjTaskClients.Put(data.Key, wsClient)
 	}

api/main.go

@@ -179,6 +179,7 @@ func main() {
 		fx.Invoke(func(s *core.AppServer, h *handler.MidJourneyHandler) {
 			s.Engine.POST("/api/mj/notify", h.Notify)
 			s.Engine.POST("/api/mj/upscale", h.Upscale)
+			s.Engine.POST("/api/mj/variation", h.Variation)
 		}),
 
 		// 管理后台控制器

api/service/function/mid_journey.go

@@ -81,7 +81,32 @@ func (f FuncMidJourney) Upscale(upReq MjUpscaleReq) error {
 	if res.Code != types.Success {
 		return errors.New(res.Message)
 	}
-	
+
+	return nil
+}
+
+type MjVariationReq struct {
+	Index       int32  `json:"index"`
+	MessageId   string `json:"message_id"`
+	MessageHash string `json:"message_hash"`
+}
+
+func (f FuncMidJourney) Variation(upReq MjVariationReq) error {
+	url := fmt.Sprintf("%s/api/mj/variation", f.config.ApiURL)
+	var res types.BizVo
+	r, err := f.client.R().
+		SetHeader("Authorization", f.config.Token).
+		SetHeader("Content-Type", "application/json").
+		SetBody(upReq).
+		SetSuccessResult(&res).Post(url)
+	if err != nil || r.IsErrorState() {
+		return fmt.Errorf("%v%v", r.String(), err)
+	}
+
+	if res.Code != types.Success {
+		return errors.New(res.Message)
+	}
+
 	return nil
 }