diff --git a/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/dal/mysql/knowledge/AiKnowledgeSegmentMapper.java b/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/dal/mysql/knowledge/AiKnowledgeSegmentMapper.java
index 00bacd9665..1b9ca867f5 100644
--- a/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/dal/mysql/knowledge/AiKnowledgeSegmentMapper.java
+++ b/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/dal/mysql/knowledge/AiKnowledgeSegmentMapper.java
@@ -30,9 +30,9 @@ public interface AiKnowledgeSegmentMapper extends BaseMapperX<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));
     }
 
diff --git a/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/chat/AiChatMessageServiceImpl.java b/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/chat/AiChatMessageServiceImpl.java
index d11844a4b5..672a3ae0c9 100644
--- a/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/chat/AiChatMessageServiceImpl.java
+++ b/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/chat/AiChatMessageServiceImpl.java
@@ -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;
diff --git a/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/knowledge/AiKnowledgeDocumentService.java b/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/knowledge/AiKnowledgeDocumentService.java
index 3ffd47a510..66155d7727 100644
--- a/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/knowledge/AiKnowledgeDocumentService.java
+++ b/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/knowledge/AiKnowledgeDocumentService.java
@@ -67,13 +67,6 @@ public interface AiKnowledgeDocumentService {
      */
     void updateKnowledgeDocumentStatus(AiKnowledgeDocumentUpdateStatusReqVO reqVO);
 
-    /**
-     * 更新文档检索次数（增加 +1）
-     *
-     * @param ids 文档编号列表
-     */
-    void updateKnowledgeDocumentRetrievalCountIncr(Collection<Long> ids);
-
     /**
      * 删除文档
      *
diff --git a/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/knowledge/AiKnowledgeDocumentServiceImpl.java b/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/knowledge/AiKnowledgeDocumentServiceImpl.java
index c03d35f064..7de51ca0f2 100644
--- a/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/knowledge/AiKnowledgeDocumentServiceImpl.java
+++ b/yudao-module-ai/yudao-module-ai-biz/src/main/java/cn/iocoder/yudao/module/ai/service/knowledge/AiKnowledgeDocumentServiceImpl.java
@@ -161,14 +161,6 @@ public class AiKnowledgeDocumentServiceImpl implements AiKnowledgeDocumentServic
         knowledgeSegmentService.deleteKnowledgeSegmentByDocumentId(id);
     }
 
-    @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);