eta_api/utils/ws/session_manager.go

package ws

import (
	"encoding/json"
	"errors"
	"eta/eta_api/models/rag"
	chatService "eta/eta_api/services/llm"
	"eta/eta_api/utils"
	"eta/eta_api/utils/llm"
	"eta/eta_api/utils/llm/eta_llm"
	"eta/eta_api/utils/llm/eta_llm/eta_llm_http"
	"fmt"
	"github.com/gorilla/websocket"
	"net/http"
	"regexp"
	"strings"
	"sync"
	"time"
)

var (
	llmService, _ = llm.GetInstance(llm.ETA_LLM_CLIENT)
)

const (
	defaultCheckInterval = 2 * time.Minute  // 检测间隔应小于心跳超时时间
	connectionTimeout    = 10 * time.Minute // 客户端超时时间
	TcpTimeout           = 20 * time.Minute // TCP超时时间，保底关闭，覆盖会话超时时间
	ReadTimeout          = 15 * time.Minute // 读取超时时间，保底关闭，覆盖会话超时时间
	writeWaitTimeout     = 60 * time.Second //写入超时时间
)

type ConnectionManager struct {
	Sessions sync.Map
	ticker   *time.Ticker
	stopChan chan struct{}
}

var (
	smOnce  sync.Once
	manager *ConnectionManager
)

func GetInstance() *ConnectionManager {
	smOnce.Do(func() {
		if manager == nil {
			manager = &ConnectionManager{
				ticker:   time.NewTicker(defaultCheckInterval),
				stopChan: make(chan struct{}),
			}
		}
	})
	return manager
}
func Manager() *ConnectionManager {
	return manager
}

// HandleMessage 消息处理核心逻辑
func (manager *ConnectionManager) HandleMessage(userID int, sessionID string, message []byte) {
	var err error
	session, exists := manager.GetSession(sessionID)
	if !exists {
		err = errors.New("session not found")
		return
	}

	if strings.ToLower(string(message)) == "pong" {
		session.UpdateActivity()
		fmt.Printf("收到心跳消息,续期长连接：%v", session.LastActive)
		return
	}
	defer func() {
		if err != nil {
			//写应答
			_ = session.writeWithTimeout("<think></think>")
			_ = session.writeWithTimeout(err.Error())
			_ = session.writeWithTimeout("<EOF/>")
		}
	}()
	var userMessage Message
	err = json.Unmarshal(message, &userMessage)
	if err != nil {
		utils.FileLog.Error(fmt.Sprintf("消息格式错误:%s", string(message)))
		fmt.Printf("消息格式错误:%s", string(message))
		err = errors.New("消息格式错误:" + err.Error())
		return
	}

	if userMessage.MessageType == `stop` {
		if session.LLMStatus == 1 {
			// 标记llm提问状态：暂停提问
			session.LLMStatus = -1
		}
		if session.CloseLlmChan != nil {
			*session.CloseLlmChan <- true
		}
		return
	}

	// 限流
	if !Allow(userID, QA_LIMITER) {
		_ = session.Conn.WriteMessage(websocket.TextMessage, []byte("<think></think>"))
		_ = session.Conn.WriteMessage(websocket.TextMessage, []byte("您提问的太频繁了,请稍后再试"))
		_ = session.Conn.WriteMessage(websocket.TextMessage, []byte("<EOF/>"))
		return
	}
	// 处理业务逻辑
	//session.History = append(session.History, userMessage.LastTopics...)
	redisHisChat, err := chatService.GetChatRecordsFromRedis(userMessage.ChatId)
	if err != nil {
		utils.FileLog.Error("获取历史对话数据失败,err:", err.Error())
	} else {
		for _, chat := range redisHisChat {
			his := eta_llm_http.HistoryContent{
				Content: chat.Content,
				Role:    chat.ChatUserType,
			}
			hisMsg, _ := json.Marshal(&his)
			if len(hisMsg) != 0 {
				session.History = append(session.History, hisMsg)
			}
		}
	}
	//修改逻辑。如果问题出现敏感词，则返回敏感词提示
	var resp *http.Response
	labels := llm.GetDAFHandlerInstance().FindTextTagLabels(userMessage.Query)
	if len(labels) > 0 {
		articles, findErr := rag.GetArticleByTags(labels, 15, 10)
		if findErr != nil {
			utils.FileLog.Warn("没有搜索到相关的研报内容,执行RAG对话 ")
			resp, err = llmService.KnowledgeBaseChat(userMessage.Query, userMessage.KbName, session.History)
		} else {
			//直接对话，不需要走RAG
			var articlesContents string
			articlesCounts := len(articles)
			for i := 0; i < articlesCounts; i++ {
				articlesContents += fmt.Sprintf("【%d】：%s\n", i+1, articles[i].TextContent)
			}
			promote := fmt.Sprintf("【问题】：%s,请基于以下%d篇研报进行回答问题，以下是研报：%s", userMessage.Query, articlesCounts, articlesContents)
			re := regexp.MustCompile(`\s+`)
			promote = re.ReplaceAllString(promote, "")
			resp, err = llmService.CompletionChat(promote, session.History)
		}
	} else {
		resp, err = llmService.KnowledgeBaseChat(userMessage.Query, userMessage.KbName, session.History)
	}
	defer func() {
		if resp != nil && resp.Body != nil && err == nil {
			_ = resp.Body.Close()
		}
	}()
	if resp == nil {
		utils.FileLog.Error("知识库问答失败: 无应答")
		err = errors.New("知识库问答失败: 无应答")
		return
	}
	if err != nil {
		utils.FileLog.Error(fmt.Sprintf("知识库问答失败: httpCode:%d,错误信息:%s", resp.StatusCode, http.StatusText(resp.StatusCode)))
		err = errors.New(fmt.Sprintf("知识库问答失败: httpCode:%d,错误信息:%s", resp.StatusCode, http.StatusText(resp.StatusCode)))
		return
	}

	if resp.StatusCode != http.StatusOK {
		utils.FileLog.Error(fmt.Sprintf("知识库问答失败: httpCode:%d,错误信息:%s", resp.StatusCode, http.StatusText(resp.StatusCode)))
		err = errors.New(fmt.Sprintf("知识库问答失败: httpCode:%d,错误信息:%s", resp.StatusCode, http.StatusText(resp.StatusCode)))
		return
	}

	// 解析流式响应
	contentChan, errChan, closeChan, closeLlmChan := eta_llm.ParseStreamResponse(resp)
	session.CloseLlmChan = &closeLlmChan
	// 标记llm提问状态：提问中
	session.LLMStatus = 1
	emptyContent := true
	// 处理流式数据并发送到 WebSocket
	for {
		select {
		case content, ok := <-contentChan:
			if !ok && session.LLMStatus != -1 {
				err = errors.New("未知的内容错误异常")

				// 标记llm提问状态：未提问
				session.LLMStatus = 0
				return
			}
			session.UpdateActivity()
			if emptyContent {
				emptyContent = false
			}
			// 发送消息到 WebSocket
			_ = session.Conn.WriteMessage(websocket.TextMessage, []byte(content))
		case chanErr, ok := <-errChan:
			if !ok && session.LLMStatus != -1 {
				err = errors.New("未知的错误异常")
			} else if chanErr != nil {
				err = errors.New(chanErr.Error())
			}
			// 标记llm提问状态：未提问
			session.LLMStatus = 0
			// 发送错误消息到 WebSocket
			return
		case <-closeChan:
			if emptyContent {
				_ = session.Conn.WriteMessage(websocket.TextMessage, []byte("<think></think>"))
				_ = session.Conn.WriteMessage(websocket.TextMessage, []byte("暂时找不到答案"))
			}
			_ = session.Conn.WriteMessage(websocket.TextMessage, []byte("<EOF/>"))
			// 标记llm提问状态：未提问
			session.LLMStatus = 0

			return
		}
	}
	// 更新最后活跃时间
	// 发送响应
	//return session.Conn.WriteMessage(websocket.TextMessage, []byte(response))
}

