Commit b709af11a138b9ebadc0877980a38071304ebe6a

Authored by YunaiV
1 parent cddaca58

【功能新增】AI:聊天记录返回时,增加 segments

yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/controller/admin/chat/AiChatMessageController.java
... ... @@ -12,9 +12,13 @@ import cn.iocoder.yudao.module.ai.controller.admin.chat.vo.message.AiChatMessage
12 12 import cn.iocoder.yudao.module.ai.controller.admin.chat.vo.message.AiChatMessageSendRespVO;
13 13 import cn.iocoder.yudao.module.ai.dal.dataobject.chat.AiChatConversationDO;
14 14 import cn.iocoder.yudao.module.ai.dal.dataobject.chat.AiChatMessageDO;
  15 +import cn.iocoder.yudao.module.ai.dal.dataobject.knowledge.AiKnowledgeDocumentDO;
  16 +import cn.iocoder.yudao.module.ai.dal.dataobject.knowledge.AiKnowledgeSegmentDO;
15 17 import cn.iocoder.yudao.module.ai.dal.dataobject.model.AiChatRoleDO;
16 18 import cn.iocoder.yudao.module.ai.service.chat.AiChatConversationService;
17 19 import cn.iocoder.yudao.module.ai.service.chat.AiChatMessageService;
  20 +import cn.iocoder.yudao.module.ai.service.knowledge.AiKnowledgeDocumentService;
  21 +import cn.iocoder.yudao.module.ai.service.knowledge.AiKnowledgeSegmentService;
18 22 import cn.iocoder.yudao.module.ai.service.model.AiChatRoleService;
19 23 import io.swagger.v3.oas.annotations.Operation;
20 24 import io.swagger.v3.oas.annotations.Parameter;
... ... @@ -32,7 +36,7 @@ import java.util.List;
32 36 import java.util.Map;
33 37  
34 38 import static cn.iocoder.yudao.framework.common.pojo.CommonResult.success;
35   -import static cn.iocoder.yudao.framework.common.util.collection.CollectionUtils.convertSet;
  39 +import static cn.iocoder.yudao.framework.common.util.collection.CollectionUtils.*;
