添加聊天角色功能

server/chat_handler.go

@@ -53,13 +53,13 @@ func (s *Server) sendMessage(userId string, text string, ws Client) error {
 		MaxTokens:   s.Config.Chat.MaxTokens,
 		Stream:      true,
 	}
-	var history []types.Message
-	if v, ok := s.History[userId]; ok && s.Config.Chat.EnableContext {
-		history = v
+	var context []types.Message
+	if v, ok := s.ChatContext[userId]; ok && s.Config.Chat.EnableContext {
+		context = v
 	} else {
-		history = make([]types.Message, 0)
+		context = make([]types.Message, 0)
 	}
-	r.Messages = append(history, types.Message{
+	r.Messages = append(context, types.Message{
 		Role:    "user",
 		Content: text,
 	})
@@ -160,13 +160,13 @@ func (s *Server) sendMessage(userId string, text string, ws Client) error {
 	}
 
 	// 追加历史消息
-	history = append(history, types.Message{
+	context = append(context, types.Message{
 		Role:    "user",
 		Content: text,
 	})
 	message.Content = strings.Join(contents, "")
-	history = append(history, message)
-	s.History[userId] = history
+	context = append(context, message)
+	s.ChatContext[userId] = context
 	return nil
 }
 

server/config_handler.go

@@ -149,3 +149,7 @@ func (s *Server) AddApiKey(c *gin.Context) {
 func (s *Server) ListApiKeys(c *gin.Context) {
 	c.JSON(http.StatusOK, types.BizVo{Code: types.Success, Message: types.OkMsg, Data: s.Config.Chat.ApiKeys})
 }
+
+func (s *Server) GetChatRoles(c *gin.Context) {
+	c.JSON(http.StatusOK, types.BizVo{Code: types.Success, Message: types.OkMsg, Data: s.Config.ChatRoles})
+}

server/server.go

@@ -31,9 +31,9 @@ func (s StaticFile) Open(name string) (fs.File, error) {
 }
 
 type Server struct {
-	Config     *types.Config
-	ConfigPath string
-	History    map[string][]types.Message
+	Config      *types.Config
+	ConfigPath  string
+	ChatContext map[string][]types.Message // 聊天上下文 [SessionID] => []Messages
 
 	// 保存 Websocket 会话 Token, 每个 Token 只能连接一次
 	// 防止第三方直接连接 socket 调用 OpenAI API
@@ -44,6 +44,9 @@ type Server struct {
 func NewServer(configPath string) (*Server, error) {
 	// load service configs
 	config, err := types.LoadConfig(configPath)
+	if config.ChatRoles == nil {
+		config.ChatRoles = types.GetDefaultChatRole()
+	}
 	if err != nil {
 		return nil, err
 	}
@@ -51,7 +54,7 @@ func NewServer(configPath string) (*Server, error) {
 	return &Server{
 		Config:           config,
 		ConfigPath:       configPath,
-		History:          make(map[string][]types.Message, 16),
+		ChatContext:      make(map[string][]types.Message, 16),
 		WsSession:        make(map[string]string),
 		ApiKeyAccessStat: make(map[string]int64),
 	}, nil
@@ -67,9 +70,10 @@ func (s *Server) Run(webRoot embed.FS, path string, debug bool) {
 	engine.Use(AuthorizeMiddleware(s))
 
 	engine.GET("/hello", Hello)
-	engine.POST("/api/session/get", s.GetSessionHandle)
+	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.POST("api/config/token/add", s.AddToken)
 	engine.POST("api/config/token/remove", s.RemoveToken)