From b81ab8412daf985dc156d621cba4ce1494b36d5b Mon Sep 17 00:00:00 2001
From: junleea <354425203@qq.com>
Date: Sat, 22 Mar 2025 17:11:42 +0800
Subject: [PATCH] =?UTF-8?q?=E9=87=8D=E5=86=99=E6=95=B0=E6=8D=AE=E8=8E=B7?=
 =?UTF-8?q?=E5=8F=96=E4=B8=8E=E6=95=B0=E6=8D=AE=E8=BF=94=E5=9B=9E=E7=BB=93?=
 =?UTF-8?q?=E6=9E=84?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 handler/im.go        | 10 ++--------
 proto/im.go          | 23 ++++++++++++++++-------
 service/imService.go | 11 +++++++----
 3 files changed, 25 insertions(+), 19 deletions(-)

diff --git a/handler/im.go b/handler/im.go
index e7a1911..644b947 100644
--- a/handler/im.go
+++ b/handler/im.go
@@ -101,7 +101,7 @@ func subscribeAndHandleIMMessages(ws *websocket.Conn, userId int) {
 				log.Println("Read error:", err)
 				break
 			}
-			var data proto.WSMessage
+			var data proto.WSMessageReq
 			err1 := json.Unmarshal(message, &data)
 			if err1 != nil {
 				log.Println("Error parsing JSON:", err1)
@@ -132,17 +132,11 @@ func subscribeAndHandleIMMessages(ws *websocket.Conn, userId int) {
 	}()
 
 	go func() {
-		var data proto.WSMessage
 		//从redis订阅消息
 		for m := range ch {
 			msg := m.Payload // 获取消息,消息格式为json
 			if msg != "" {
-				data.Msg = msg
-				data.Type = "msg"
-				data.SessionID = int(sessionID)
-				data.ToID = 0
-				msg_, _ := json.Marshal(data)
-				err2 := ws.WriteMessage(websocket.TextMessage, msg_)
+				err2 := ws.WriteMessage(websocket.TextMessage, []byte(msg))
 				if err2 != nil {
 					// 发生错误，删除连接
 					clientsMux.Lock()
diff --git a/proto/im.go b/proto/im.go
index 6751861..0524738 100644
--- a/proto/im.go
+++ b/proto/im.go
@@ -1,18 +1,27 @@
 package proto
 
+import "github.com/ollama/ollama/api"
+
 type AIQueueMessage struct {
-	Type string `json:"type"` //声明不同消息类型
-	ID   uint   `json:"id"`   //消息id
-	Msg  string `json:"msg"`  //原始json消息
+	Type string               `json:"type"` //声明不同消息类型
+	ID   uint                 `json:"id"`   //消息id
+	Msg  api.GenerateResponse `json:"msg"`  //原始消息
 }
 
 type WSMessage struct {
-	Type      string `json:"type"`       //接收及发送消息类型
-	Msg       string `json:"msg"`        //消息内容,只进行转发，不做处理
-	SessionID int    `json:"session_id"` //应用层会话id
-	ToID      int    `json:"to_id"`      //接收者id
+	Type      string         `json:"type"`       //接收及发送消息类型
+	Msg       AIQueueMessage `json:"msg"`        //消息内容,只进行转发，不做处理
+	SessionID int            `json:"session_id"` //应用层会话id
+	ToID      int            `json:"to_id"`      //接收者id
 }
 
 type ModelParam struct {
 	Model string `json:"model"` //模型名称
 }
+
+type WSMessageReq struct {
+	Type      string `json:"type"`       //接收及发送消息类型
+	Msg       string `json:"msg"`        //消息内容,只进行转发，不做处理
+	SessionID int    `json:"session_id"` //应用层会话id
+	ToID      int    `json:"to_id"`      //接收者id
+}
diff --git a/service/imService.go b/service/imService.go
index 1ed5d50..3033f6a 100644
--- a/service/imService.go
+++ b/service/imService.go
@@ -79,7 +79,7 @@ func CreateGeneralMessageService(fromID, toID, msgType, sessionID int, msg strin
 var client *api.Client
 var actxMap map[int][]int
 
-func WSReceiveMessageService(userID, sessionID int, channel string, msg proto.WSMessage) (error, uint) {
+func WSReceiveMessageService(userID, sessionID int, channel string, msg proto.WSMessageReq) (error, uint) {
 	var resErr error
 	var resID uint
 	//处理消息
@@ -101,10 +101,10 @@ func WSReceiveMessageService(userID, sessionID int, channel string, msg proto.WS
 		ctx := context.Background()
 		robotMsg := ""
 		var aiMsg proto.AIQueueMessage
+		var data proto.WSMessage
 		respFunc := func(resp api.GenerateResponse) error {
 			aiMsg.Type = "ollama"
-			respStr, _ := json.Marshal(resp)
-			aiMsg.Msg = string(respStr)
+			aiMsg.Msg = resp
 			robotMsg += resp.Response
 			if resp.Done { //该消息完成
 				actx = resp.Context
@@ -121,7 +121,10 @@ func WSReceiveMessageService(userID, sessionID int, channel string, msg proto.WS
 				resID = msgID
 			}
 			//发送消息
-			aiMsgStr, _ := json.Marshal(aiMsg)
+			data.Msg = aiMsg
+			data.SessionID = sessionID
+			data.ToID = userID
+			aiMsgStr, _ := json.Marshal(data)
 			worker.Publish(channel, string(aiMsgStr), time.Second*60)
 			return nil
 		}