aiChat/internal/model/ws.go

package model

import (
	"aiChat/utility/fsw"
	. "app.yhyue.com/moapp/jybase/common"
	"app.yhyue.com/moapp/jybase/date"
	"app.yhyue.com/moapp/jybase/encrypt"
	"bufio"
	"context"
	"encoding/json"
	"fmt"
	"github.com/gogf/gf/v2/encoding/gjson"
	"github.com/gogf/gf/v2/frame/g"
	"github.com/gogf/gf/v2/net/ghttp"
	"github.com/gogf/gf/v2/os/glog"
	"github.com/gogf/gf/v2/util/gconv"
	"io"
	"time"
)

type WsChat struct {
	Ctx context.Context
}

func NewMessage(ctx context.Context) *WsChat {
	return &WsChat{
		Ctx: ctx,
	}
}

// Handle 处理消息
func (m *WsChat) Handle(ws *ghttp.WebSocket, msg []byte) {
	defer Catch()
	jSession := SessionCtx.Get(m.Ctx).JSession
	if jSession.PositionId == 0 {
		_ = ws.WriteJSON(g.Map{
			"error_code": -1,
			"error_msg":  "请登录",
		})
		return
	}
	req, from := &QuestionReq{}, 0
	if err := gjson.Unmarshal(msg, req); err != nil {
		glog.Errorf(m.Ctx, "%d 接收消息Unmarshal出错:%v", jSession.PositionId, err)
		return
	}
	questionId := ChatHistory.Save(m.Ctx, &ChatRecord{
		Content:    req.Prompt,
		Type:       1,
		Refer:      req.Href,
		PersonId:   jSession.PositionId,
		CreateTime: time.Now().Format(date.Date_Full_Layout),
	})
	content, buf, replyId, errMsg := func() (string, *bufio.Reader, int64, error) {
		var (
			err   error
			buf   *bufio.Reader
			reply string
		)
		errReply := func() string {
			// 校验是否在黑名单,黑名单不返回内容
			if UserBlackList.CheckBlackList(m.Ctx, jSession.PositionId) {
				return g.Cfg().MustGet(m.Ctx, "limit.blackMsg").String()
			}
			// 校验问答频率
			if ChatLimit.GetBucket(m.Ctx, jSession.PositionId).TakeAvailable(1) == 0 {
				return g.Cfg().MustGet(m.Ctx, "limit.exceedMsg").String()
			}
			// 问题敏感词过滤
			if fsw.Match(req.Prompt) {
				return g.Cfg().MustGet(m.Ctx, "limit.fswMsg").String()
			}
			return ""
		}()
		if errReply != "" {
			reply, from = errReply, -1
		} else {
			reply, buf, from, err = Question.DetailQuestion(m.Ctx, req)
			if err != nil {
				g.Log().Error(m.Ctx, "问答异常", err)
				reply, from = g.Cfg().MustGet(m.Ctx, "limit.errMsg").String(), -1
			}
		}

		if from != Answer_ChatGPT {
			if reply == "" {
				reply, from = g.Cfg().MustGet(m.Ctx, "limit.emptyMsg").String(), -1
			}
			replyId := ChatHistory.Save(m.Ctx, &ChatRecord{
				Content:    reply,
				Type:       2,
				Actions:    gconv.Int(If(errReply == "", 1, 0)),
				QuestionId: questionId,
				PersonId:   jSession.PositionId,
				Item:       gconv.Int(If(errReply == "", from, -1)),
				CreateTime: time.Now().Format(date.Date_Full_Layout),
			})
			if replyId <= 0 {
				g.Log().Error(m.Ctx, "问答存储存储异常")
			}
			return reply, nil, replyId, nil
		}
		return reply, buf, 0, nil
	}()
	if from != Answer_ChatGPT {
		if errMsg != nil {
			_ = ws.WriteJSON(g.Map{"error_code": -1, "error_msg": errMsg.Error(), "data": nil})
		} else {
			_ = ws.WriteJSON(g.Map{"error_code": 0, "error_msg": "", "data": g.Map{"id": encrypt.SE.Encode2Hex(fmt.Sprintf("%d", replyId)), "reply": content, "isEnd": true}})
		}
	} else if buf != nil {
		for {
			line, _, err := buf.ReadLine()
			if err == io.EOF {
				break
			}
			if _, data := parseEventStream(line); data != nil {
				data.Response = fsw.Repl(data.Response)
				if data.Finished {
					replyId := ChatHistory.Save(m.Ctx, &ChatRecord{
						Content:    data.Response,
						Type:       2,
						Actions:    1,
						QuestionId: questionId,
						PersonId:   jSession.PositionId,
						Item:       Answer_ChatGPT,
						CreateTime: time.Now().Format(date.Date_Full_Layout),
					})
					_ = ws.WriteJSON(g.Map{"error_code": 0, "error_msg": "", "data": g.Map{"id": encrypt.SE.Encode2Hex(fmt.Sprintf("%d", replyId)), "reply": data.Response, "isEnd": data.Finished}})
				} else {
					_ = ws.WriteJSON(g.Map{"error_code": 0, "error_msg": "", "data": g.Map{"reply": data.Response, "isEnd": data.Finished}})
				}
			}
		}
	}
}

func parseEventStream(line []byte) (event string, date *BufRes) {
	// 如果行以 "event:" 开头，表示这是一个事件的标识符
	if len(line) > 6 && string(line[:6]) == "event:" {
		event = string(line[6 : len(line)-1])
		return event, nil
	}
	// 如果行以 "data:" 开头，表示这是事件的数据部分
	if len(line) > 5 && string(line[:5]) == "data:" {
		date = &BufRes{}
		if err := json.Unmarshal(line[5:len(line)], date); err == nil {
			return
		}
		return event, nil
	}
	return "", nil
}