xinyin025/geekai
1
0
Fork 0
mirror of https://github.com/yangjian102621/geekai.git synced 2025-11-08 10:13:44 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

feat: chat chrawing function is refactored

Browse Source
This commit is contained in:
RockYang
2023-09-13 06:57:25 +08:00
parent fa341bab30
commit 1d3acc8ed3
12 changed files with 121 additions and 279 deletions
Download Patch File Download Diff File Expand all files Collapse all files

152
api/handler/mj_handler.go
View File

@@ -4,7 +4,6 @@ import (
"chatplus/core"
"chatplus/core/types"
"chatplus/service"
"chatplus/service/function"
"chatplus/service/oss"
"chatplus/store/model"
"chatplus/utils"
@@ -22,8 +21,6 @@ import (
type TaskStatus string
const (
Start = TaskStatus("Started")
Running = TaskStatus("Running")
Stopped = TaskStatus("Stopped")
Finished = TaskStatus("Finished")
)
@@ -64,42 +61,58 @@ func NewMidJourneyHandler(
return &h
}
type notifyData struct {
MessageId string `json:"message_id"`
ReferenceId string `json:"reference_id"`
Image Image `json:"image"`
Content string `json:"content"`
Prompt string `json:"prompt"`
Status TaskStatus `json:"status"`
Progress int `json:"progress"`
}
func (h *MidJourneyHandler) Notify(c *gin.Context) {
token := c.GetHeader("Authorization")
if token != h.App.Config.ExtConfig.Token {
resp.NotAuth(c)
return
}
var data struct {
MessageId string `json:"message_id"`
ReferenceId string `json:"reference_id"`
Image Image `json:"image"`
Content string `json:"content"`
Prompt string `json:"prompt"`
Status TaskStatus `json:"status"`
Progress int `json:"progress"`
}
var data notifyData
if err := c.ShouldBindJSON(&data); err != nil || data.Prompt == "" {
resp.ERROR(c, types.InvalidArgs)
return
}
logger.Debugf("收到 MidJourney 回调请求：%+v", data)
h.lock.Lock()
defer h.lock.Unlock()
taskString, err := h.redis.Get(c, service.MjRunningJobKey).Result()
err := h.notifyHandler(c, data)
if err != nil {
resp.SUCCESS(c) // 过期任务，丢弃
resp.ERROR(c, err.Error())
return
}
// 解除任务锁定
if data.Status == Finished || data.Status == Stopped {
h.redis.Del(c, service.MjRunningJobKey)
}
resp.SUCCESS(c)
}
func (h *MidJourneyHandler) notifyHandler(c *gin.Context, data notifyData) error {
taskString, err := h.redis.Get(c, service.MjRunningJobKey).Result()
if err != nil { // 过期任务，丢弃
logger.Warn("任务已过期：", err)
return nil
}
var task service.MjTask
err = utils.JsonDecode(taskString, &task)
if err != nil {
resp.SUCCESS(c) // 非标准任务，丢弃
return
if err != nil { // 非标准任务，丢弃
logger.Warn("任务解析失败：", err)
return nil
}
if task.Src == service.TaskSrcImg { // 绘画任务
@@ -107,19 +120,20 @@ func (h *MidJourneyHandler) Notify(c *gin.Context) {
var job model.MidJourneyJob
res := h.db.First(&job, task.Id)
if res.Error != nil {
resp.SUCCESS(c) // 非法任务，丢弃
return
logger.Warn("非法任务：", err)
return nil
}
job.MessageId = data.MessageId
job.ReferenceId = data.ReferenceId
job.Progress = data.Progress
job.Prompt = data.Prompt
// download image
if data.Progress == 100 {
imgURL, err := h.uploaderManager.GetUploadHandler().PutImg(data.Image.URL)
if err != nil {
resp.ERROR(c, "error with download img: "+err.Error())
return
logger.Error("error with download img: ", err.Error())
return err
}
job.ImgURL = imgURL
} else {
@@ -128,18 +142,16 @@ func (h *MidJourneyHandler) Notify(c *gin.Context) {
}
res = h.db.Updates(&job)
if res.Error != nil {
resp.ERROR(c, "error with update job: "+err.Error())
return
logger.Error("error with update job: ", err.Error())
return res.Error
}
resp.SUCCESS(c)
} else if task.Src == service.TaskSrcChat { // 聊天任务
var job model.MidJourneyJob
res := h.db.Where("message_id = ?", data.MessageId).First(&job)
if res.Error == nil {
resp.SUCCESS(c)
return
logger.Warn("重复消息：", data.MessageId)
return nil
}
wsClient := h.App.MjTaskClients.Get(task.Id)
@@ -156,10 +168,10 @@ func (h *MidJourneyHandler) Notify(c *gin.Context) {
content := fmt.Sprintf("**%s** 图片下载失败：%s", data.Prompt, err.Error())
utils.ReplyMessage(wsClient, content)
}
resp.ERROR(c, err.Error())
return
return err
}
tx := h.db.Begin()
data.Image.URL = imgURL
message := model.HistoryMessage{
UserId: uint(task.UserId),
@@ -171,9 +183,9 @@ func (h *MidJourneyHandler) Notify(c *gin.Context) {
Tokens: 0,
UseContext: false,
}
res := h.db.Create(&message)
res = tx.Create(&message)
if res.Error != nil {
logger.Error("error with save chat history message: ", res.Error)
return res.Error
}
// save the job
@@ -184,16 +196,17 @@ func (h *MidJourneyHandler) Notify(c *gin.Context) {
job.ImgURL = imgURL
job.Progress = data.Progress
job.CreatedAt = time.Now()
res = h.db.Create(&job)
res = tx.Create(&job)
if res.Error != nil {
logger.Error("error with save MidJourney Job: ", res.Error)
tx.Rollback()
return res.Error
}
tx.Commit()
}
if wsClient == nil { // 客户端断线，则丢弃
logger.Errorf("Client is offline: %+v", data)
resp.SUCCESS(c, "Client is offline")
return
return nil
}
if data.Status == Finished {
@@ -202,22 +215,17 @@ func (h *MidJourneyHandler) Notify(c *gin.Context) {
// delete client
h.App.MjTaskClients.Delete(task.Id)
} else {
//// 使用代理临时转发图片
//if data.Image.URL != "" {
// image, err := utils.DownloadImage(data.Image.URL, h.App.Config.ProxyURL)
// if err == nil {
// data.Image.URL = "data:image/png;base64," + base64.StdEncoding.EncodeToString(image)
// }
//}
data.Image.URL = fmt.Sprintf("/api/mj/proxy?url=%s", data.Image.URL)
utils.ReplyChunkMessage(wsClient, types.WsMessage{Type: types.WsMjImg, Content: data})
}
resp.SUCCESS(c, "SUCCESS")
}
return nil
}
func (h *MidJourneyHandler) Proxy(c *gin.Context) {
logger.Info(c.Request.Host, c.Request.Proto)
return
url := c.Query("url")
image, err := utils.DownloadImage(url, h.App.Config.ProxyURL)
if err != nil {
@@ -232,16 +240,16 @@ type reqVo struct {
MessageId string `json:"message_id"`
MessageHash string `json:"message_hash"`
SessionId string `json:"session_id"`
Key string `json:"key"`
Prompt string `json:"prompt"`
ChatId string `json:"chat_id"`
RoleId int `json:"role_id"`
Icon string `json:"icon"`
}
// Upscale send upscale command to MidJourney Bot
func (h *MidJourneyHandler) Upscale(c *gin.Context) {
var data reqVo
if err := c.ShouldBindJSON(&data); err != nil ||
data.SessionId == "" ||
data.Key == "" {
if err := c.ShouldBindJSON(&data); err != nil || data.SessionId == "" {
resp.ERROR(c, types.InvalidArgs)
return
}
@@ -250,35 +258,32 @@ func (h *MidJourneyHandler) Upscale(c *gin.Context) {
resp.ERROR(c, "No Websocket client online")
return
}
userId, _ := c.Get(types.LoginUserID)
h.mjService.PushTask(service.MjTask{
Id: data.SessionId,
Src: service.TaskSrcChat,
Type: service.Upscale,
Prompt: data.Prompt,
UserId: utils.IntValue(utils.InterfaceToString(userId), 0),
RoleId: data.RoleId,
Icon: data.Icon,
ChatId: data.ChatId,
Index: data.Index,
MessageId: data.MessageId,
MessageHash: data.MessageHash,
})
err := n.Upscale(function.MjUpscaleReq{
Index: data.Index,
MessageId: data.MessageId,
MessageHash: data.MessageHash,
})
if err != nil {
resp.ERROR(c, err.Error())
return
}
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)
if h.App.MjTaskClients.Get(data.SessionId) == nil {
h.App.MjTaskClients.Put(data.SessionId, 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 == "" {
if err := c.ShouldBindJSON(&data); err != nil || data.SessionId == "" {
resp.ERROR(c, types.InvalidArgs)
return
}
@@ -288,19 +293,24 @@ func (h *MidJourneyHandler) Variation(c *gin.Context) {
return
}
err := h.mjFunc.Variation(function.MjVariationReq{
userId, _ := c.Get(types.LoginUserID)
h.mjService.PushTask(service.MjTask{
Id: data.SessionId,
Src: service.TaskSrcChat,
Type: service.Variation,
Prompt: data.Prompt,
UserId: utils.IntValue(utils.InterfaceToString(userId), 0),
RoleId: data.RoleId,
Icon: data.Icon,
ChatId: data.ChatId,
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)
if h.App.MjTaskClients.Get(data.SessionId) == nil {
h.App.MjTaskClients.Put(data.SessionId, wsClient)
}
resp.SUCCESS(c)
}