36 40 import static cn.iocoder.yudao.framework.security.core.util.SecurityFrameworkUtils.getLoginUserId;
37 41  
38 42 @Tag(name = "管理后台 - 聊天消息")
... ... @@ -47,6 +51,10 @@ public class AiChatMessageController {
47 51 private AiChatConversationService chatConversationService;
48 52 @Resource
49 53 private AiChatRoleService chatRoleService;
  54 + @Resource
  55 + private AiKnowledgeSegmentService knowledgeSegmentService;
  56 + @Resource
  57 + private AiKnowledgeDocumentService knowledgeDocumentService;
50 58  
51 59 @Operation(summary = "发送消息(段式)", description = "一次性返回,响应较慢")
52 60 @PostMapping("/send")
... ... @@ -56,7 +64,8 @@ public class AiChatMessageController {
56 64  
57 65 @Operation(summary = "发送消息(流式)", description = "流式返回,响应较快")
58 66 @PostMapping(value = "/send-stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
59   - public Flux<CommonResult<AiChatMessageSendRespVO>> sendChatMessageStream(@Valid @RequestBody AiChatMessageSendReqVO sendReqVO) {
  67 + public Flux<CommonResult<AiChatMessageSendRespVO>> sendChatMessageStream(
  68 + @Valid @RequestBody AiChatMessageSendReqVO sendReqVO) {
60 69 return chatMessageService.sendChatMessageStream(sendReqVO, getLoginUserId());
61 70 }
62 71  
... ... @@ -69,8 +78,38 @@ public class AiChatMessageController {
69 78 if (conversation == null || ObjUtil.notEqual(conversation.getUserId(), getLoginUserId())) {
70 79 return success(Collections.emptyList());
71 80 }
  81 + // 1. 获取消息列表
72 82 List<AiChatMessageDO> messageList = chatMessageService.getChatMessageListByConversationId(conversationId);
73   - return success(BeanUtils.toBean(messageList, AiChatMessageRespVO.class));
  83 + if (CollUtil.isEmpty(messageList)) {
  84 + return success(Collections.emptyList());
  85 + }
  86 +
  87 + // 2. 拼接数据,主要是知识库段落信息
  88 + Map<Long, AiKnowledgeSegmentDO> segmentMap = knowledgeSegmentService.getKnowledgeSegmentMap(convertListByFlatMap(messageList,
  89 + message -> CollUtil.isEmpty(message.getSegmentIds()) ? null : message.getSegmentIds().stream()));
  90 + Map<Long, AiKnowledgeDocumentDO> documentMap = knowledgeDocumentService.getKnowledgeDocumentMap(
  91 + convertList(segmentMap.values(), AiKnowledgeSegmentDO::getDocumentId));
  92 + List<AiChatMessageRespVO> messageVOList = BeanUtils.toBean(messageList, AiChatMessageRespVO.class);
  93 + for (int i = 0; i < messageList.size(); i++) {
  94 + AiChatMessageDO message = messageList.get(i);
  95 + if (CollUtil.isEmpty(message.getSegmentIds())) {
  96 + continue;
  97 + }
  98 + // 设置知识库段落信息
  99 + messageVOList.get(i).setSegments(convertList(message.getSegmentIds(), segmentId -> {
  100 + AiKnowledgeSegmentDO segment = segmentMap.get(segmentId);
  101 + if (segment == null) {
  102 + return null;
  103 + }
  104 + AiKnowledgeDocumentDO document = documentMap.get(segment.getDocumentId());
  105 + if (document == null) {
  106 + return null;
  107 + }
  108 + return new AiChatMessageRespVO.KnowledgeSegment().setId(segment.getId()).setContent(segment.getContent())
  109 + .setDocumentId(segment.getDocumentId()).setDocumentName(document.getName());
  110 + }));
  111 + }
  112 + return success(messageVOList);
74 113 }
75 114  
76 115 @Operation(summary = "删除消息")
... ... @@ -84,7 +123,8 @@ public class AiChatMessageController {
84 123 @Operation(summary = "删除指定对话的消息")
85 124 @DeleteMapping("/delete-by-conversation-id")
86 125 @Parameter(name = "conversationId", required = true, description = "对话编号", example = "1024")
87   - public CommonResult<Boolean> deleteChatMessageByConversationId(@RequestParam("conversationId") Long conversationId) {
  126 + public CommonResult<Boolean> deleteChatMessageByConversationId(
  127 + @RequestParam("conversationId") Long conversationId) {
88 128 chatMessageService.deleteChatMessageByConversationId(conversationId, getLoginUserId());
89 129 return success(true);
90 130 }
... ... @@ -103,7 +143,8 @@ public class AiChatMessageController {
103 143 Map<Long, AiChatRoleDO> roleMap = chatRoleService.getChatRoleMap(
104 144 convertSet(pageResult.getList(), AiChatMessageDO::getRoleId));
105 145 return success(BeanUtils.toBean(pageResult, AiChatMessageRespVO.class,
106   - respVO -> MapUtils.findAndThen(roleMap, respVO.getRoleId(), role -> respVO.setRoleName(role.getName()))));
  146 + respVO -> MapUtils.findAndThen(roleMap, respVO.getRoleId(),
  147 + role -> respVO.setRoleName(role.getName()))));
107 148 }
108 149  
109 150 @Operation(summary = "管理员删除消息")
... ...
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/controller/admin/chat/vo/message/AiChatMessageRespVO.java
... ... @@ -4,6 +4,7 @@ import io.swagger.v3.oas.annotations.media.Schema;
4 4 import lombok.Data;
5 5  
6 6 import java.time.LocalDateTime;
  7 +import java.util.List;
7 8  
8 9 @Schema(description = "管理后台 - AI 聊天消息 Response VO")
9 10 @Data
... ... @@ -39,6 +40,12 @@ public class AiChatMessageRespVO {
39 40 @Schema(description = "是否携带上下文", requiredMode = Schema.RequiredMode.REQUIRED, example = "true")
40 41 private Boolean useContext;
41 42  
  43 + @Schema(description = "知识库段落编号数组", example = "[1,2,3]")
  44 + private List<Long> segmentIds;
  45 +
  46 + @Schema(description = "知识库段落数组")
  47 + private List<KnowledgeSegment> segments;
  48 +
42 49 @Schema(description = "创建时间", requiredMode = Schema.RequiredMode.REQUIRED, example = "2024-05-12 12:51")
43 50 private LocalDateTime createTime;
44 51  
... ... @@ -47,4 +54,22 @@ public class AiChatMessageRespVO {
47 54 @Schema(description = "角色名字", example = "小黄")
48 55 private String roleName;
49 56  
  57 + @Schema(description = "知识库段落", example = "Java 开发手册")
  58 + @Data
  59 + public static class KnowledgeSegment {
  60 +
  61 + @Schema(description = "段落编号", requiredMode = Schema.RequiredMode.REQUIRED, example = "1024")
  62 + private Long id;
  63 +
  64 + @Schema(description = "切片内容", requiredMode = Schema.RequiredMode.REQUIRED, example = "Java 开发手册")
  65 + private String content;
  66 +
  67 + @Schema(description = "文档编号", requiredMode = Schema.RequiredMode.REQUIRED, example = "24790")
  68 + private Long documentId;
  69 +
  70 + @Schema(description = "文档名称", requiredMode = Schema.RequiredMode.REQUIRED, example = "产品使用手册")
  71 + private String documentName;
  72 +
  73 + }
  74 +
50 75 }
... ...
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/controller/admin/knowledge/AiKnowledgeController.java
... ... @@ -23,7 +23,6 @@ import java.util.List;
23 23 import static cn.iocoder.yudao.framework.common.pojo.CommonResult.success;
24 24 import static cn.iocoder.yudao.framework.common.util.collection.CollectionUtils.convertList;
25 25  
26   -// TODO @芋艿:增加权限标识
27 26 @Tag(name = "管理后台 - AI 知识库")
28 27 @RestController
29 28 @RequestMapping("/ai/knowledge")
... ...
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/dal/dataobject/chat/AiChatMessageDO.java
... ... @@ -20,7 +20,7 @@ import java.util.List;
20 20 * @since 2024/4/14 17:35
21 21 * @since 2024/4/14 17:35
22 22 */
23   -@TableName("ai_chat_message")
  23 +@TableName(value = "ai_chat_message", autoResultMap = true)
24 24 @KeySequence("ai_chat_conversation_seq") // 用于 Oracle、PostgreSQL、Kingbase、DB2、H2 数据库的主键自增。如果是 MySQL 等数据库,可不写。
25 25 @Data
26 26 @EqualsAndHashCode(callSuper = true)
... ... @@ -72,14 +72,6 @@ public class AiChatMessageDO extends BaseDO {
72 72 private Long roleId;
73 73  
74 74 /**
75   - * 知识库段落编号数组
76   - *
77   - * 关联 {@link AiKnowledgeSegmentDO#getId()} 字段
78   - */
79   - @TableField(typeHandler = LongListTypeHandler.class)
80   - private List<Long> segmentIds;
81   -
82   - /**
83 75 * 模型标志
84 76 *
85 77 * 冗余 {@link AiModelDO#getModel()}
... ... @@ -102,4 +94,12 @@ public class AiChatMessageDO extends BaseDO {
102 94 */
103 95 private Boolean useContext;
104 96  
  97 + /**
  98 + * 知识库段落编号数组
  99 + *
  100 + * 关联 {@link AiKnowledgeSegmentDO#getId()} 字段
  101 + */
  102 + @TableField(typeHandler = LongListTypeHandler.class)
  103 + private List<Long> segmentIds;
  104 +
105 105 }
... ...
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/chat/AiChatMessageServiceImpl.java
... ... @@ -212,7 +212,7 @@ public class AiChatMessageServiceImpl implements AiChatMessageService {
212 212 // 1.4 知识库,通过 UserMessage 实现
213 213 if (CollUtil.isNotEmpty(knowledgeSegments)) {
214 214 String reference = knowledgeSegments.stream()
215   - .map(segment -> "<Reference>\n" + segment.getContent() + "</Reference>")
  215 + .map(segment -> "<Reference>" + segment.getContent() + "</Reference>")
216 216 .collect(Collectors.joining("\n\n"));
217 217 chatMessages.add(new UserMessage(String.format(KNOWLEDGE_USER_MESSAGE_TEMPLATE, reference)));
218 218 }
... ...
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/knowledge/AiKnowledgeDocumentServiceImpl.java
... ... @@ -26,6 +26,7 @@ import org.springframework.transaction.annotation.Transactional;
26 26  
27 27 import java.util.ArrayList;
28 28 import java.util.Collection;
  29 +import java.util.Collections;
29 30 import java.util.List;
30 31  
31 32 import static cn.iocoder.yudao.framework.common.exception.util.ServiceExceptionUtil.exception;
... ... @@ -205,9 +206,9 @@ public class AiKnowledgeDocumentServiceImpl implements AiKnowledgeDocumentServic
205 206 @Override
206 207 public List<AiKnowledgeDocumentDO> getKnowledgeDocumentList(Collection<Long> ids) {
207 208 if (CollUtil.isEmpty(ids)) {
208   - return new ArrayList<>();
  209 + return Collections.emptyList();
209 210 }
210   - return knowledgeDocumentMapper.selectByIds(ids);
  211 + return knowledgeDocumentMapper.selectBatchIds(ids);
211 212 }
212 213  
213 214 }
... ...
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/knowledge/AiKnowledgeSegmentService.java
... ... @@ -10,7 +10,11 @@ import cn.iocoder.yudao.module.ai.service.knowledge.bo.AiKnowledgeSegmentSearchR
10 10 import cn.iocoder.yudao.module.ai.service.knowledge.bo.AiKnowledgeSegmentSearchRespBO;
11 11 import org.springframework.scheduling.annotation.Async;
12 12  
  13 +import java.util.Collection;
13 14 import java.util.List;
  15 +import java.util.Map;
  16 +
  17 +import static cn.iocoder.yudao.framework.common.util.collection.CollectionUtils.convertMap;
14 18  
15 19 /**
16 20 * AI 知识库段落 Service 接口
... ... @@ -28,6 +32,24 @@ public interface AiKnowledgeSegmentService {
28 32 AiKnowledgeSegmentDO getKnowledgeSegment(Long id);
29 33  
30 34 /**
  35 + * 获取知识库段落列表
  36 + *
  37 + * @param ids 段落编号列表
  38 + * @return 段落列表
  39 + */
  40 + List<AiKnowledgeSegmentDO> getKnowledgeSegmentList(Collection<Long> ids);
  41 +
  42 + /**
  43 + * 获取知识库段落 Map
  44 + *
  45 + * @param ids 段落编号列表
  46 + * @return 段落 Map
  47 + */
  48 + default Map<Long, AiKnowledgeSegmentDO> getKnowledgeSegmentMap(Collection<Long> ids) {
  49 + return convertMap(getKnowledgeSegmentList(ids), AiKnowledgeSegmentDO::getId);
  50 + }
  51 +
  52 + /**
31 53 * 获取段落分页
32 54 *
33 55 * @param pageReqVO 分页查询
... ...
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/knowledge/AiKnowledgeSegmentServiceImpl.java
... ... @@ -30,10 +30,7 @@ import org.springframework.ai.vectorstore.filter.FilterExpressionBuilder;
30 30 import org.springframework.context.annotation.Lazy;
31 31 import org.springframework.stereotype.Service;
32 32  
33   -import java.util.Collections;
34   -import java.util.List;
35   -import java.util.Map;
36   -import java.util.Objects;
  33 +import java.util.*;
37 34  
38 35 import static cn.iocoder.yudao.framework.common.exception.util.ServiceExceptionUtil.exception;
39 36 import static cn.iocoder.yudao.framework.common.util.collection.CollectionUtils.convertList;
... ... @@ -322,7 +319,15 @@ public class AiKnowledgeSegmentServiceImpl implements AiKnowledgeSegmentService
322 319  
323 320 @Override
324 321 public AiKnowledgeSegmentDO getKnowledgeSegment(Long id) {
325   - return validateKnowledgeSegmentExists(id);
  322 + return segmentMapper.selectById(id);
  323 + }
  324 +
  325 + @Override
  326 + public List<AiKnowledgeSegmentDO> getKnowledgeSegmentList(Collection<Long> ids) {
  327 + if (CollUtil.isEmpty(ids)) {
  328 + return Collections.emptyList();
  329 + }
  330 + return segmentMapper.selectBatchIds(ids);
326 331 }
327 332  
328 333 }
... ...