add ws handler

api/handler/chatimpl/chat_handler.go

@@ -151,9 +151,9 @@ func (h *ChatHandler) ChatHandle(c *gin.Context) {
 			err = h.sendMessage(ctx, session, chatRole, utils.InterfaceToString(message.Content), client)
 			if err != nil {
 				logger.Error(err)
-				utils.ReplyMessage(client, err.Error())
+				utils.SendMessage(client, err.Error())
 			} else {
-				utils.ReplyChunkMessage(client, types.ReplyMessage{Type: types.WsEnd})
+				utils.SendChunkMessage(client, types.ReplyMessage{Type: types.WsMsgTypeEnd})
 				logger.Infof("回答完毕: %v", message.Content)
 			}
 

api/handler/chatimpl/openai_handler.go

@@ -108,7 +108,7 @@ func (h *ChatHandler) sendOpenAiMessage(
 			}
 
 			if responseBody.Choices[0].FinishReason == "stop" && len(contents) == 0 {
-				utils.ReplyMessage(ws, "抱歉😔😔😔，AI助手由于未知原因已经停止输出内容。")
+				utils.SendMessage(ws, "抱歉😔😔😔，AI助手由于未知原因已经停止输出内容。")
 				break
 			}
 
@@ -136,7 +136,7 @@ func (h *ChatHandler) sendOpenAiMessage(
 				if res.Error == nil {
 					toolCall = true
 					callMsg := fmt.Sprintf("正在调用工具 `%s` 作答 ...\n\n", function.Label)
-					utils.ReplyChunkMessage(ws, types.ReplyMessage{Type: types.WsContent, Content: callMsg})
+					utils.SendChunkMessage(ws, types.ReplyMessage{Type: types.WsMsgTypeContent, Content: callMsg})
 					contents = append(contents, callMsg)
 				}
 				continue
@@ -153,8 +153,8 @@ func (h *ChatHandler) sendOpenAiMessage(
 			} else {
 				content := responseBody.Choices[0].Delta.Content
 				contents = append(contents, utils.InterfaceToString(content))
-				utils.ReplyChunkMessage(ws, types.ReplyMessage{
-					Type:    types.WsContent,
+				utils.SendChunkMessage(ws, types.ReplyMessage{
+					Type:    types.WsMsgTypeContent,
 					Content: utils.InterfaceToString(responseBody.Choices[0].Delta.Content),
 				})
 			}
@@ -186,14 +186,14 @@ func (h *ChatHandler) sendOpenAiMessage(
 			}
 			if errMsg != "" || apiRes.Code != types.Success {
 				msg := "调用函数工具出错：" + apiRes.Message + errMsg
-				utils.ReplyChunkMessage(ws, types.ReplyMessage{
-					Type:    types.WsContent,
+				utils.SendChunkMessage(ws, types.ReplyMessage{
+					Type:    types.WsMsgTypeContent,
 					Content: msg,
 				})
 				contents = append(contents, msg)
 			} else {
-				utils.ReplyChunkMessage(ws, types.ReplyMessage{
-					Type:    types.WsContent,
+				utils.SendChunkMessage(ws, types.ReplyMessage{
+					Type:    types.WsMsgTypeContent,
 					Content: apiRes.Data,
 				})
 				contents = append(contents, utils.InterfaceToString(apiRes.Data))
@@ -226,7 +226,7 @@ func (h *ChatHandler) sendOpenAiMessage(
 		if strings.HasPrefix(req.Model, "o1-") {
 			content = fmt.Sprintf("AI思考结束，耗时：%d 秒。\n%s", time.Now().Unix()-session.Start, respVo.Choices[0].Message.Content)
 		}
-		utils.ReplyMessage(ws, content)
+		utils.SendMessage(ws, content)
 		respVo.Usage.Prompt = prompt
 		respVo.Usage.Content = content
 		h.saveChatHistory(req, respVo.Usage, respVo.Choices[0].Message, chatCtx, session, role, userVo, promptCreatedAt, time.Now())

api/handler/markmap_handler.go

@@ -87,7 +87,7 @@ func (h *MarkMapHandler) Client(c *gin.Context) {
 				logger.Error(err)
 				utils.ReplyErrorMessage(client, err.Error())
 			} else {
-				utils.ReplyMessage(client, types.ReplyMessage{Type: types.WsEnd})
+				utils.SendMessage(client, types.ReplyMessage{Type: types.WsMsgTypeEnd})
 			}
 
 		}
@@ -170,13 +170,13 @@ func (h *MarkMapHandler) sendMessage(client *types.WsClient, prompt string, mode
 				break
 			}
 
-			utils.ReplyChunkMessage(client, types.ReplyMessage{
-				Type:    types.WsContent,
+			utils.SendChunkMessage(client, types.ReplyMessage{
+				Type:    types.WsMsgTypeContent,
 				Content: utils.InterfaceToString(responseBody.Choices[0].Delta.Content),
 			})
 		} // end for
 
-		utils.ReplyChunkMessage(client, types.ReplyMessage{Type: types.WsEnd})
+		utils.SendChunkMessage(client, types.ReplyMessage{Type: types.WsMsgTypeEnd})
 
 	} else {
 		body, _ := io.ReadAll(response.Body)

api/handler/ws_handler.go

@@ -0,0 +1,85 @@
+package handler
+
+// * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
+// * Copyright 2023 The Geek-AI Authors. All rights reserved.
+// * Use of this source code is governed by a Apache-2.0 license
+// * that can be found in the LICENSE file.
+// * @Author yangjian102621@163.com
+// * +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
+
+import (
+	"geekai/core"
+	"geekai/core/types"
+	"geekai/service"
+	"geekai/store/model"
+	"geekai/utils"
+	"github.com/gin-gonic/gin"
+	"github.com/gorilla/websocket"
+	"net/http"
+)
+
+// Websocket 连接处理 handler
+
+type WebsocketHandler struct {
+	BaseHandler
+	wsService *service.WebsocketService
+}
+
+func NewWebsocketHandler(app *core.AppServer, s *service.WebsocketService) *WebsocketHandler {
+	return &WebsocketHandler{
+		BaseHandler: BaseHandler{App: app},
+		wsService:   s,
+	}
+}
+
+func (h *WebsocketHandler) Client(c *gin.Context) {
+	ws, err := (&websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return true }}).Upgrade(c.Writer, c.Request, nil)
+	if err != nil {
+		logger.Error(err)
+		c.Abort()
+		return
+	}
+
+	userId := h.GetInt(c, "user_id", 0)
+	clientId := c.Query("client")
+	client := types.NewWsClient(ws)
+	if userId == 0 {
+		_ = client.Send([]byte("Invalid user_id"))
+		c.Abort()
+		return
+	}
+	var user model.User
+	if err := h.DB.Where("id", userId).First(&user).Error; err != nil {
+		_ = client.Send([]byte("Invalid user_id"))
+		c.Abort()
+		return
+	}
+
+	h.wsService.Clients.Put(clientId, client)
+	logger.Infof("New websocket connected, IP: %s", c.RemoteIP())
+	go func() {
+		for {
+			_, msg, err := client.Receive()
+			if err != nil {
+				logger.Debugf("close connection: %s", client.Conn.RemoteAddr())
+				client.Close()
+			}
+
+			var message types.InputMessage
+			err = utils.JsonDecode(string(msg), &message)
+			if err != nil {
+				continue
+			}
+
+			logger.Infof("Receive a message:%+v", message)
+			if message.Type == types.WsMsgTypePing {
+				_ = client.Send([]byte(`{"type":"pong"}`))
+			}
+
+			switch message.Channel {
+
+			}
+
+		}
+	}()
+}