// AddSession Add 添加一个新的会话
func (manager *ConnectionManager) AddSession(session *Session) {
	manager.Sessions.Store(session.Id, session)
}
func (manager *ConnectionManager) GetSessionId(userId int, sessionId string) (sessionID string) {
	return fmt.Sprintf("%d_%s", userId, sessionId)
}

// RemoveSession Remove 移除一个会话
func (manager *ConnectionManager) RemoveSession(sessionCode string) {
	fmt.Printf("移除会话: SessionID=%s, UserID=%s", sessionCode, sessionCode)
	manager.Sessions.Delete(sessionCode)
}

// GetSession 获取一个会话
func (manager *ConnectionManager) GetSession(sessionCode string) (session *Session, exists bool) {
	if data, ok := manager.Sessions.Load(sessionCode); ok {
		session = data.(*Session)
		exists = ok
	}
	return
}

// CheckAll 批量检测所有连接
func (manager *ConnectionManager) CheckAll() {
	manager.Sessions.Range(func(key, value interface{}) bool {
		session := value.(*Session)
		// 判断超时
		if time.Since(session.LastActive) > 2*connectionTimeout {
			fmt.Printf("连接超时关闭: SessionID=%s, UserID=%s", session.Id, session.UserId)
			utils.FileLog.Warn("连接超时关闭: SessionID=%s, UserID=%s", session.Id, session.UserId)
			session.Close()
			return true
		}
		// 发送心跳
		go func(s *Session) {
			err := s.Conn.WriteControl(websocket.PingMessage,
				nil, time.Now().Add(writeWaitTimeout))
			if err != nil {
				fmt.Printf("心跳发送失败: SessionID=%s, Error=%v", s.Id, err)
				utils.FileLog.Warn("心跳发送失败: SessionID=%s, Error=%v",
					s.Id, err)
				fmt.Println("心跳无响应，退出请求")
				session.Close()
			}
		}(session)
		return true
	})
}

// Start 启动心跳检测
func (manager *ConnectionManager) Start() {
	defer manager.ticker.Stop()
	for {
		select {
		case <-manager.ticker.C:
			manager.CheckAll()
		case <-manager.stopChan:
			return
		}
	}
}

// Stop 停止心跳检测
func (manager *ConnectionManager) Stop() {
	close(manager.stopChan)
}