Commit 2a01f6b9fca6826331e3cea2de549d8514667896
1 parent
ccd7928a
【增加】AI Chat stream 增加头像返回
Showing
2 changed files
with
23 additions
and
2 deletions
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 | } | ... | ... |