feat: add midjourney message receive handler

api/core/app_server.go

@@ -34,7 +34,7 @@ type AppServer struct {
 	ChatClients   *types.LMap[string, *types.WsClient]    // map[sessionId]Websocket 连接集合
 	ReqCancelFunc *types.LMap[string, context.CancelFunc] // HttpClient 请求取消 handle function
 	Functions     map[string]function.Function
-	MjTasks       *types.LMap[string, types.MjTask]
+	MjTaskClients *types.LMap[string, *types.WsClient]
 }
 
 func NewServer(appConfig *types.AppConfig, functions map[string]function.Function) *AppServer {
@@ -48,7 +48,7 @@ func NewServer(appConfig *types.AppConfig, functions map[string]function.Functio
 		ChatSession:   types.NewLMap[string, types.ChatSession](),
 		ChatClients:   types.NewLMap[string, *types.WsClient](),
 		ReqCancelFunc: types.NewLMap[string, context.CancelFunc](),
-		MjTasks:       types.NewLMap[string, types.MjTask](),
+		MjTaskClients: types.NewLMap[string, *types.WsClient](),
 		Functions:     functions,
 	}
 }

api/core/types/web.go

@@ -12,8 +12,8 @@ type BizVo struct {
 
 // WsMessage Websocket message
 type WsMessage struct {
-	Type    WsMsgType `json:"type"` // 消息类别，start, end
+	Type    WsMsgType   `json:"type"` // 消息类别，start, end, img
-	Content string    `json:"content"`
+	Content interface{} `json:"content"`
 }
 type WsMsgType string
 
@@ -21,6 +21,7 @@ const (
 	WsStart  = WsMsgType("start")
 	WsMiddle = WsMsgType("middle")
 	WsEnd    = WsMsgType("end")
+	WsImg    = WsMsgType("img")
 )
 
 type BizCode int

api/handler/chat_handler.go

@@ -5,6 +5,7 @@ import (
 	"bytes"
 	"chatplus/core"
 	"chatplus/core/types"
+	"chatplus/store"
 	"chatplus/store/model"
 	"chatplus/store/vo"
 	"chatplus/utils"
@@ -26,14 +27,16 @@ import (
 )
 
 const ErrorMsg = "抱歉，AI 助手开小差了，请稍后再试。"
+const TaskStorePrefix = "/tasks/"
 
 type ChatHandler struct {
 	BaseHandler
-	db *gorm.DB
+	db      *gorm.DB
+	leveldb *store.LevelDB
 }
 
-func NewChatHandler(app *core.AppServer, db *gorm.DB) *ChatHandler {
+func NewChatHandler(app *core.AppServer, db *gorm.DB, levelDB *store.LevelDB) *ChatHandler {
-	handler := ChatHandler{db: db}
+	handler := ChatHandler{db: db, leveldb: levelDB}
 	handler.App = app
 	return &handler
 }
@@ -133,7 +136,7 @@ func (h *ChatHandler) ChatHandle(c *gin.Context) {
 }
 
 // 将消息发送给 ChatGPT 并获取结果，通过 WebSocket 推送到客户端
-func (h *ChatHandler) sendMessage(ctx context.Context, session types.ChatSession, role model.ChatRole, prompt string, ws types.Client) error {
+func (h *ChatHandler) sendMessage(ctx context.Context, session types.ChatSession, role model.ChatRole, prompt string, ws *types.WsClient) error {
 	promptCreatedAt := time.Now() // 记录提问时间
 
 	var user model.User
@@ -340,13 +343,18 @@ func (h *ChatHandler) sendMessage(ctx context.Context, session types.ChatSession
 					if functionName == types.FuncMidJourney {
 						key := utils.Sha256(data)
 						// add task for MidJourney
-						h.App.MjTasks.Put(key, types.MjTask{
+						h.App.MjTaskClients.Put(key, ws)
+						task := types.MjTask{
 							UserId: userVo.Id,
 							RoleId: role.Id,
 							Icon:   role.Icon,
 							Client: ws,
 							ChatId: session.ChatId,
-						})
+						}
+						err := h.leveldb.Put(TaskStorePrefix+key, task)
+						if err != nil {
+							logger.Error("error with store MidJourney task: ", err)
+						}
 						content = fmt.Sprintf("绘画提示词：%s 已推送任务到 MidJourney 机器人，请耐心等待任务执行...", data)
 					}
 

api/handler/mj_handler.go

@@ -24,6 +24,7 @@ type Image struct {
 	Width    int    `json:"width"`
 	Height   int    `json:"height"`
 	Size     int    `json:"size"`
+	Hash     string `json:"hash"`
 }
 
 type MidJourneyHandler struct {
@@ -44,18 +45,30 @@ func (h *MidJourneyHandler) Notify(c *gin.Context) {
 	}
 
 	var data struct {
-		Image   Image      `json:"image"`
+		Type      string     `json:"type"`
-		Content string     `json:"content"`
+		MessageId string     `json:"message_id"`
-		Status  TaskStatus `json:"status"`
+		Image     Image      `json:"image"`
+		Content   string     `json:"content"`
+		Prompt    string     `json:"prompt"`
+		Status    TaskStatus `json:"status"`
+		Key       string     `json:"key"`
 	}
-	if err := c.ShouldBindJSON(&data); err != nil {
+	if err := c.ShouldBindJSON(&data); err != nil || data.Prompt == "" {
 		resp.ERROR(c, types.InvalidArgs)
 		return
 	}
+	key := utils.Sha256(data.Prompt)
+	data.Key = key
+	// TODO: 如果绘画任务完成了则将该消息保存到当前会话的聊天历史记录
 
-	sessionId := "u7blnft9zqisyrwidjb22j6b78iqc30lv9jtud3k9o"
+	wsClient := h.App.MjTaskClients.Get(key)
-	wsClient := h.App.ChatClients.Get(sessionId)
+	if wsClient == nil { // 客户端断线，则丢弃
-	utils.ReplyMessage(wsClient, "![](https://cdn.discordapp.com/attachments/1138713254718361633/1139482452579070053/lal603743923_A_Chinese_girl_walking_barefoot_on_the_beach_weari_df8b6dc0-3b13-478c-8dbb-983015d21661.png)")
+		resp.SUCCESS(c)
-	logger.Infof("Data: %+v", data)
+		return
+	}
+
+	// 推送消息到客户端
+	// TODO: 增加绘画消息类型
+	utils.ReplyChunkMessage(wsClient, types.WsMessage{Type: types.WsImg, Content: data})
 	resp.ERROR(c, "Error with CallBack")
 }

api/handler/sms_handler.go

@@ -14,13 +14,13 @@ const CodeStorePrefix = "/verify/codes/"
 
 type SmsHandler struct {
 	BaseHandler
-	db      *store.LevelDB
+	leveldb *store.LevelDB
 	sms     *service.AliYunSmsService
 	captcha *service.CaptchaService
 }
 
 func NewSmsHandler(app *core.AppServer, db *store.LevelDB, sms *service.AliYunSmsService, captcha *service.CaptchaService) *SmsHandler {
-	handler := &SmsHandler{db: db, sms: sms, captcha: captcha}
+	handler := &SmsHandler{leveldb: db, sms: sms, captcha: captcha}
 	handler.App = app
 	return handler
 }
@@ -50,7 +50,7 @@ func (h *SmsHandler) SendCode(c *gin.Context) {
 	}
 
 	// 存储验证码，等待后面注册验证
-	err = h.db.Put(CodeStorePrefix+data.Mobile, code)
+	err = h.leveldb.Put(CodeStorePrefix+data.Mobile, code)
 	if err != nil {
 		resp.ERROR(c, "验证码保存失败")
 		return

api/handler/user_handler.go

@@ -22,11 +22,11 @@ type UserHandler struct {
 	BaseHandler
 	db       *gorm.DB
 	searcher *xdb.Searcher
-	levelDB  *store.LevelDB
+	leveldb  *store.LevelDB
 }
 
 func NewUserHandler(app *core.AppServer, db *gorm.DB, searcher *xdb.Searcher, levelDB *store.LevelDB) *UserHandler {
-	handler := &UserHandler{db: db, searcher: searcher, levelDB: levelDB}
+	handler := &UserHandler{db: db, searcher: searcher, leveldb: levelDB}
 	handler.App = app
 	return handler
 }
@@ -60,7 +60,7 @@ func (h *UserHandler) Register(c *gin.Context) {
 	key := CodeStorePrefix + data.Mobile
 	if h.App.SysConfig.EnabledMsgService {
 		var code int
-		err := h.levelDB.Get(key, &code)
+		err := h.leveldb.Get(key, &code)
 		if err != nil || code != data.Code {
 			logger.Info(code)
 			resp.ERROR(c, "短信验证码错误")
@@ -118,7 +118,7 @@ func (h *UserHandler) Register(c *gin.Context) {
 	}
 
 	if h.App.SysConfig.EnabledMsgService {
-		_ = h.levelDB.Delete(key) // 注册成功，删除短信验证码
+		_ = h.leveldb.Delete(key) // 注册成功，删除短信验证码
 	}
 	resp.SUCCESS(c, user)
 }
@@ -366,7 +366,7 @@ func (h *UserHandler) BindMobile(c *gin.Context) {
 	// 检查验证码
 	key := CodeStorePrefix + data.Mobile
 	var code int
-	err := h.levelDB.Get(key, &code)
+	err := h.leveldb.Get(key, &code)
 	if err != nil || code != data.Code {
 		resp.ERROR(c, "短信验证码错误")
 		return
@@ -384,6 +384,6 @@ func (h *UserHandler) BindMobile(c *gin.Context) {
 		return
 	}
 
-	_ = h.levelDB.Delete(key) // 删除短信验证码
+	_ = h.leveldb.Delete(key) // 删除短信验证码
 	resp.SUCCESS(c)
 }

api/service/function/mid_journey.go

@@ -36,12 +36,13 @@ func (f FuncMidJourney) Invoke(params map[string]interface{}) (string, error) {
 		delete(params, "ar")
 	}
 	prompt = prompt + " --niji 5"
+	url := fmt.Sprintf("%s/api/mj/image", f.config.ApiURL)
 	var res types.BizVo
 	r, err := f.client.R().
 		SetHeader("Authorization", f.config.Token).
 		SetHeader("Content-Type", "application/json").
 		SetBody(params).
-		SetSuccessResult(&res).Post(f.config.ApiURL)
+		SetSuccessResult(&res).Post(url)
 	if err != nil || r.IsErrorState() {
 		return "", fmt.Errorf("%v%v", r.String(), err)
 	}

api/utils/websocket.go

@@ -9,20 +9,20 @@ import (
 var logger = logger2.GetLogger()
 
 // ReplyChunkMessage 回复客户片段端消息
-func ReplyChunkMessage(client types.Client, message types.WsMessage) {
+func ReplyChunkMessage(client *types.WsClient, message types.WsMessage) {
 	msg, err := json.Marshal(message)
 	if err != nil {
 		logger.Errorf("Error for decoding json data: %v", err.Error())
 		return
 	}
-	err = client.(*types.WsClient).Send(msg)
+	err = client.Send(msg)
 	if err != nil {
 		logger.Errorf("Error for reply message: %v", err.Error())
 	}
 }
 
 // ReplyMessage 回复客户端一条完整的消息
-func ReplyMessage(ws types.Client, message string) {
+func ReplyMessage(ws *types.WsClient, message interface{}) {
 	ReplyChunkMessage(ws, types.WsMessage{Type: types.WsStart})
 	ReplyChunkMessage(ws, types.WsMessage{Type: types.WsMiddle, Content: message})
 	ReplyChunkMessage(ws, types.WsMessage{Type: types.WsEnd})