feat: midjourney drawing image function is ready

api/handler/chat_handler.go

@@ -27,7 +27,6 @@ import (
 )
 
 const ErrorMsg = "抱歉，AI 助手开小差了，请稍后再试。"
-const TaskStorePrefix = "/tasks/"
 
 type ChatHandler struct {
 	BaseHandler
@@ -342,16 +341,16 @@ func (h *ChatHandler) sendMessage(ctx context.Context, session types.ChatSession
 					content := data
 					if functionName == types.FuncMidJourney {
 						key := utils.Sha256(data)
+						//logger.Info(data, ",", key)
 						// add task for MidJourney
 						h.App.MjTaskClients.Put(key, ws)
 						task := types.MjTask{
 							UserId: userVo.Id,
 							RoleId: role.Id,
-							Icon:   role.Icon,
-							Client: ws,
+							Icon:   "/images/avatar/mid_journey.png",
 							ChatId: session.ChatId,
 						}
-						err := h.leveldb.Put(TaskStorePrefix+key, task)
+						err := h.leveldb.Put(types.TaskStorePrefix+key, task)
 						if err != nil {
 							logger.Error("error with store MidJourney task: ", err)
 						}

api/handler/mj_handler.go

@@ -3,9 +3,12 @@ package handler
 import (
 	"chatplus/core"
 	"chatplus/core/types"
+	"chatplus/store"
+	"chatplus/store/model"
 	"chatplus/utils"
 	"chatplus/utils/resp"
 	"github.com/gin-gonic/gin"
+	"gorm.io/gorm"
 )
 
 type TaskStatus string
@@ -29,10 +32,12 @@ type Image struct {
 
 type MidJourneyHandler struct {
 	BaseHandler
+	leveldb *store.LevelDB
+	db      *gorm.DB
 }
 
-func NewMidJourneyHandler(app *core.AppServer) *MidJourneyHandler {
-	h := MidJourneyHandler{}
+func NewMidJourneyHandler(app *core.AppServer, leveldb *store.LevelDB, db *gorm.DB) *MidJourneyHandler {
+	h := MidJourneyHandler{leveldb: leveldb, db: db}
 	h.App = app
 	return &h
 }
@@ -57,18 +62,54 @@ func (h *MidJourneyHandler) Notify(c *gin.Context) {
 		resp.ERROR(c, types.InvalidArgs)
 		return
 	}
+
 	key := utils.Sha256(data.Prompt)
 	data.Key = key
-	// TODO: 如果绘画任务完成了则将该消息保存到当前会话的聊天历史记录
+	//logger.Info(data.Prompt, ",", key)
+	if data.Status == Finished {
+		var task types.MjTask
+		err := h.leveldb.Get(types.TaskStorePrefix+key, &task)
+		if err != nil {
+			logger.Error("error with get MidJourney task: ", err)
+			resp.ERROR(c)
+			return
+		}
 
-	wsClient := h.App.MjTaskClients.Get(key)
-	if wsClient == nil { // 客户端断线，则丢弃
-		resp.SUCCESS(c)
-		return
+		// TODO: 是否需要把图片下载到本地服务器？
+
+		historyUserMsg := model.HistoryMessage{
+			UserId:     task.UserId,
+			ChatId:     task.ChatId,
+			RoleId:     task.RoleId,
+			Type:       types.MjMsg,
+			Icon:       task.Icon,
+			Content:    utils.JsonEncode(data),
+			Tokens:     0,
+			UseContext: false,
+		}
+		res := h.db.Save(&historyUserMsg)
+		if res.Error != nil {
+			logger.Error("error with save MidJourney message: ", res.Error)
+		}
+
+		// delete task from leveldb
+		_ = h.leveldb.Delete(types.TaskStorePrefix + key)
 	}
 
 	// 推送消息到客户端
-	// TODO: 增加绘画消息类型
-	utils.ReplyChunkMessage(wsClient, types.WsMessage{Type: types.WsImg, Content: data})
-	resp.ERROR(c, "Error with CallBack")
+	wsClient := h.App.MjTaskClients.Get(key)
+	if wsClient == nil { // 客户端断线，则丢弃
+		resp.SUCCESS(c, "Client is offline")
+		return
+	}
+
+	if data.Status == Finished {
+		utils.ReplyChunkMessage(wsClient, types.WsMessage{Type: types.WsMjImg, Content: data})
+		utils.ReplyChunkMessage(wsClient, types.WsMessage{Type: types.WsEnd})
+		// delete client
+		h.App.MjTaskClients.Delete(key)
+	} else {
+		utils.ReplyChunkMessage(wsClient, types.WsMessage{Type: types.WsMjImg, Content: data})
+	}
+	resp.SUCCESS(c, "SUCCESS")
 }