reactor：【AI 大模型】优化消息发送后，读取文档的逻辑，避免 N 次读取

YunaiV
1 parent 3e8596ee
Showing 4 changed files with 13 additions and 27 deletions
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/dal/mysql/knowledge/AiKnowledgeSegmentMapper.java
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/chat/AiChatMessageServiceImpl.java
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/knowledge/AiKnowledgeDocumentService.java
yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/knowledge/AiKnowledgeDocumentServiceImpl.java
@@ -30,9 +30,9 @@ public interface AiKnowledgeSegmentMapper extends BaseMapperX&lt;AiKnowledgeSegment
                 .orderByDesc(AiKnowledgeSegmentDO::getId));
     }
-    default List<AiKnowledgeSegmentDO> selectListByVectorIds(List<String> vectorIdList) {
+    default List<AiKnowledgeSegmentDO> selectListByVectorIds(List<String> vectorIds) {
         return selectList(new LambdaQueryWrapperX<AiKnowledgeSegmentDO>()
-                .in(AiKnowledgeSegmentDO::getVectorId, vectorIdList)
+                .in(AiKnowledgeSegmentDO::getVectorId, vectorIds)
                 .orderByDesc(AiKnowledgeSegmentDO::getId));
     }
@@ -121,11 +121,11 @@ public class AiChatMessageServiceImpl implements AiChatMessageService {
         String newContent = chatResponse.getResult().getOutput().getText();
         chatMessageMapper.updateById(new AiChatMessageDO().setId(assistantMessage.getId()).setContent(newContent));
         // 3.4 响应结果
+        Map<Long, AiKnowledgeDocumentDO> documentMap = knowledgeDocumentService.getKnowledgeDocumentMap(
+                convertSet(knowledgeSegments, AiKnowledgeSegmentSearchRespBO::getDocumentId));
         List<AiChatMessageRespVO.KnowledgeSegment> segments = BeanUtils.toBean(knowledgeSegments,
-                AiChatMessageRespVO.KnowledgeSegment.class,
-                segment -> {
-                    AiKnowledgeDocumentDO document = knowledgeDocumentService
-                            .getKnowledgeDocument(segment.getDocumentId());
+                AiChatMessageRespVO.KnowledgeSegment.class, segment -> {
+                    AiKnowledgeDocumentDO document = documentMap.get(segment.getDocumentId());
                     segment.setDocumentName(document != null ? document.getName() : null);
                 });
         return new AiChatMessageSendRespVO()
@@ -172,12 +172,13 @@ public class AiChatMessageServiceImpl implements AiChatMessageService {
             // 处理知识库的返回，只有首次才有
             List<AiChatMessageRespVO.KnowledgeSegment> segments = null;
             if (StrUtil.isEmpty(contentBuffer)) {
-                segments = BeanUtils.toBean(knowledgeSegments, AiChatMessageRespVO.KnowledgeSegment.class,
-                        segment -> TenantUtils.executeIgnore(() -> {
-                            AiKnowledgeDocumentDO document = knowledgeDocumentService
-                                    .getKnowledgeDocument(segment.getDocumentId());
-                            segment.setDocumentName(document != null ? document.getName() : null);
-                        }));
+                Map<Long, AiKnowledgeDocumentDO> documentMap = TenantUtils.executeIgnore(() ->
+                        knowledgeDocumentService.getKnowledgeDocumentMap(
+                                convertSet(knowledgeSegments, AiKnowledgeSegmentSearchRespBO::getDocumentId)));
+                segments = BeanUtils.toBean(knowledgeSegments, AiChatMessageRespVO.KnowledgeSegment.class, segment ->  {
+                    AiKnowledgeDocumentDO document = documentMap.get(segment.getDocumentId());
+                    segment.setDocumentName(document != null ? document.getName() : null);
+                });
             }
             // 响应结果
             String newContent = chunk.getResult() != null ? chunk.getResult().getOutput().getText() : null;
@@ -68,13 +68,6 @@ public interface AiKnowledgeDocumentService {
     void updateKnowledgeDocumentStatus(AiKnowledgeDocumentUpdateStatusReqVO reqVO);
     /**
-     * 更新文档检索次数（增加 +1）
-     *
-     * @param ids 文档编号列表
-     */
-    void updateKnowledgeDocumentRetrievalCountIncr(Collection<Long> ids);
-
-    /**
      * 删除文档
      *
      * @param id 文档编号
@@ -162,14 +162,6 @@ public class AiKnowledgeDocumentServiceImpl implements AiKnowledgeDocumentServic
     }
     @Override
-    public void updateKnowledgeDocumentRetrievalCountIncr(Collection<Long> ids) {
-        if (CollUtil.isEmpty(ids)) {
-            return;
-        }
-        knowledgeDocumentMapper.updateRetrievalCountIncr(ids);
-    }
-
-    @Override
     public AiKnowledgeDocumentDO validateKnowledgeDocumentExists(Long id) {
         AiKnowledgeDocumentDO knowledgeDocument = knowledgeDocumentMapper.selectById(id);
         if (knowledgeDocument == null) {