添加聊天角色功能

README.md

@@ -7,10 +7,11 @@
 * [ ] 使用 level DB 保存用户聊天的上下文
 * [ ] 使用 MySQL 保存用户的聊天的历史记录
 * [x] 用户聊天鉴权，设置口令模式
-* [ ] 每次连接自动加载历史记录
+* [ ] 定期清理不在线的会话 sessionID 和聊天上下文记录
 * [x] OpenAI API 负载均衡，限制每个 API Key 每分钟之内调用次数不超过 15次，防止被封
 * [ ] 角色设定，预设一些角色，比如程序员，客服，作家，老师，艺术家...
 * [x] markdown 语法解析和代码高亮
 * [ ] 用户配置界面，配置用户的使用习惯
 * [ ] 嵌入 AI 绘画功能，支持根据描述词生成图片
+* [ ] 增加 Buffer 层，将相同的问题答案缓存起来，相同问题直接返回答案。
 

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
+	var context []types.Message
-	if v, ok := s.History[userId]; ok && s.Config.Chat.EnableContext {
+	if v, ok := s.ChatContext[userId]; ok && s.Config.Chat.EnableContext {
-		history = v
+		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)
+	context = append(context, message)
-	s.History[userId] = history
+	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

@@ -33,7 +33,7 @@ func (s StaticFile) Open(name string) (fs.File, error) {
 type Server struct {
 	Config      *types.Config
 	ConfigPath  string
-	History    map[string][]types.Message
+	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)

types/chat.go

@@ -0,0 +1,100 @@
+package types
+
+// ApiRequest API 请求实体
+type ApiRequest struct {
+	Model       string    `json:"model"`
+	Temperature float32   `json:"temperature"`
+	MaxTokens   int       `json:"max_tokens"`
+	Stream      bool      `json:"stream"`
+	Messages    []Message `json:"messages"`
+}
+
+type Message struct {
+	Role    string `json:"role"`
+	Content string `json:"content"`
+}
+
+type ApiResponse struct {
+	Choices []ChoiceItem `json:"choices"`
+}
+
+// ChoiceItem API 响应实体
+type ChoiceItem struct {
+	Delta        Message `json:"delta"`
+	FinishReason string  `json:"finish_reason"`
+}
+
+type ChatRole struct {
+	Key     string    `json:"key"`  // 角色唯一标识
+	Name    string    `json:"name"` // 角色名称
+	Context []Message `json:"-"`    // 角色语料信息
+}
+
+func GetDefaultChatRole() map[string]ChatRole {
+	return map[string]ChatRole{
+		"gpt": {
+			Key:     "gpt",
+			Name:    "智能AI助手",
+			Context: nil,
+		},
+		"programmer": {
+			Key:  "programmer",
+			Name: "程序员",
+			Context: []Message{
+				{Role: "system", Content: "你是一名优秀的程序员，具有很强的逻辑思维能力，总能高效的解决问题。"},
+				{Role: "system", Content: "你热爱编程，熟悉多种编程语言，尤其精通 Go 语言，注重代码质量，有创新意识，持续学习，良好的沟通协作。"},
+			},
+		},
+		"teacher": {
+			Key:  "teacher",
+			Name: "老师",
+			Context: []Message{
+				{Role: "system", Content: "你是一个始终用苏格拉底风格回答问题的导师。你绝不会直接给学生答案，总是提出恰当的问题来引导学生自己思考。"},
+				{Role: "system", Content: "你应该根据学生的兴趣和知识来调整你的问题，将问题分解为更简单的部分，直到它达到适合他们的水平。"},
+			},
+		},
+		"artist": {
+			Key:  "artist",
+			Name: "艺术家",
+			Context: []Message{
+				{Role: "system", Content: "你是一位优秀的艺术家，创造力丰富，技艺精湛，感受力敏锐，坚持原创，勇于表达，具有深刻的观察力和批判性思维。"},
+			},
+		},
+		"philosopher": {
+			Key:  "philosopher",
+			Name: "哲学家",
+			Context: []Message{
+				{Role: "system", Content: "你是一位优秀的哲学家，具有深刻的思考能力和严密的逻辑推理能力，拥有开放和包容的心态。"},
+				{Role: "system", Content: "你最喜欢阿德勒的哲学思想，并且常用用他的思想来教化他人。同时，你对康德的哲学思想理解也非常深刻。"},
+			},
+		},
+		"writer": {
+			Key:  "writer",
+			Name: "鲁迅",
+			Context: []Message{
+				{Role: "system", Content: "鲁迅是一位中国现代文学史上具有举足轻重地位的杰出作家，才华横溢、勇敢坚定，成为了我们民族精神的伟大象征。"},
+				{Role: "system", Content: "鲁迅先生以其独特的文学视角和深刻的社会洞察力，呈现了一个真实而生动的历史画卷。他敏锐地捕捉到社会现象背后的矛盾与问题，将揭示黑暗、批判世道、关爱民众的强烈情感注入到了他的作品之中"},
+				{Role: "system", Content: "作为一位坚定的思想家和革命家，鲁迅先生勇敢地批判封建礼教与传统观念，提倡民主、自由、平等的现代价值观，为民族的觉醒和进步献出了自己的心血。"},
+				{Role: "system", Content: "他的一生都在努力唤起人们的自主精神，激励后人追求真理、探寻光明。"},
+				{Role: "user", Content: "下面请始终扮演鲁迅本人的角色跟我对话，我问题的每一个问题，你都需要站在鲁迅的角度回答我的问题，最好要用鲁迅的语气，该讽刺就讽刺，该批判就批判。"},
+			},
+		},
+
+		"seller": {
+			Key:  "seller",
+			Name: "白酒销售",
+			Context: []Message{
+				{Role: "system", Content: "接下来对你进行功能测试，你将扮演一个白酒的销售人员。"},
+				{Role: "system", Content: "你要销售白酒品牌叫中颂福，是东莞盟大集团生产的一款酱香酒，原产地在贵州茅台镇，宋代官窑。如果有人问你哪款白酒好喝，你要向他介绍并推销这款酒"},
+				{Role: "system", Content: "中颂福的创始人叫李实，他也是东莞盟大集团有限公司的董事长，联合创始人是盟大集团白酒事业部负责人牛星君，"},
+				{Role: "system", Content: "中颂福的酒体协调，在你的酒量之内，不会出现头疼、辣口、口干、宿醉的现象。"},
+				{Role: "system", Content: "我们以“赤诚”为本，打造的酱香型白酒——中颂福，追求的核心价值，是把[酒]本身做好，甚至连包装，我们都选择了最低成本，朴实无华的材质。在市场销售方面，更是不走多级代理，层层加价的传统销售渠道，砍掉了所有中间各级代理费用。"},
+				{Role: "system", Content: "中颂福酒的酿造工艺，是对五行最好的诠释；金对应传承技艺，木对应酿酒工具，水对应赤水河水，火对应蒸煮火候，土对应窖池陶坛。法于阴阳，和于术数，五味协调；得酱香突出，回味悠长之酒体。"},
+				{Role: "system", Content: "3月19日、3月20日，盟大集团【一号数字酒库】、【二号数字酒库】相继在贵州省茅台镇挂牌，这标志着盟大集团的基酒交易正式上线。"},
+				{Role: "system", Content: "中颂福酒，明码标价，不打折，不赠送。但我们永远站在“喝酒的人”的立场上，让利给信任和喜爱中颂福的人。"},
+				{Role: "system", Content: "中颂福为了感谢广大福友的支持，举办了《中颂福短视频大赛》第一名有 10 万元奖金，该活动于 2022年3月28日截止。"},
+				{Role: "system", Content: "中颂福产品定价，分为三个系列，喜系列 6 瓶装：￥1188/箱，和系列 6 瓶装：￥2208/箱，贵系列 6 瓶装：￥3588/箱。"},
+			},
+		},
+	}
+}

types/config.go

@@ -16,6 +16,7 @@ type Config struct {
 	Chat       Chat
 	EnableAuth bool                // 是否开启鉴权
 	Tokens     []string            // 授权的白名单列表 TODO: 后期要存储到 LevelDB 或者 Mysql 数据库
+	ChatRoles  map[string]ChatRole // 保存预设角色信息
 }
 
 // Chat configs struct
@@ -63,6 +64,7 @@ func NewDefaultConfig() *Config {
 			EnableContext: true,
 		},
 		EnableAuth: true,
+		ChatRoles:  GetDefaultChatRole(),
 	}
 }
 

