package rag import ( "eta/eta_api/controllers" "eta/eta_api/models" "eta/eta_api/models/system" "eta/eta_api/services/llm/facade" "eta/eta_api/utils" "eta/eta_api/utils/ws" "github.com/gorilla/websocket" "net" "net/http" "time" ) type ChatController struct { controllers.BaseAuthController } func (cc *ChatController) Prepare() { cc.SysUser = cc.Ctx.Input.GetData("admin").(*system.Admin) if cc.SysUser == nil || cc.SysUser.AdminId == 0 { utils.FileLog.Error("用户信息不存在") cc.Ctx.ResponseWriter.WriteHeader(http.StatusBadRequest) return } } // NewChat @Title 新建对话框 // @Description 测试知识库问答 // @Success 101 {object} response.ListResp // @router /chat/new_chat [post] func (kbctrl *KbController) NewChat() { br := new(models.BaseResponse).Init() defer func() { kbctrl.Data["json"] = br kbctrl.ServeJSON() }() sysUser := kbctrl.SysUser if sysUser == nil { br.Msg = "请登录" br.ErrMsg = "请登录,SysUser Is Empty" br.Ret = 408 return } //searchResp, err := facade.LLMKnowledgeBaseChat(req) //if err != nil { // br.Msg = "知识库问答" // br.ErrMsg = "知识库问答:" + err.Error() // return //} br.Ret = 200 br.Success = true br.Msg = "创建成功" } // ChatConnect @Title 知识库问答创建对话连接 // @Description 知识库问答创建对话连接 // @Success 101 {object} response.ListResp // @router /chat/connect [get] func (cc *ChatController) ChatConnect() { if !ws.Allow(cc.SysUser.AdminId, ws.CONNECT_LIMITER) { utils.FileLog.Error("WebSocket连接太频繁,主动拒绝链接") cc.Ctx.ResponseWriter.WriteHeader(http.StatusTooManyRequests) return } wsCon, err := webSocketHandler(cc.Ctx.ResponseWriter, cc.Ctx.Request) if err != nil { utils.FileLog.Error("WebSocket连接失败:", err) cc.Ctx.ResponseWriter.WriteHeader(http.StatusBadRequest) return } facade.AddSession(cc.SysUser.AdminId, wsCon) } // upGrader 用于将HTTP连接升级为WebSocket连接 var upGrader = websocket.Upgrader{ ReadBufferSize: 1024, WriteBufferSize: 1024, CheckOrigin: func(r *http.Request) bool { return true }, } // WebSocketHandler 处理WebSocket连接 func webSocketHandler(w http.ResponseWriter, r *http.Request) (conn *websocket.Conn, err error) { conn, err = upGrader.Upgrade(w, r, nil) if err != nil { utils.FileLog.Error("升级协议失败:WebSocket:%s", err.Error()) return } // 获取底层 TCP 连接并设置保活 if tcpConn, ok := conn.NetConn().(*net.TCPConn); ok { _ = tcpConn.SetKeepAlive(true) _ = tcpConn.SetKeepAlivePeriod(ws.TcpTimeout) utils.FileLog.Info("TCP KeepAlive 已启用") } _ = conn.SetReadDeadline(time.Now().Add(ws.ReadTimeout)) return }