【增加】AI message 增加用户头像、role头像

cherishsince
1 parent 747b2ebc
Showing 4 changed files with 36 additions and 7 deletions
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/controller/admin/chat/vo/message/AiChatMessageRespVO.java
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/impl/AiChatServiceImpl.java
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/model/AiChatRoleService.java
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/model/AiChatRoleServiceImpl.java
@@ -36,4 +36,11 @@ public class AiChatMessageRespVO {
     @Schema(description = "创建时间", requiredMode = Schema.RequiredMode.REQUIRED, example = "2024-05-12 12:51")
     private LocalDateTime createTime;
+    // ========= 扩展字段
+
+    @Schema(description = "用户头像", requiredMode = Schema.RequiredMode.REQUIRED, example = "http://xxx")
+    private String userAvatar;
+
+    @Schema(description = "角色头像", requiredMode = Schema.RequiredMode.REQUIRED, example = "http://xxx")
+    private String roleAvatar;
 }
@@ -13,6 +13,7 @@ import cn.iocoder.yudao.module.ai.convert.AiChatMessageConvert;
 import cn.iocoder.yudao.module.ai.dal.dataobject.chat.AiChatConversationDO;
 import cn.iocoder.yudao.module.ai.dal.dataobject.chat.AiChatMessageDO;
 import cn.iocoder.yudao.module.ai.dal.dataobject.model.AiChatModelDO;
+import cn.iocoder.yudao.module.ai.dal.dataobject.model.AiChatRoleDO;
 import cn.iocoder.yudao.module.ai.dal.mysql.AiChatMessageMapper;
 import cn.iocoder.yudao.module.ai.service.AiChatService;
 import cn.iocoder.yudao.module.ai.service.chat.AiChatConversationService;
@@ -216,16 +217,22 @@ public class AiChatServiceImpl implements AiChatService {
         // 获取对话所有 message
         List<AiChatMessageDO> aiChatMessageDOList = chatMessageMapper.selectByConversationId(conversationId);
         // 获取模型信息
-        Set<Long> modalIds = aiChatMessageDOList.stream().map(AiChatMessageDO::getModelId).collect(Collectors.toSet());
-        List<AiChatModelDO> modalList = chatModalService.getModalByIds(modalIds);
-        Map<Long, AiChatModelDO> modalIdMap = modalList.stream().collect(Collectors.toMap(AiChatModelDO::getId, o -> o));
+        Set<Long> roleIds = aiChatMessageDOList.stream().map(AiChatMessageDO::getRoleId).collect(Collectors.toSet());
+        List<AiChatRoleDO> roleList;
+        if (!CollUtil.isEmpty(roleIds)) {
+            roleList = chatRoleService.getChatRoles(roleIds);
+        } else {
+            roleList = Collections.emptyList();
+        }
+        Map<Long, AiChatRoleDO> roleMap = roleList.stream().collect(Collectors.toMap(AiChatRoleDO::getId, o -> o));
         // 转换 AiChatMessageRespVO
         List<AiChatMessageRespVO> aiChatMessageRespList = AiChatMessageConvert.INSTANCE.convertAiChatMessageRespVOList(aiChatMessageDOList);
-        // 设置用户头像 和 模型头像 todo @芋艿 这里需要转换 用户头像、模型头像
+        // 设置用户头像 和 模型头像
         return aiChatMessageRespList.stream().map(item -> {
-            if (modalIdMap.containsKey(item.getModelId())) {
-//                modalIdMap.get(item.getModelId());
-//                item.setModelImage()
+            // 设置 role 头像
+            if (roleMap.containsKey(item.getRoleId())) {
+                AiChatRoleDO role = roleMap.get(item.getRoleId());
+                item.setRoleAvatar(role.getAvatar());
             }
             return item;
         }).collect(Collectors.toList());
@@ -9,6 +9,7 @@ import cn.iocoder.yudao.module.ai.dal.dataobject.model.AiChatRoleDO;
 import jakarta.validation.Valid;
 import java.util.List;
+import java.util.Set;
 /**
  * AI 聊天角色 Service 接口
@@ -73,6 +74,14 @@ public interface AiChatRoleService {
     AiChatRoleDO getChatRole(Long id);
     /**
+     * 获得聊天角色 - 根据 ids
+     *
+     * @param roleIds
+     * @return
+     */
+    List<AiChatRoleDO> getChatRoles(Set<Long> roleIds);
+
+    /**
      * 校验聊天角色是否合法
      *
      * @param id 角色编号
@@ -16,6 +16,7 @@ import org.springframework.stereotype.Service;
 import java.util.List;
 import java.util.Objects;
+import java.util.Set;
 import java.util.function.Function;
 import java.util.function.Predicate;
 import java.util.stream.Collectors;
@@ -106,6 +107,11 @@ public class AiChatRoleServiceImpl implements AiChatRoleService {
     }
     @Override
+    public List<AiChatRoleDO> getChatRoles(Set<Long> roleIds) {
+        return chatRoleMapper.selectBatchIds(roleIds);
+    }
+
+    @Override
     public AiChatRoleDO validateChatRole(Long id) {
         AiChatRoleDO chatRole = validateChatRoleExists(id);
         if (CommonStatusEnum.isDisable(chatRole.getStatus())) {