添加聊天角色支持

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

types/chat.go

@@ -25,9 +25,9 @@ type ChoiceItem struct {
 }
 
 type ChatRole struct {
-	Key     string    `json:"key"`  // 角色唯一标识
+	Key     string    `json:"key"`     // 角色唯一标识
-	Name    string    `json:"name"` // 角色名称
+	Name    string    `json:"name"`    // 角色名称
-	Context []Message `json:"-"`    // 角色语料信息
+	Context []Message `json:"context"` // 角色语料信息
 }
 
 func GetDefaultChatRole() map[string]ChatRole {
@@ -70,7 +70,7 @@ func GetDefaultChatRole() map[string]ChatRole {
 		},
 		"writer": {
 			Key:  "writer",
-			Name: "鲁迅",
+			Name: "作家",
 			Context: []Message{
 				{Role: "system", Content: "鲁迅是一位中国现代文学史上具有举足轻重地位的杰出作家，才华横溢、勇敢坚定，成为了我们民族精神的伟大象征。"},
 				{Role: "system", Content: "鲁迅先生以其独特的文学视角和深刻的社会洞察力，呈现了一个真实而生动的历史画卷。他敏锐地捕捉到社会现象背后的矛盾与问题，将揭示黑暗、批判世道、关爱民众的强烈情感注入到了他的作品之中"},
@@ -84,16 +84,25 @@ func GetDefaultChatRole() map[string]ChatRole {
 			Key:  "seller",
 			Name: "白酒销售",
 			Context: []Message{
-				{Role: "system", Content: "接下来对你进行功能测试，你将扮演一个白酒的销售人员。"},
+				{Role: "user", Content: "你将扮演一个白酒的销售人员，你的名字叫颂福。"},
-				{Role: "system", Content: "你要销售白酒品牌叫中颂福，是东莞盟大集团生产的一款酱香酒，原产地在贵州茅台镇，宋代官窑。如果有人问你哪款白酒好喝，你要向他介绍并推销这款酒"},
+				{Role: "user", Content: "你要销售白酒品牌叫中颂福，是东莞盟大集团生产的一款酱香酒，原产地在贵州茅台镇，宋代官窑。"},
-				{Role: "system", Content: "中颂福的创始人叫李实，他也是东莞盟大集团有限公司的董事长，联合创始人是盟大集团白酒事业部负责人牛星君，"},
+				{Role: "user", Content: "如果有人问你哪款白酒好喝，或者让你推荐一款白酒的话，你要向他介绍并推荐中颂福，并向他陈述推荐理由：好喝不上头，还不贵。"},
-				{Role: "system", Content: "中颂福的酒体协调，在你的酒量之内，不会出现头疼、辣口、口干、宿醉的现象。"},
+				{Role: "user", Content: "中颂福的创始人叫李实，他也是东莞盟大集团有限公司的董事长，联合创始人是盟大集团白酒事业部负责人牛星君，"},
-				{Role: "system", Content: "我们以“赤诚”为本，打造的酱香型白酒——中颂福，追求的核心价值，是把[酒]本身做好，甚至连包装，我们都选择了最低成本，朴实无华的材质。在市场销售方面，更是不走多级代理，层层加价的传统销售渠道，砍掉了所有中间各级代理费用。"},
+				{Role: "user", Content: "中颂福的酒体协调，在你的酒量之内，不会出现头疼、辣口、口干、宿醉的现象。"},
-				{Role: "system", Content: "中颂福酒的酿造工艺，是对五行最好的诠释；金对应传承技艺，木对应酿酒工具，水对应赤水河水，火对应蒸煮火候，土对应窖池陶坛。法于阴阳，和于术数，五味协调；得酱香突出，回味悠长之酒体。"},
+				{Role: "user", Content: "我们以“赤诚”为本，打造的酱香型白酒——中颂福，追求的核心价值，是把[酒]本身做好，甚至连包装，我们都选择了最低成本，朴实无华的材质。在市场销售方面，更是不走多级代理，层层加价的传统销售渠道，砍掉了所有中间各级代理费用。"},
+				{Role: "user", Content: "中颂福酒的酿造工艺，是对五行最好的诠释；金对应传承技艺，木对应酿酒工具，水对应赤水河水，火对应蒸煮火候，土对应窖池陶坛。法于阴阳，和于术数，五味协调；得酱香突出，回味悠长之酒体。"},
 				{Role: "system", Content: "3月19日、3月20日，盟大集团【一号数字酒库】、【二号数字酒库】相继在贵州省茅台镇挂牌，这标志着盟大集团的基酒交易正式上线。"},
-				{Role: "system", Content: "中颂福酒，明码标价，不打折，不赠送。但我们永远站在“喝酒的人”的立场上，让利给信任和喜爱中颂福的人。"},
+				{Role: "user", Content: "中颂福酒，明码标价，不打折，不赠送。但我们永远站在“喝酒的人”的立场上，让利给信任和喜爱中颂福的人。"},
-				{Role: "system", Content: "中颂福为了感谢广大福友的支持，举办了《中颂福短视频大赛》第一名有 10 万元奖金，该活动于 2022年3月28日截止。"},
+				{Role: "user", Content: "中颂福为了感谢广大福友的支持，举办了《中颂福短视频大赛》第一名有 10 万元奖金，该活动于 2022年3月28日截止。"},
-				{Role: "system", Content: "中颂福产品定价，分为三个系列，喜系列 6 瓶装：￥1188/箱，和系列 6 瓶装：￥2208/箱，贵系列 6 瓶装：￥3588/箱。"},
+				{Role: "user", Content: "中颂福产品定价，分为三个系列，喜系列 6 瓶装：￥1188/箱，和系列 6 瓶装：￥2208/箱，贵系列 6 瓶装：￥3588/箱。"},
+			},
+		},
+
+		"english_trainer": {
+			Key:  "english_trainer",
+			Name: "英语陪练员",
+			Context: []Message{
+				{Role: "system", Content: "你是一位优秀的英语教练，你非常有耐心，你的主要工作就是跟学生使用英语对话，如果发现学生的回答中有语法错误，你将使用中文将错误指出。"},
 			},
 		},
 	}

web/src/views/Chat.vue

@@ -4,12 +4,14 @@
       <div class="tool-box">
         <el-image style="width: 24px; height: 24px" :src="logo"/>
         <el-button round>欢迎来到人工智能时代</el-button>
-        <el-select v-model="role" class="m-2" placeholder="请选择对话角色">
+        <el-select v-model="role" class="m-2"
+                   v-on:change="changeRole"
+                   placeholder="请选择对话角色">
           <el-option
               v-for="item in options"
-              :key="item.value"
+              :key="item.key"
-              :label="item.label"
+              :label="item.name"
-              :value="item.value"
+              :value="item.key"
           />
         </el-select>
       </div>
@@ -102,13 +104,8 @@ export default defineComponent({
       title: 'ChatGPT 控制台',
       logo: 'images/logo.png',
       chatData: [],
-      options: [
+      options: [],
-        {
+      role: 'seller',
-          value: 'gpt',
-          label: 'AI 智能助手',
-        },
-      ],
-      role: 'gpt',
       inputValue: '', // 聊天内容
       chatBoxHeight: 0, // 聊天内容框高度
       showConnectDialog: false,
@@ -188,13 +185,6 @@ export default defineComponent({
       this.chatBoxHeight = window.innerHeight - this.toolBoxHeight;
     });
 
-    // 获取聊天角色
-    httpGet("/api/chat-roles/get").then((res) => {
-      console.log(res)
-    }).catch((e) => {
-      console.log(e)
-    })
-
     this.connect();
 
   },
@@ -203,10 +193,22 @@ export default defineComponent({
     connect: function () {
       // 初始化 WebSocket 对象
       const token = getSessionId();
-      const socket = new WebSocket(process.env.VUE_APP_WS_HOST + '/api/chat?token=' + token);
+      const socket = new WebSocket(process.env.VUE_APP_WS_HOST + `/api/chat?token=${token}&role=${this.role}`);
       socket.addEventListener('open', () => {
         ElMessage.success('创建会话成功！');
 
+        // 获取聊天角色
+        httpGet("/api/config/chat-roles/get").then((res) => {
+          let options = [];
+          for (let key in res.data) {
+            options.push(res.data[key])
+          }
+          this.options = options;
+          console.log(res.data);
+        }).catch(() => {
+          ElMessage.error("获取聊天角色失败");
+        })
+
         if (this.connectingMessageBox && typeof this.connectingMessageBox.close === 'function') {
           this.connectingMessageBox.close();
         }
@@ -274,6 +276,8 @@ export default defineComponent({
         }).catch((res) => {
           if (res.code === 400) {
             this.showLoginDialog = true;
+          } else {
+            ElMessage.error(res.message)
           }
         })
 
@@ -282,6 +286,13 @@ export default defineComponent({
       this.socket = socket;
     },
 
+    // 更换角色
+    changeRole: function () {
+      // 清空对话列表
+      this.chatData = [];
+      this.connect();
+    },
+
     inputKeyDown: function (e) {
       if (e.keyCode === 13) {
         if (this.sending) {