定期清理用户聊天会话上下文 ChatContext

server/chat_handler.go

@@ -80,9 +80,9 @@ func (s *Server) sendMessage(session types.ChatSession, role types.ChatRole, pro
 		Stream:      true,
 	}
 	var context []types.Message
-	var key = session.SessionId + role.Name
-	if v, ok := s.ChatContext[key]; ok && s.Config.Chat.EnableContext {
-		context = v
+	var ctxKey = fmt.Sprintf("%s-%s", session.SessionId, role.Key)
+	if v, ok := s.ChatContexts[ctxKey]; ok && s.Config.Chat.EnableContext {
+		context = v.Messages
 	} else {
 		context = role.Context
 	}
@@ -206,7 +206,11 @@ func (s *Server) sendMessage(session types.ChatSession, role types.ChatRole, pro
 	context = append(context, useMsg)
 	message.Content = strings.Join(contents, "")
 	context = append(context, message)
-	s.ChatContext[key] = context
+	// 更新上下文消息
+	s.ChatContexts[ctxKey] = types.ChatContext{
+		Messages:       context,
+		LastAccessTime: time.Now().Unix(),
+	}
 
 	// 追加历史消息
 	if user.EnableHistory {

server/config_handler.go

@@ -63,6 +63,18 @@ func (s *Server) ConfigSetHandle(c *gin.Context) {
 		s.Config.Chat.EnableContext = v
 	}
 
+	if expireTime, ok := data["chat_context_expire_time"]; ok {
+		v, err := strconv.Atoi(expireTime)
+		if err != nil {
+			c.JSON(http.StatusOK, types.BizVo{
+				Code:    types.InvalidParams,
+				Message: "chat_context_expire_time must be a integer parameter",
+			})
+			return
+		}
+		s.Config.Chat.ChatContextExpireTime = v
+	}
+
 	// enable auth
 	if enableAuth, ok := data["enable_auth"]; ok {
 		v, err := strconv.ParseBool(enableAuth)
@@ -123,7 +135,7 @@ func (s *Server) AddUserHandle(c *gin.Context) {
 		return
 	}
 
-	user := types.User{Name: data.Name, MaxCalls: data.MaxCalls, RemainingCalls: data.MaxCalls}
+	user := types.User{Name: data.Name, MaxCalls: data.MaxCalls, RemainingCalls: data.MaxCalls, EnableHistory: data.EnableHistory}
 	err = PutUser(user)
 	if err != nil {
 		c.JSON(http.StatusOK, types.BizVo{Code: types.Failed, Message: "Failed to save configs"})

server/server.go

@@ -15,6 +15,7 @@ import (
 	"path/filepath"
 	"runtime/debug"
 	"strings"
+	"time"
 )
 
 var logger = logger2.GetLogger()
@@ -31,9 +32,9 @@ func (s StaticFile) Open(name string) (fs.File, error) {
 }
 
 type Server struct {
-	Config      *types.Config
-	ConfigPath  string
-	ChatContext map[string][]types.Message // 聊天上下文 [SessionID] => []Messages
+	Config       *types.Config
+	ConfigPath   string
+	ChatContexts map[string]types.ChatContext // 聊天上下文 [SessionID+ChatRole] => ChatContext
 
 	// 保存 Websocket 会话 Username, 每个 Username 只能连接一次
 	// 防止第三方直接连接 socket 调用 OpenAI API
@@ -61,7 +62,7 @@ func NewServer(configPath string) (*Server, error) {
 	return &Server{
 		Config:           config,
 		ConfigPath:       configPath,
-		ChatContext:      make(map[string][]types.Message, 16),
+		ChatContexts:     make(map[string]types.ChatContext, 16),
 		ChatSession:      make(map[string]types.ChatSession),
 		ApiKeyAccessStat: make(map[string]int64),
 	}, nil
@@ -111,6 +112,20 @@ func (s *Server) Run(webRoot embed.FS, path string, debug bool) {
 		path:    path,
 	}))
 
+	// 定时清理过期的会话
+	go func() {
+		for {
+			for key, context := range s.ChatContexts {
+				// 清理超过 60min 没有更新，则表示为过期会话
+				if time.Now().Unix()-context.LastAccessTime > 3600 {
+					logger.Infof("清理会话上下文: %s", key)
+					delete(s.ChatContexts, key)
+				}
+			}
+			time.Sleep(time.Second * 5) // 每隔 5 秒钟清理一次
+		}
+	}()
+
 	logger.Infof("http://%s", s.Config.Listen)
 	err := engine.Run(s.Config.Listen)
 

types/chat.go

@@ -48,6 +48,12 @@ type ChatSession struct {
 	Username  string `json:"user"`      // 当前登录的 user
 }
 
+// ChatContext 聊天上下文
+type ChatContext struct {
+	Messages       []Message
+	LastAccessTime int64 // 最后一次访问上下文时间
+}
+
 func GetDefaultChatRole() map[string]ChatRole {
 	return map[string]ChatRole{
 		"gpt": {

types/config.go

@@ -22,12 +22,13 @@ type User struct {
 
 // Chat configs struct
 type Chat struct {
-	ApiURL        string
-	ApiKeys       []string
-	Model         string
-	Temperature   float32
-	MaxTokens     int
-	EnableContext bool // 是否保持聊天上下文
+	ApiURL                string
+	ApiKeys               []string
+	Model                 string
+	Temperature           float32
+	MaxTokens             int
+	EnableContext         bool // 是否保持聊天上下文
+	ChatContextExpireTime int  // 聊天上下文过期时间，单位：秒
 }
 
 // Session configs struct

utils/config.go

@@ -27,12 +27,13 @@ func NewDefaultConfig() *types.Config {
 			SameSite:  http.SameSiteLaxMode,
 		},
 		Chat: types.Chat{
-			ApiURL:        "https://api.openai.com/v1/chat/completions",
-			ApiKeys:       []string{""},
-			Model:         "gpt-3.5-turbo",
-			MaxTokens:     1024,
-			Temperature:   0.9,
-			EnableContext: true,
+			ApiURL:                "https://api.openai.com/v1/chat/completions",
+			ApiKeys:               []string{""},
+			Model:                 "gpt-3.5-turbo",
+			MaxTokens:             1024,
+			Temperature:           0.9,
+			EnableContext:         true,
+			ChatContextExpireTime: 3600,
 		},
 	}
 }