From daa36f308b9d3b62a44618ef3ebdd4b5d90e89e8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=9D=8E=E5=AF=BB=E6=AC=A2?= <lxh@cxh.cn>
Date: Tue, 16 Apr 2024 17:19:08 +0800
Subject: [PATCH] =?UTF-8?q?:art:=20=E4=BC=98=E5=8C=96AI=E5=AF=B9=E8=AF=9D?=
 =?UTF-8?q?=E8=A7=92=E8=89=B2=E8=AE=BE=E7=BD=AE=EF=BC=8C=E6=94=AF=E6=8C=81?=
 =?UTF-8?q?=E4=B8=BA=E7=BE=A4=E6=88=96=E8=80=85=E5=A5=BD=E5=8F=8B=E5=8D=95?=
 =?UTF-8?q?=E7=8B=AC=E8=AE=BE=E7=BD=AE=E8=A7=92=E8=89=B2(=E9=9C=80?=
 =?UTF-8?q?=E6=89=8B=E5=8A=A8=E4=BF=AE=E6=94=B9=E6=95=B0=E6=8D=AE=E5=BA=93?=
 =?UTF-8?q?)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 entity/friend.go     | 1 +
 plugin/plugins/ai.go | 8 +++++++-
 2 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/entity/friend.go b/entity/friend.go
index 0cb57be1..04539615 100644
--- a/entity/friend.go
+++ b/entity/friend.go
@@ -15,6 +15,7 @@ type Friend struct {
 	LastActive     time.Time `json:"lastActive"`                                               // 最后活跃时间
 	EnableAi       bool      `json:"enableAI" gorm:"type:tinyint(1) default 0 not null"`       // 是否使用AI
 	AiModel        string    `json:"aiModel"`                                                  // AI模型
+	Prompt         string    `json:"prompt"`                                                   // 提示词
 	EnableChatRank bool      `json:"enableChatRank" gorm:"type:tinyint(1) default 0 not null"` // 是否使用聊天排行
 	EnableWelcome  bool      `json:"enableWelcome" gorm:"type:tinyint(1) default 0 not null"`  // 是否启用迎新
 	EnableSummary  bool      `json:"enableSummary" gorm:"type:tinyint(1) default 0 not null"`  // 是否启用总结
diff --git a/plugin/plugins/ai.go b/plugin/plugins/ai.go
index 2e6bfe01..dec8b5e3 100644
--- a/plugin/plugins/ai.go
+++ b/plugin/plugins/ai.go
@@ -46,13 +46,19 @@ func AI(m *plugin.MessageContext) {
 		m.Content = strings.Replace(m.Content, matches[0], "", 1)
 	}
 
+	// 处理预设角色，默认是配置文件里的，如果数据库配置不为空，则使用数据库配置
+	prompt := config.Conf.Ai.Personality
+	if friendInfo.Prompt != "" {
+		prompt = friendInfo.Prompt
+	}
+
 	// 组装消息体
 	messages := make([]openai.ChatCompletionMessage, 0)
 	if config.Conf.Ai.Personality != "" {
 		// 填充人设
 		messages = append(messages, openai.ChatCompletionMessage{
 			Role:    openai.ChatMessageRoleSystem,
-			Content: config.Conf.Ai.Personality,
+			Content: prompt,
 		})
 	}