session_manager.go 6.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223
  1. package ws
  2. import (
  3. "encoding/json"
  4. "errors"
  5. chatService "eta/eta_api/services/llm"
  6. "eta/eta_api/utils"
  7. "eta/eta_api/utils/llm"
  8. "eta/eta_api/utils/llm/eta_llm"
  9. "eta/eta_api/utils/llm/eta_llm/eta_llm_http"
  10. "fmt"
  11. "github.com/gorilla/websocket"
  12. "net/http"
  13. "strings"
  14. "sync"
  15. "time"
  16. )
  17. var (
  18. llmService, _ = llm.GetInstance(llm.ETA_LLM_CLIENT)
  19. )
  20. const (
  21. defaultCheckInterval = 2 * time.Minute // 检测间隔应小于心跳超时时间
  22. connectionTimeout = 10 * time.Minute // 客户端超时时间
  23. TcpTimeout = 20 * time.Minute // TCP超时时间,保底关闭,覆盖会话超时时间
  24. ReadTimeout = 15 * time.Minute // 读取超时时间,保底关闭,覆盖会话超时时间
  25. writeWaitTimeout = 60 * time.Second //写入超时时间
  26. )
  27. type ConnectionManager struct {
  28. Sessions sync.Map
  29. ticker *time.Ticker
  30. stopChan chan struct{}
  31. }
  32. var (
  33. smOnce sync.Once
  34. manager *ConnectionManager
  35. )
  36. func GetInstance() *ConnectionManager {
  37. smOnce.Do(func() {
  38. if manager == nil {
  39. manager = &ConnectionManager{
  40. ticker: time.NewTicker(defaultCheckInterval),
  41. stopChan: make(chan struct{}),
  42. }
  43. }
  44. })
  45. return manager
  46. }
  47. func Manager() *ConnectionManager {
  48. return manager
  49. }
  50. // HandleMessage 消息处理核心逻辑
  51. func (manager *ConnectionManager) HandleMessage(userID int, sessionID string, message []byte) error {
  52. session, exists := manager.GetSession(sessionID)
  53. if !exists {
  54. return errors.New("session not found")
  55. }
  56. if strings.ToLower(string(message)) == "pong" {
  57. session.UpdateActivity()
  58. fmt.Printf("收到心跳消息,续期长连接:%v", session.LastActive)
  59. return nil
  60. }
  61. if !Allow(userID, QA_LIMITER) {
  62. _ = session.Conn.WriteMessage(websocket.TextMessage, []byte("您提问的太频繁了,请稍后再试"))
  63. _ = session.Conn.WriteMessage(websocket.TextMessage, []byte("<EOF/>"))
  64. return nil
  65. }
  66. var userMessage Message
  67. err := json.Unmarshal(message, &userMessage)
  68. if err != nil {
  69. utils.FileLog.Error(fmt.Sprintf("消息格式错误:%s", string(message)))
  70. fmt.Printf("消息格式错误:%s", string(message))
  71. return errors.New("消息格式错误:" + err.Error())
  72. }
  73. // 处理业务逻辑
  74. //session.History = append(session.History, userMessage.LastTopics...)
  75. redisHisChat, err := chatService.GetChatRecordsFromRedis(userMessage.ChatId)
  76. if err != nil {
  77. utils.FileLog.Error("获取历史对话数据失败,err:", err.Error())
  78. } else {
  79. for _, chat := range redisHisChat {
  80. his := eta_llm_http.HistoryContent{
  81. Content: chat.Content,
  82. Role: chat.ChatUserType,
  83. }
  84. hisMsg, _ := json.Marshal(&his)
  85. if len(hisMsg) != 0 {
  86. session.History = append(session.History, hisMsg)
  87. }
  88. }
  89. }
  90. resp, err := llmService.KnowledgeBaseChat(userMessage.Query, userMessage.KbName, session.History)
  91. defer func() {
  92. if resp != nil && resp.Body != nil && err == nil {
  93. _ = resp.Body.Close()
  94. }
  95. }()
  96. if resp == nil {
  97. utils.FileLog.Error("知识库问答失败: 无应答")
  98. return errors.New("知识库问答失败: 无应答")
  99. }
  100. if err != nil {
  101. utils.FileLog.Error(fmt.Sprintf("知识库问答失败: httpCode:%d,错误信息:%s", resp.StatusCode, http.StatusText(resp.StatusCode)))
  102. err = errors.New(fmt.Sprintf("知识库问答失败: httpCode:%d,错误信息:%s", resp.StatusCode, http.StatusText(resp.StatusCode)))
  103. return err
  104. }
  105. if resp.StatusCode != http.StatusOK {
  106. utils.FileLog.Error(fmt.Sprintf("知识库问答失败: httpCode:%d,错误信息:%s", resp.StatusCode, http.StatusText(resp.StatusCode)))
  107. err = errors.New(fmt.Sprintf("知识库问答失败: httpCode:%d,错误信息:%s", resp.StatusCode, http.StatusText(resp.StatusCode)))
  108. return err
  109. }
  110. // 解析流式响应
  111. contentChan, errChan, closeChan := eta_llm.ParseStreamResponse(resp)
  112. emptyContent := true
  113. // 处理流式数据并发送到 WebSocket
  114. for {
  115. select {
  116. case content, ok := <-contentChan:
  117. if !ok {
  118. err = errors.New("未知的错误异常")
  119. return err
  120. }
  121. session.UpdateActivity()
  122. if emptyContent {
  123. emptyContent = false
  124. }
  125. // 发送消息到 WebSocket
  126. _ = session.Conn.WriteMessage(websocket.TextMessage, []byte(content))
  127. case chanErr, ok := <-errChan:
  128. if !ok {
  129. err = errors.New("未知的错误异常")
  130. } else {
  131. err = errors.New(chanErr.Error())
  132. }
  133. // 发送错误消息到 WebSocket
  134. return err
  135. case <-closeChan:
  136. if emptyContent {
  137. _ = session.Conn.WriteMessage(websocket.TextMessage, []byte("<think></think>"))
  138. }
  139. _ = session.Conn.WriteMessage(websocket.TextMessage, []byte("<EOF/>"))
  140. return nil
  141. }
  142. }
  143. // 更新最后活跃时间
  144. // 发送响应
  145. //return session.Conn.WriteMessage(websocket.TextMessage, []byte(response))
  146. }
  147. // AddSession Add 添加一个新的会话
  148. func (manager *ConnectionManager) AddSession(session *Session) {
  149. manager.Sessions.Store(session.Id, session)
  150. }
  151. func (manager *ConnectionManager) GetSessionId(userId int, sessionId string) (sessionID string) {
  152. return fmt.Sprintf("%d_%s", userId, sessionId)
  153. }
  154. // RemoveSession Remove 移除一个会话
  155. func (manager *ConnectionManager) RemoveSession(sessionCode string) {
  156. fmt.Printf("移除会话: SessionID=%s, UserID=%s", sessionCode, sessionCode)
  157. manager.Sessions.Delete(sessionCode)
  158. }
  159. // GetSession 获取一个会话
  160. func (manager *ConnectionManager) GetSession(sessionCode string) (session *Session, exists bool) {
  161. if data, ok := manager.Sessions.Load(sessionCode); ok {
  162. session = data.(*Session)
  163. exists = ok
  164. }
  165. return
  166. }
  167. // CheckAll 批量检测所有连接
  168. func (manager *ConnectionManager) CheckAll() {
  169. manager.Sessions.Range(func(key, value interface{}) bool {
  170. session := value.(*Session)
  171. // 判断超时
  172. if time.Since(session.LastActive) > 2*connectionTimeout {
  173. fmt.Printf("连接超时关闭: SessionID=%s, UserID=%s", session.Id, session.UserId)
  174. utils.FileLog.Warn("连接超时关闭: SessionID=%s, UserID=%s", session.Id, session.UserId)
  175. session.Close()
  176. return true
  177. }
  178. // 发送心跳
  179. go func(s *Session) {
  180. err := s.Conn.WriteControl(websocket.PingMessage,
  181. nil, time.Now().Add(writeWaitTimeout))
  182. if err != nil {
  183. fmt.Printf("心跳发送失败: SessionID=%s, Error=%v", s.Id, err)
  184. utils.FileLog.Warn("心跳发送失败: SessionID=%s, Error=%v",
  185. s.Id, err)
  186. fmt.Println("心跳无响应,退出请求")
  187. session.Close()
  188. }
  189. }(session)
  190. return true
  191. })
  192. }
  193. // Start 启动心跳检测
  194. func (manager *ConnectionManager) Start() {
  195. defer manager.ticker.Stop()
  196. for {
  197. select {
  198. case <-manager.ticker.C:
  199. manager.CheckAll()
  200. case <-manager.stopChan:
  201. return
  202. }
  203. }
  204. }
  205. // Stop 停止心跳检测
  206. func (manager *ConnectionManager) Stop() {
  207. close(manager.stopChan)
  208. }