添加聊天角色支持

server/chat_handler.go

@@ -24,6 +24,7 @@ func (s *Server) ChatHandle(c *gin.Context) {
 		return
 	}
 	token := c.Query("token")
+	role := c.Query("role")
 	logger.Infof("New websocket connected, IP: %s", c.Request.RemoteAddr)
 	client := NewWsClient(ws)
 	go func() {
@@ -37,7 +38,7 @@ func (s *Server) ChatHandle(c *gin.Context) {
 
 			logger.Info(string(message))
 			// TODO: 当前只保持当前会话的上下文，部保存用户的所有的聊天历史记录，后期要考虑保存所有的历史记录
-			err = s.sendMessage(token, string(message), client)
+			err = s.sendMessage(token, role, string(message), client)
 			if err != nil {
 				logger.Error(err)
 			}
@@ -46,7 +47,7 @@ func (s *Server) ChatHandle(c *gin.Context) {
 }
 
 // 将消息发送给 ChatGPT 并获取结果，通过 WebSocket 推送到客户端
-func (s *Server) sendMessage(userId string, text string, ws Client) error {
+func (s *Server) sendMessage(sessionId string, role string, text string, ws Client) error {
 	var r = types.ApiRequest{
 		Model:       s.Config.Chat.Model,
 		Temperature: s.Config.Chat.Temperature,
@@ -54,11 +55,13 @@ func (s *Server) sendMessage(userId string, text string, ws Client) error {
 		Stream:      true,
 	}
 	var context []types.Message
-	if v, ok := s.ChatContext[userId]; ok && s.Config.Chat.EnableContext {
+	var key = sessionId + role
+	if v, ok := s.ChatContext[key]; ok && s.Config.Chat.EnableContext {
 		context = v
 	} else {
-		context = make([]types.Message, 0)
+		context = s.Config.ChatRoles[role].Context
 	}
+	logger.Info(context)
 	r.Messages = append(context, types.Message{
 		Role:    "user",
 		Content: text,
@@ -166,7 +169,8 @@ func (s *Server) sendMessage(userId string, text string, ws Client) error {
 	})
 	message.Content = strings.Join(contents, "")
 	context = append(context, message)
-	s.ChatContext[userId] = context
+	// 保存上下文
+	s.ChatContext[key] = context
 	return nil
 }
 

server/config_handler.go

@@ -151,5 +151,15 @@ func (s *Server) ListApiKeys(c *gin.Context) {
 }
 
 func (s *Server) GetChatRoles(c *gin.Context) {
-	c.JSON(http.StatusOK, types.BizVo{Code: types.Success, Message: types.OkMsg, Data: s.Config.ChatRoles})
+	var roles = make(map[string]interface{})
+	for k, v := range s.Config.ChatRoles {
+		roles[k] = struct {
+			Key  string `json:"key"`
+			Name string `json:"name"`
+		}{
+			Key:  v.Key,
+			Name: v.Name,
+		}
+	}
+	c.JSON(http.StatusOK, types.BizVo{Code: types.Success, Message: types.OkMsg, Data: roles})
 }

server/server.go

@@ -73,8 +73,8 @@ func (s *Server) Run(webRoot embed.FS, path string, debug bool) {
 	engine.GET("/api/session/get", s.GetSessionHandle)
 	engine.POST("/api/login", s.LoginHandle)
 	engine.Any("/api/chat", s.ChatHandle)
-	engine.GET("/api/chat-roles/get", s.GetChatRoles)
 	engine.POST("/api/config/set", s.ConfigSetHandle)
+	engine.GET("/api/config/chat-roles/get", s.GetChatRoles)
 	engine.POST("api/config/token/add", s.AddToken)
 	engine.POST("api/config/token/remove", s.RemoveToken)
 	engine.POST("api/config/apikey/add", s.AddApiKey)
@@ -154,6 +154,7 @@ func AuthorizeMiddleware(s *Server) gin.HandlerFunc {
 	return func(c *gin.Context) {
 		if !s.Config.EnableAuth ||
 			c.Request.URL.Path == "/api/login" ||
+			c.Request.URL.Path == "/api/config/chat-roles/get" ||
 			!strings.HasPrefix(c.Request.URL.Path, "/api") {
 			c.Next()
 			return