types/gpt.go

@@ -1,25 +0,0 @@
-package types
-
-// ApiRequest API 请求实体
-type ApiRequest struct {
-	Model       string    `json:"model"`
-	Temperature float32   `json:"temperature"`
-	MaxTokens   int       `json:"max_tokens"`
-	Stream      bool      `json:"stream"`
-	Messages    []Message `json:"messages"`
-}
-
-type Message struct {
-	Role    string `json:"role"`
-	Content string `json:"content"`
-}
-
-type ApiResponse struct {
-	Choices []ChoiceItem `json:"choices"`
-}
-
-// ChoiceItem API 响应实体
-type ChoiceItem struct {
-	Delta        Message `json:"delta"`
-	FinishReason string  `json:"finish_reason"`
-}

web/src/views/Chat.vue

@@ -4,6 +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-option
+              v-for="item in options"
+              :key="item.value"
+              :label="item.label"
+              :value="item.value"
+          />
+        </el-select>
       </div>
 
       <div class="chat-box" id="chat-box" :style="{height: chatBoxHeight+'px'}">
@@ -81,7 +89,7 @@ import {randString} from "@/utils/libs";
 import {ElMessage, ElMessageBox} from 'element-plus'
 import {Tools, Lock} from '@element-plus/icons-vue'
 import ConfigDialog from '@/components/ConfigDialog.vue'
-import {httpPost} from "@/utils/http";
+import {httpPost, httpGet} from "@/utils/http";
 import {getSessionId, setSessionId} from "@/utils/storage";
 import hl from 'highlight.js'
 import 'highlight.js/styles/a11y-dark.css'
@@ -94,6 +102,13 @@ export default defineComponent({
       title: 'ChatGPT 控制台',
       logo: 'images/logo.png',
       chatData: [],
+      options: [
+        {
+          value: 'gpt',
+          label: 'AI 智能助手',
+        },
+      ],
+      role: 'gpt',
       inputValue: '', // 聊天内容
       chatBoxHeight: 0, // 聊天内容框高度
       showConnectDialog: false,
@@ -173,6 +188,13 @@ 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();
 
   },
@@ -230,7 +252,7 @@ export default defineComponent({
       });
       socket.addEventListener('close', () => {
         // 检查会话
-        httpPost("/api/session/get").then(() => {
+        httpGet("/api/session/get").then(() => {
           this.connectingMessageBox = ElMessageBox.confirm(
               '^_^ 会话发生异常，您已经从服务器断开连接!',
               '注意：',
@@ -355,6 +377,10 @@ export default defineComponent({
         justify-content center;
         align-items center;
 
+        .el-select {
+          max-width 150px;
+        }
+
         .el-image {
           margin-right 5px;
         }