BJ_ZTEITS-Developers / parking-cloud-yudao

Commit cddaca586360a3cb924190c9ade7c5a55a1a420c

Authored by YunaiV
1 parent 3bb57edc

【功能新增】AI：聊天时，增加知识库的拼接

Inline Side-by-side
Showing 5 changed files with 110 additions and 85 deletions
yudao-module-ai/yudao-module-ai-api/src/main/java/cn/iocoder/yudao/module/ai/enums/AiChatRoleEnum.java
  View file @cddaca5
... ... @@ -35,11 +35,7 @@ public enum AiChatRoleEnum {
35 35 ### 微信
36 36 除此之外不要任何解释性语句。
37 37 """),
38   -
39   - AI_KNOWLEDGE_ROLE("知识库助手", """
40   - 给你提供一些数据参考：{info},请回答我的问题。
41   - 请你跟进数据参考与工具返回结果回复用户的请求。
42   - """);
  38 + ;
43 39  
44 40 /**
45 41 * 角色名
... ...
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/dal/dataobject/chat/AiChatConversationDO.java
  View file @cddaca5
1 1 package cn.iocoder.yudao.module.ai.dal.dataobject.chat;
2 2  
3 3 import cn.iocoder.yudao.framework.mybatis.core.dataobject.BaseDO;
4   -import cn.iocoder.yudao.module.ai.dal.dataobject.knowledge.AiKnowledgeDO;
5   -import cn.iocoder.yudao.module.ai.dal.dataobject.model.AiModelDO;
6 4 import cn.iocoder.yudao.module.ai.dal.dataobject.model.AiChatRoleDO;
  5 +import cn.iocoder.yudao.module.ai.dal.dataobject.model.AiModelDO;
7 6 import com.baomidou.mybatisplus.annotation.KeySequence;
8 7 import com.baomidou.mybatisplus.annotation.TableId;
9 8 import com.baomidou.mybatisplus.annotation.TableName;
... ... @@ -65,14 +64,6 @@ public class AiChatConversationDO extends BaseDO {
65 64 */
66 65 private Long roleId;
67 66  
68   - // TODO @芋艿：可优化，绑定多个知识库。前提，spring ai 支持 RerankModel 的封装
69   - /**
70   - * 知识库编号
71   - * <p>
72   - * 关联 {@link AiKnowledgeDO#getId()}
73   - */
74   - private Long knowledgeId;
75   -
76 67 /**
77 68 * 模型编号
78 69 *
... ...
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/dal/dataobject/chat/AiChatMessageDO.java
  View file @cddaca5
1 1 package cn.iocoder.yudao.module.ai.dal.dataobject.chat;
2 2  
3 3 import cn.iocoder.yudao.framework.mybatis.core.dataobject.BaseDO;
  4 +import cn.iocoder.yudao.framework.mybatis.core.type.LongListTypeHandler;
4 5 import cn.iocoder.yudao.module.ai.dal.dataobject.knowledge.AiKnowledgeSegmentDO;
5   -import cn.iocoder.yudao.module.ai.dal.dataobject.model.AiModelDO;
6 6 import cn.iocoder.yudao.module.ai.dal.dataobject.model.AiChatRoleDO;
  7 +import cn.iocoder.yudao.module.ai.dal.dataobject.model.AiModelDO;
7 8 import com.baomidou.mybatisplus.annotation.KeySequence;
8 9 import com.baomidou.mybatisplus.annotation.TableField;
9 10 import com.baomidou.mybatisplus.annotation.TableId;
10 11 import com.baomidou.mybatisplus.annotation.TableName;
11   -import com.baomidou.mybatisplus.extension.handlers.JacksonTypeHandler;
12 12 import lombok.*;
13 13 import org.springframework.ai.chat.messages.MessageType;
14 14  
... ... @@ -71,13 +71,12 @@ public class AiChatMessageDO extends BaseDO {
71 71 */
72 72 private Long roleId;
73 73  
74   -
75 74 /**
76   - * 段落编号数组
  75 + * 知识库段落编号数组
77 76 *
78 77 * 关联 {@link AiKnowledgeSegmentDO#getId()} 字段
79 78 */
80   - @TableField(typeHandler = JacksonTypeHandler.class)
  79 + @TableField(typeHandler = LongListTypeHandler.class)
81 80 private List<Long> segmentIds;
82 81  
83 82 /**
... ...
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/chat/AiChatConversationServiceImpl.java
  View file @cddaca5
... ... @@ -68,7 +68,7 @@ public class AiChatConversationServiceImpl implements AiChatConversationService
68 68  
69 69 // 2. 创建 AiChatConversationDO 聊天对话
70 70 AiChatConversationDO conversation = new AiChatConversationDO().setUserId(userId).setPinned(false)
71   - .setModelId(model.getId()).setModel(model.getModel()).setKnowledgeId(createReqVO.getKnowledgeId())
  71 + .setModelId(model.getId()).setModel(model.getModel())
72 72 .setTemperature(model.getTemperature()).setMaxTokens(model.getMaxTokens()).setMaxContexts(model.getMaxContexts());
73 73 if (role != null) {
74 74 conversation.setTitle(role.getName()).setRoleId(role.getId()).setSystemMessage(role.getSystemMessage());
... ...
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/chat/AiChatMessageServiceImpl.java
  View file @cddaca5
... ... @@ -14,12 +14,14 @@ import cn.iocoder.yudao.module.ai.controller.admin.chat.vo.message.AiChatMessage
14 14 import cn.iocoder.yudao.module.ai.controller.admin.chat.vo.message.AiChatMessageSendRespVO;
15 15 import cn.iocoder.yudao.module.ai.dal.dataobject.chat.AiChatConversationDO;
16 16 import cn.iocoder.yudao.module.ai.dal.dataobject.chat.AiChatMessageDO;
17   -import cn.iocoder.yudao.module.ai.dal.dataobject.knowledge.AiKnowledgeSegmentDO;
  17 +import cn.iocoder.yudao.module.ai.dal.dataobject.model.AiChatRoleDO;
18 18 import cn.iocoder.yudao.module.ai.dal.dataobject.model.AiModelDO;
19 19 import cn.iocoder.yudao.module.ai.dal.mysql.chat.AiChatMessageMapper;
20   -import cn.iocoder.yudao.module.ai.enums.AiChatRoleEnum;
21 20 import cn.iocoder.yudao.module.ai.enums.ErrorCodeConstants;
22 21 import cn.iocoder.yudao.module.ai.service.knowledge.AiKnowledgeSegmentService;
  22 +import cn.iocoder.yudao.module.ai.service.knowledge.bo.AiKnowledgeSegmentSearchReqBO;
  23 +import cn.iocoder.yudao.module.ai.service.knowledge.bo.AiKnowledgeSegmentSearchRespBO;
  24 +import cn.iocoder.yudao.module.ai.service.model.AiChatRoleService;
23 25 import cn.iocoder.yudao.module.ai.service.model.AiModelService;
24 26 import jakarta.annotation.Resource;
25 27 import lombok.extern.slf4j.Slf4j;
... ... @@ -32,13 +34,13 @@ import org.springframework.ai.chat.model.ChatResponse;
32 34 import org.springframework.ai.chat.model.StreamingChatModel;
33 35 import org.springframework.ai.chat.prompt.ChatOptions;
34 36 import org.springframework.ai.chat.prompt.Prompt;
35   -import org.springframework.ai.chat.prompt.PromptTemplate;
36 37 import org.springframework.stereotype.Service;
37 38 import org.springframework.transaction.annotation.Transactional;
38 39 import reactor.core.publisher.Flux;
39 40  
40 41 import java.time.LocalDateTime;
41 42 import java.util.*;
  43 +import java.util.stream.Collectors;
42 44  
43 45 import static cn.iocoder.yudao.framework.common.exception.util.ServiceExceptionUtil.exception;
44 46 import static cn.iocoder.yudao.framework.common.pojo.CommonResult.error;
... ... @@ -56,12 +58,21 @@ import static cn.iocoder.yudao.module.ai.enums.ErrorCodeConstants.CHAT_MESSAGE_N
56 58 @Slf4j
57 59 public class AiChatMessageServiceImpl implements AiChatMessageService {
58 60  
  61 + /**
  62 + * 知识库转 {@link UserMessage} 的内容模版
  63 + */
  64 + private static final String KNOWLEDGE_USER_MESSAGE_TEMPLATE = "使用 <Reference></Reference> 标记中的内容作为本次对话的参考:\n\n" +
  65 + "%s\n\n" + // 多个 <Reference></Reference> 的拼接
  66 + "回答要求：\n- 避免提及你是从 <Reference></Reference> 获取的知识。";
  67 +
59 68 @Resource
60 69 private AiChatMessageMapper chatMessageMapper;
61 70  
62 71 @Resource
63 72 private AiChatConversationService chatConversationService;
64 73 @Resource
  74 + private AiChatRoleService chatRoleService;
  75 + @Resource
65 76 private AiModelService modalService;
66 77 @Resource
67 78 private AiKnowledgeSegmentService knowledgeSegmentService;
... ... @@ -69,118 +80,143 @@ public class AiChatMessageServiceImpl implements AiChatMessageService {
69 80 @Transactional(rollbackFor = Exception.class)
70 81 public AiChatMessageSendRespVO sendMessage(AiChatMessageSendReqVO sendReqVO, Long userId) {
71 82 // 1.1 校验对话存在
72   - AiChatConversationDO conversation = chatConversationService.validateChatConversationExists(sendReqVO.getConversationId());
  83 + AiChatConversationDO conversation = chatConversationService
  84 + .validateChatConversationExists(sendReqVO.getConversationId());
73 85 if (ObjUtil.notEqual(conversation.getUserId(), userId)) {
74 86 throw exception(CHAT_CONVERSATION_NOT_EXISTS);
75 87 }
76 88 List<AiChatMessageDO> historyMessages = chatMessageMapper.selectListByConversationId(conversation.getId());
77 89 // 1.2 校验模型
78 90 AiModelDO model = modalService.validateModel(conversation.getModelId());
79   - ChatModel chatModel = modalService.getChatModel(model.getKeyId());
  91 + ChatModel chatModel = modalService.getChatModel(model.getId());
  92 +
  93 + // 2. 知识库找回
  94 + List<AiKnowledgeSegmentSearchRespBO> knowledgeSegments = recallKnowledgeSegment(sendReqVO.getContent(),
  95 + conversation);
80 96  
81   - // 2. 插入 user 发送消息
  97 + // 3. 插入 user 发送消息
82 98 AiChatMessageDO userMessage = createChatMessage(conversation.getId(), null, model,
83   - userId, conversation.getRoleId(), MessageType.USER, sendReqVO.getContent(), sendReqVO.getUseContext());
  99 + userId, conversation.getRoleId(), MessageType.USER, sendReqVO.getContent(), sendReqVO.getUseContext(),
  100 + null);
84 101  
85 102 // 3.1 插入 assistant 接收消息
86 103 AiChatMessageDO assistantMessage = createChatMessage(conversation.getId(), userMessage.getId(), model,
87   - userId, conversation.getRoleId(), MessageType.ASSISTANT, "", sendReqVO.getUseContext());
  104 + userId, conversation.getRoleId(), MessageType.ASSISTANT, "", sendReqVO.getUseContext(),
  105 + knowledgeSegments);
88 106  
89   - // 3.2 召回段落
90   - List<AiKnowledgeSegmentDO> segmentList = recallSegment(sendReqVO.getContent(), conversation.getKnowledgeId());
91   -
92   - // 3.3 创建 chat 需要的 Prompt
93   - Prompt prompt = buildPrompt(conversation, historyMessages, segmentList, model, sendReqVO);
  107 + // 3.2 创建 chat 需要的 Prompt
  108 + Prompt prompt = buildPrompt(conversation, historyMessages, knowledgeSegments, model, sendReqVO);
94 109 ChatResponse chatResponse = chatModel.call(prompt);
95 110  
96   - // 3.4 段式返回
  111 + // 3.3 段式返回
97 112 String newContent = chatResponse.getResult().getOutput().getText();
98   - chatMessageMapper.updateById(new AiChatMessageDO().setId(assistantMessage.getId()).setSegmentIds(convertList(segmentList, AiKnowledgeSegmentDO::getId)).setContent(newContent));
99   - return new AiChatMessageSendRespVO().setSend(BeanUtils.toBean(userMessage, AiChatMessageSendRespVO.Message.class))
100   - .setReceive(BeanUtils.toBean(assistantMessage, AiChatMessageSendRespVO.Message.class).setContent(newContent));
  113 + chatMessageMapper.updateById(new AiChatMessageDO().setId(assistantMessage.getId()).setContent(newContent));
  114 + return new AiChatMessageSendRespVO()
  115 + .setSend(BeanUtils.toBean(userMessage, AiChatMessageSendRespVO.Message.class))
  116 + .setReceive(BeanUtils.toBean(assistantMessage, AiChatMessageSendRespVO.Message.class)
  117 + .setContent(newContent));
101 118 }
102 119  
103 120 @Override
104   - public Flux<CommonResult<AiChatMessageSendRespVO>> sendChatMessageStream(AiChatMessageSendReqVO sendReqVO, Long userId) {
  121 + public Flux<CommonResult<AiChatMessageSendRespVO>> sendChatMessageStream(AiChatMessageSendReqVO sendReqVO,
  122 + Long userId) {
105 123 // 1.1 校验对话存在
106   - AiChatConversationDO conversation = chatConversationService.validateChatConversationExists(sendReqVO.getConversationId());
  124 + AiChatConversationDO conversation = chatConversationService
  125 + .validateChatConversationExists(sendReqVO.getConversationId());
107 126 if (ObjUtil.notEqual(conversation.getUserId(), userId)) {
108 127 throw exception(CHAT_CONVERSATION_NOT_EXISTS);
109 128 }
110 129 List<AiChatMessageDO> historyMessages = chatMessageMapper.selectListByConversationId(conversation.getId());
111 130 // 1.2 校验模型
112 131 AiModelDO model = modalService.validateModel(conversation.getModelId());
113   - StreamingChatModel chatModel = modalService.getChatModel(model.getKeyId());
  132 + StreamingChatModel chatModel = modalService.getChatModel(model.getId());
  133 +
  134 + // 2. 知识库找回
  135 + List<AiKnowledgeSegmentSearchRespBO> knowledgeSegments = recallKnowledgeSegment(sendReqVO.getContent(),
  136 + conversation);
114 137  
115   - // 2. 插入 user 发送消息
  138 + // 3. 插入 user 发送消息
116 139 AiChatMessageDO userMessage = createChatMessage(conversation.getId(), null, model,
117   - userId, conversation.getRoleId(), MessageType.USER, sendReqVO.getContent(), sendReqVO.getUseContext());
  140 + userId, conversation.getRoleId(), MessageType.USER, sendReqVO.getContent(), sendReqVO.getUseContext(),
  141 + null);
118 142  
119   - // 3.1 插入 assistant 接收消息
  143 + // 4.1 插入 assistant 接收消息
120 144 AiChatMessageDO assistantMessage = createChatMessage(conversation.getId(), userMessage.getId(), model,
121   - userId, conversation.getRoleId(), MessageType.ASSISTANT, "", sendReqVO.getUseContext());
  145 + userId, conversation.getRoleId(), MessageType.ASSISTANT, "", sendReqVO.getUseContext(),
  146 + knowledgeSegments);
122 147  
123   - // 3.2 召回段落
124   - List<AiKnowledgeSegmentDO> segmentList = recallSegment(sendReqVO.getContent(), conversation.getKnowledgeId());
125   -
126   - // 3.3 构建 Prompt，并进行调用
127   - Prompt prompt = buildPrompt(conversation, historyMessages, segmentList, model, sendReqVO);
  148 + // 4.2 构建 Prompt，并进行调用
  149 + Prompt prompt = buildPrompt(conversation, historyMessages, knowledgeSegments, model, sendReqVO);
128 150 Flux<ChatResponse> streamResponse = chatModel.stream(prompt);
129 151  
130   - // 3.4 流式返回
  152 + // 4.3 流式返回
131 153 StringBuffer contentBuffer = new StringBuffer();
132 154 return streamResponse.map(chunk -> {
133 155 String newContent = chunk.getResult() != null ? chunk.getResult().getOutput().getText() : null;
134 156 newContent = StrUtil.nullToDefault(newContent, ""); // 避免 null 的 情况
135 157 contentBuffer.append(newContent);
136 158 // 响应结果
137   - return success(new AiChatMessageSendRespVO().setSend(BeanUtils.toBean(userMessage, AiChatMessageSendRespVO.Message.class))
138   - .setReceive(BeanUtils.toBean(assistantMessage, AiChatMessageSendRespVO.Message.class).setContent(newContent)));
  159 + return success(new AiChatMessageSendRespVO()
  160 + .setSend(BeanUtils.toBean(userMessage, AiChatMessageSendRespVO.Message.class))
  161 + .setReceive(BeanUtils.toBean(assistantMessage, AiChatMessageSendRespVO.Message.class)
  162 + .setContent(newContent)));
139 163 }).doOnComplete(() -> {
140 164 // 忽略租户，因为 Flux 异步无法透传租户
141   - TenantUtils.executeIgnore(() ->
142   - chatMessageMapper.updateById(new AiChatMessageDO().setId(assistantMessage.getId()).setSegmentIds(convertList(segmentList, AiKnowledgeSegmentDO::getId))
143   - .setContent(contentBuffer.toString())));
  165 + TenantUtils.executeIgnore(() -> chatMessageMapper.updateById(
  166 + new AiChatMessageDO().setId(assistantMessage.getId()).setContent(contentBuffer.toString())));
144 167 }).doOnError(throwable -> {
145 168 log.error("[sendChatMessageStream][userId({}) sendReqVO({}) 发生异常]", userId, sendReqVO, throwable);
146 169 // 忽略租户，因为 Flux 异步无法透传租户
147   - TenantUtils.executeIgnore(() ->
148   - chatMessageMapper.updateById(new AiChatMessageDO().setId(assistantMessage.getId()).setContent(throwable.getMessage())));
  170 + TenantUtils.executeIgnore(() -> chatMessageMapper.updateById(
  171 + new AiChatMessageDO().setId(assistantMessage.getId()).setContent(throwable.getMessage())));
149 172 }).onErrorResume(error -> Flux.just(error(ErrorCodeConstants.CHAT_STREAM_ERROR)));
150 173 }
151 174  
152   - private List<AiKnowledgeSegmentDO> recallSegment(String content, Long knowledgeId) {
153   - if (Objects.isNull(knowledgeId)) {
  175 + private List<AiKnowledgeSegmentSearchRespBO> recallKnowledgeSegment(String content,
  176 + AiChatConversationDO conversation) {
  177 + // 1. 查询聊天角色
  178 + if (conversation == null || conversation.getRoleId() == null) {
  179 + return Collections.emptyList();
  180 + }
  181 + AiChatRoleDO role = chatRoleService.getChatRole(conversation.getRoleId());
  182 + if (role == null || CollUtil.isEmpty(role.getKnowledgeIds())) {
154 183 return Collections.emptyList();
155 184 }
156   -// return knowledgeSegmentService.similaritySearch(new AiKnowledgeSegmentSearchReqVO().setKnowledgeId(knowledgeId).setContent(content));
157   - return null;
158   - }
159   -
160   - private Prompt buildPrompt(AiChatConversationDO conversation, List<AiChatMessageDO> messages, List<AiKnowledgeSegmentDO> segmentList,
161   - AiModelDO model, AiChatMessageSendReqVO sendReqVO) {
162   - // 1. 构建 Prompt Message 列表
163   - List<Message> chatMessages = new ArrayList<>();
164 185  
165   - // 1.1 召回内容消息构建
166   - if (CollUtil.isNotEmpty(segmentList)) {
167   - PromptTemplate promptTemplate = new PromptTemplate(AiChatRoleEnum.AI_KNOWLEDGE_ROLE.getSystemMessage());
168   - StringBuilder infoBuilder = StrUtil.builder();
169   - segmentList.forEach(segment -> infoBuilder.append(System.lineSeparator()).append(segment.getContent()));
170   - Message message = promptTemplate.createMessage(Map.of("info", infoBuilder.toString()));
171   - chatMessages.add(message);
  186 + // 2. 遍历找回
  187 + List<AiKnowledgeSegmentSearchRespBO> knowledgeSegments = new ArrayList<>();
  188 + for (Long knowledgeId : role.getKnowledgeIds()) {
  189 + knowledgeSegments.addAll(knowledgeSegmentService.searchKnowledgeSegment(new AiKnowledgeSegmentSearchReqBO()
  190 + .setKnowledgeId(knowledgeId).setContent(content)));
172 191 }
  192 + return knowledgeSegments;
  193 + }
173 194  
174   - // 1.2 system context 角色设定
  195 + private Prompt buildPrompt(AiChatConversationDO conversation, List<AiChatMessageDO> messages,
  196 + List<AiKnowledgeSegmentSearchRespBO> knowledgeSegments,
  197 + AiModelDO model, AiChatMessageSendReqVO sendReqVO) {
  198 + List<Message> chatMessages = new ArrayList<>();
  199 + // 1.1 System Context 角色设定
175 200 if (StrUtil.isNotBlank(conversation.getSystemMessage())) {
176 201 chatMessages.add(new SystemMessage(conversation.getSystemMessage()));
177 202 }
178   - // 1.3 history message 历史消息
  203 +
  204 + // 1.2 历史 history message 历史消息
179 205 List<AiChatMessageDO> contextMessages = filterContextMessages(messages, conversation, sendReqVO);
180   - contextMessages.forEach(message -> chatMessages.add(AiUtils.buildMessage(message.getType(), message.getContent())));
181   - // 1.4 user message 新发送消息
  206 + contextMessages
  207 + .forEach(message -> chatMessages.add(AiUtils.buildMessage(message.getType(), message.getContent())));
  208 +
  209 + // 1.3 当前 user message 新发送消息
182 210 chatMessages.add(new UserMessage(sendReqVO.getContent()));
183 211  
  212 + // 1.4 知识库，通过 UserMessage 实现
  213 + if (CollUtil.isNotEmpty(knowledgeSegments)) {
  214 + String reference = knowledgeSegments.stream()
  215 + .map(segment -> "<Reference>\n" + segment.getContent() + "</Reference>")
  216 + .collect(Collectors.joining("\n\n"));
  217 + chatMessages.add(new UserMessage(String.format(KNOWLEDGE_USER_MESSAGE_TEMPLATE, reference)));
  218 + }
  219 +
184 220 // 2. 构建 ChatOptions 对象
185 221 AiPlatformEnum platform = AiPlatformEnum.validatePlatform(model.getPlatform());
186 222 ChatOptions chatOptions = AiUtils.buildChatOptions(platform, model.getModel(),
... ... @@ -199,8 +235,8 @@ public class AiChatMessageServiceImpl implements AiChatMessageService {
199 235 * @return 消息上下文
200 236 */
201 237 private List<AiChatMessageDO> filterContextMessages(List<AiChatMessageDO> messages,
202   - AiChatConversationDO conversation,
203   - AiChatMessageSendReqVO sendReqVO) {
  238 + AiChatConversationDO conversation,
  239 + AiChatMessageSendReqVO sendReqVO) {
204 240 if (conversation.getMaxContexts() == null || ObjUtil.notEqual(sendReqVO.getUseContext(), Boolean.TRUE)) {
205 241 return Collections.emptyList();
206 242 }
... ... @@ -211,7 +247,8 @@ public class AiChatMessageServiceImpl implements AiChatMessageService {
211 247 continue;
212 248 }
213 249 AiChatMessageDO userMessage = CollUtil.get(messages, i - 1);
214   - if (userMessage == null || ObjUtil.notEqual(assistantMessage.getReplyId(), userMessage.getId())
  250 + if (userMessage == null
  251 + || ObjUtil.notEqual(assistantMessage.getReplyId(), userMessage.getId())
215 252 || StrUtil.isEmpty(assistantMessage.getContent())) {
216 253 continue;
217 254 }
... ... @@ -228,11 +265,13 @@ public class AiChatMessageServiceImpl implements AiChatMessageService {
228 265 }
229 266  
230 267 private AiChatMessageDO createChatMessage(Long conversationId, Long replyId,
231   - AiModelDO model, Long userId, Long roleId,
232   - MessageType messageType, String content, Boolean useContext) {
  268 + AiModelDO model, Long userId, Long roleId,
  269 + MessageType messageType, String content, Boolean useContext,
  270 + List<AiKnowledgeSegmentSearchRespBO> knowledgeSegments) {
233 271 AiChatMessageDO message = new AiChatMessageDO().setConversationId(conversationId).setReplyId(replyId)
234 272 .setModel(model.getModel()).setModelId(model.getId()).setUserId(userId).setRoleId(roleId)
235   - .setType(messageType.getValue()).setContent(content).setUseContext(useContext);
  273 + .setType(messageType.getValue()).setContent(content).setUseContext(useContext)
  274 + .setSegmentIds(convertList(knowledgeSegments, AiKnowledgeSegmentSearchRespBO::getId));
236 275 message.setCreateTime(LocalDateTime.now());
237 276 chatMessageMapper.insert(message);
238 277 return message;
... ...