saw-go/handler/im.go

package handler

import (
	"StuAcaWorksAI/dao"
	"StuAcaWorksAI/proto"
	"StuAcaWorksAI/service"
	"StuAcaWorksAI/worker"
	"context"
	"encoding/json"
	"errors"
	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"log"
	"net/http"
	"strconv"
	"sync"
)

type SMessage struct {
	Type      int    `json:"type" form:"type"`
	Msg       string `json:"msg" form:"msg"`
	ToID      int    `json:"to_id" form:"to_id"`
	SessionID int    `json:"session_id" form:"session_id"`
}

var (
	upgrader = websocket.Upgrader{
		ReadBufferSize:  1024,
		WriteBufferSize: 1024,
		CheckOrigin: func(r *http.Request) bool {
			// 允许所有来源的连接
			return true
		},
	}
)

// 创建一个用于存储WebSocket连接的map和互斥锁
var (
	clients    = make(map[*websocket.Conn]bool)
	clientsMux sync.Mutex
)

func SetUpIMGroup(router *gin.Engine) {
	imGroup := router.Group("/im")
	imGroup.POST("/send_message", SendMessage)
	imGroup.GET("/ai_chat_ws", SRMessage)
	imGroup.POST("/get_message", GetMessage)
}

func GetMessage(c *gin.Context) {
	var req SMessage
	id, _ := c.Get("id")
	userID := int(id.(float64))
	if err := c.ShouldBind(&req); err == nil {
		err2, messages := service.FindMessageBySessionID(req.SessionID, userID)
		if err2 == nil {
			c.JSON(http.StatusOK, gin.H{"code": proto.SuccessCode, "message": "success", "data": messages})
		} else {
			c.JSON(http.StatusOK, gin.H{"error": err2.Error(), "code": proto.MsgSendFailed, "message": "failed"})
		}
	} else {
		c.JSON(http.StatusOK, gin.H{"error": err.Error(), "code": proto.ParameterError, "message": "failed"})
	}
}

func SendMessage(c *gin.Context) {
	var req SMessage
	id, _ := c.Get("id")
	userID := int(id.(float64))
	if err := c.ShouldBind(&req); err == nil {
		var err2 error
		var mid uint
		err2, mid = service.CreateGeneralMessageService(userID, req.ToID, req.Type, req.SessionID, req.Msg)
		if err2 == nil {
			c.JSON(http.StatusOK, gin.H{"code": proto.SuccessCode, "message": "success", "data": mid})
		} else {
			c.JSON(http.StatusOK, gin.H{"error": err2.Error(), "code": proto.MsgSendFailed, "message": "failed"})
		}
	} else {
		c.JSON(http.StatusOK, gin.H{"error": err.Error(), "code": proto.ParameterError, "message": "failed"})
	}
}

// ws长连接
func SRMessage(c *gin.Context) {
	id, _ := c.Get("id")
	userID := int(id.(float64))
	// 升级HTTP连接为WebSocket连接
	ws, err1 := upgrader.Upgrade(c.Writer, c.Request, nil)
	clients[ws] = true
	if err1 != nil {
		log.Println("升级为WebSocket时发生错误:", err1)
		return
	}
	subscribeAndHandleIMMessages(ws, userID)
}

func subscribeAndHandleIMMessages(ws *websocket.Conn, userId int) {
	ctx := context.Background()
	chanel := "user_" + strconv.Itoa(userId) + "_ai_chat_msg"
	pubsub := worker.RedisClient.Subscribe(ctx, chanel)
	defer pubsub.Close()
	defer ws.Close()
	ch := pubsub.Channel()
	//定时器，每3秒发送一次心跳
	//ticker := time.NewTicker(3 * time.Second)
	//defer ticker.Stop()

	var sessionID uint
	sessionID = 0

	//接收客户端消息,发送消息
	done := make(chan struct{})
	go func() {
		defer close(done)
		for {
			_, message, err := ws.ReadMessage()
			log.Println("rec message:", string(message))
			if err != nil {
				log.Println("Read error:", err)
				break
			}
			var data proto.WSMessageReq
			err1 := json.Unmarshal(message, &data)
			if err1 != nil {
				log.Println("Error parsing JSON:", err1)
			}
			err = doReceiveGenChatMessage(userId, &sessionID, &data)
			if err != nil {
				log.Println("do receive gen chat message error:", err)
				return
			}
		}
	}()

	go func() {
		//defer close(done)
		//从redis订阅消息
		for m := range ch {
			msg := m.Payload // 获取消息,消息格式为json
			if msg != "" {
				err2 := ws.WriteMessage(websocket.TextMessage, []byte(msg))
				if err2 != nil {
					// 发生错误，删除连接
					clientsMux.Lock()
					delete(clients, ws)
					clientsMux.Unlock()
					break
				}
			}
		}
		//查看done是否关闭
	}()

	//发送心跳包
	for {
		select {
		case <-done:
			//结束会话将context保存
			ctx2, ok := service.SessionContextMap[int(sessionID)]
			if !ok {
				log.Println("session context not found")
				return
			}
			err := service.UpdateSessionByID(int(sessionID), userId, "", ctx2)
			if err != nil {
				log.Println("update session context error:", err)
			}
			//保存的会话
			delete(service.SessionContextMap, int(sessionID))
			//客户端断开连接
			return
			//case t := <-ticker.C:
			//	//发送心跳
			//	err := ws.WriteMessage(websocket.PingMessage, nil)
			//	if err != nil {
			//		log.Println("Ping error:", err, t)
			//		return
			//	}
			//}
		}
	}

}

func doReceiveGenChatMessage(userId int, sessionID *uint, data *proto.WSMessageReq) error {
	var err error
	chanel := "user_" + strconv.Itoa(userId) + "_ai_chat_msg"
	if *sessionID == 0 {
		//创建会话
		err, *sessionID = service.CreateSession(userId, data.Msg)
		log.Println("create session id:", sessionID)
		if err != nil {
			log.Println("Create session error:", err)
		}
	}
	var model dao.Model
	//查看请求功能类型
	if data.Function == "gen-ai-chat" {
		model = service.FindFuncModelByFunction(data.Function, userId)[0]
	} else {
		return errors.New("function not exist")
	}
	//解析模型参数
	var modelParam proto.ModelParam
	err = json.Unmarshal([]byte(model.Parameter), &modelParam)
	if err != nil {
		return err
	}
	modelParam.Url = model.Url

	var userMsgID uint
	//将消息存入数据库
	err, userMsgID = service.CreateMessage(proto.UserToModelMsgType, int(*sessionID), userId, int(model.ID), data.Msg, proto.MsgHasRead)
	log.Println("create user message id:", userMsgID)
	if err != nil {
		return err
	} //机器人id为2,消息类型为1,状态为1
	//需要后续处理
	err, _ = service.WSReceiveMessageService(userId, int(*sessionID), chanel, *data, model, modelParam)
	return err
}