Commit 2a01f6b9fca6826331e3cea2de549d8514667896

Authored by cherishsince
1 parent ccd7928a

【增加】AI Chat stream 增加头像返回

yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/controller/admin/chat/vo/message/AiChatMessageSendRespVO.java
... ... @@ -31,6 +31,13 @@ public class AiChatMessageSendRespVO {
31 31 @Schema(description = "创建时间", requiredMode = Schema.RequiredMode.REQUIRED)
32 32 private LocalDateTime createTime;
33 33  
  34 + // ========= 扩展字段
  35 +
  36 + @Schema(description = "用户头像", requiredMode = Schema.RequiredMode.REQUIRED, example = "http://xxx")
  37 + private String userAvatar;
  38 +
  39 + @Schema(description = "角色头像", requiredMode = Schema.RequiredMode.REQUIRED, example = "http://xxx")
  40 + private String roleAvatar;
34 41 }
35 42  
36 43 }
... ...
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/impl/AiChatServiceImpl.java
... ... @@ -5,6 +5,7 @@ import cn.hutool.core.util.ObjUtil;
5 5 import cn.hutool.core.util.StrUtil;
6 6 import cn.iocoder.yudao.framework.ai.core.enums.AiPlatformEnum;
7 7 import cn.iocoder.yudao.framework.common.util.object.BeanUtils;
  8 +import cn.iocoder.yudao.framework.security.core.util.SecurityFrameworkUtils;
8 9 import cn.iocoder.yudao.module.ai.config.AiChatClientFactory;
9 10 import cn.iocoder.yudao.module.ai.controller.admin.chat.vo.message.AiChatMessageRespVO;
10 11 import cn.iocoder.yudao.module.ai.controller.admin.chat.vo.message.AiChatMessageSendReqVO;
... ... @@ -19,6 +20,8 @@ import cn.iocoder.yudao.module.ai.service.AiChatService;
19 20 import cn.iocoder.yudao.module.ai.service.chat.AiChatConversationService;
20 21 import cn.iocoder.yudao.module.ai.service.model.AiChatModelService;
21 22 import cn.iocoder.yudao.module.ai.service.model.AiChatRoleService;
  23 +import cn.iocoder.yudao.module.system.api.user.AdminUserApi;
  24 +import cn.iocoder.yudao.module.system.api.user.dto.AdminUserRespDTO;
22 25 import lombok.AllArgsConstructor;
23 26 import lombok.extern.slf4j.Slf4j;
24 27 import org.springframework.ai.chat.ChatResponse;
... ... @@ -56,6 +59,7 @@ public class AiChatServiceImpl implements AiChatService {
56 59 private final AiChatConversationService chatConversationService;
57 60 private final AiChatModelService chatModalService;
58 61 private final AiChatRoleService chatRoleService;
  62 + private final AdminUserApi adminUserApi;
59 63  
60 64 @Transactional(rollbackFor = Exception.class)
61 65 public AiChatMessageRespVO chat(AiChatMessageSendReqVO req) {
... ... @@ -124,13 +128,17 @@ public class AiChatServiceImpl implements AiChatService {
124 128 // 3.3 流式返回
125 129 // 注意:Schedulers.immediate() 目的是,避免默认 Schedulers.parallel() 并发消费 chunk 导致 SSE 响应前端会乱序问题
126 130 StringBuffer contentBuffer = new StringBuffer();
  131 +
  132 + // 3.4 获取用户头像、角色头像
  133 + AdminUserRespDTO user = adminUserApi.getUser(SecurityFrameworkUtils.getLoginUserId());
  134 + AiChatRoleDO chatRole = chatRoleService.getChatRole(assistantMessage.getRoleId());
127 135 return streamResponse.publishOn(Schedulers.immediate()).map(chunk -> {
128 136 String newContent = chunk.getResult() != null ? chunk.getResult().getOutput().getContent() : null;
129 137 newContent = StrUtil.nullToDefault(newContent, ""); // 避免 null 的 情况
130 138 contentBuffer.append(newContent);
131 139 // 响应结果
132   - return new AiChatMessageSendRespVO().setSend(BeanUtils.toBean(userMessage, AiChatMessageSendRespVO.Message.class))
133   - .setReceive(BeanUtils.toBean(assistantMessage, AiChatMessageSendRespVO.Message.class).setContent(newContent));
  140 + return new AiChatMessageSendRespVO().setSend(BeanUtils.toBean(userMessage, AiChatMessageSendRespVO.Message.class).setUserAvatar(user.getAvatar()))
  141 + .setReceive(BeanUtils.toBean(assistantMessage, AiChatMessageSendRespVO.Message.class).setContent(newContent).setRoleAvatar(chatRole == null ? null : chatRole.getAvatar()));
134 142 }).doOnComplete(() -> {
135 143 chatMessageMapper.updateById(new AiChatMessageDO().setId(assistantMessage.getId()).setContent(contentBuffer.toString()));
136 144 }).doOnError(throwable -> {
... ... @@ -227,6 +235,8 @@ public class AiChatServiceImpl implements AiChatService {
227 235 Map<Long, AiChatRoleDO> roleMap = roleList.stream().collect(Collectors.toMap(AiChatRoleDO::getId, o -> o));
228 236 // 转换 AiChatMessageRespVO
229 237 List<AiChatMessageRespVO> aiChatMessageRespList = AiChatMessageConvert.INSTANCE.convertAiChatMessageRespVOList(aiChatMessageDOList);
  238 + // 获取用户信息
  239 + AdminUserRespDTO user = adminUserApi.getUser(SecurityFrameworkUtils.getLoginUserId());
230 240 // 设置用户头像 和 模型头像
231 241 return aiChatMessageRespList.stream().map(item -> {
232 242 // 设置 role 头像
... ... @@ -234,6 +244,10 @@ public class AiChatServiceImpl implements AiChatService {
234 244 AiChatRoleDO role = roleMap.get(item.getRoleId());
235 245 item.setRoleAvatar(role.getAvatar());
236 246 }
  247 + // 设置 user 头像
  248 + if (user != null) {
  249 + item.setUserAvatar(user.getAvatar());
  250 + }
237 251 return item;
238 252 }).collect(Collectors.toList());
239 253 }
... ...