xinyin025/geekai
1
0
Fork 0
mirror of https://github.com/yangjian102621/geekai.git synced 2025-09-17 16:56:38 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

feat: add midjourney message receive handler

Browse Source
This commit is contained in:
RockYang 2023-08-14 07:09:52 +08:00
parent 302cb8a5be
commit 6e40f92aaf
8 changed files with 54 additions and 31 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
api/core/app_server.go
View File

@ -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,
}
}

5
api/core/types/web.go
View File

@ -12,8 +12,8 @@ type BizVo struct {
// WsMessage Websocket message
type WsMessage struct {
Type WsMsgType `json:"type"` // 消息类别，start, end
Content string `json:"content"`
Type WsMsgType `json:"type"` // 消息类别，start, end, img
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

20
api/handler/chat_handler.go
View File

@ -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 {
handler := ChatHandler{db: db}
func NewChatHandler(app *core.AppServer, db *gorm.DB, levelDB *store.LevelDB) *ChatHandler {
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)
}

29
api/handler/mj_handler.go
View File

@ -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"`
Content string `json:"content"`
Status TaskStatus `json:"status"`
Type string `json:"type"`
MessageId string `json:"message_id"`
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.ChatClients.Get(sessionId)
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)")
logger.Infof("Data: %+v", data)
wsClient := h.App.MjTaskClients.Get(key)
if wsClient == nil { // 客户端断线，则丢弃
resp.SUCCESS(c)
return
}
// 推送消息到客户端
// TODO: 增加绘画消息类型
utils.ReplyChunkMessage(wsClient, types.WsMessage{Type: types.WsImg, Content: data})
resp.ERROR(c, "Error with CallBack")
}

6
api/handler/sms_handler.go
View File

@ -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

12
api/handler/user_handler.go
View File

@ -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)
}

3
api/service/function/mid_journey.go
View File

@ -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)
}

6
api/utils/websocket.go
View File

@ -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})