BJ_ZTEITS-Developers / parking-cloud-yudao

Commit 9a3bfe89e51de755b4168f945866398341d93e7b

Authored by YunaiV
1 parent be56c834

feat：【ai】226 修改知识库和联网搜索仅初始化一次

Inline Side-by-side
Showing 1 changed file with 18 additions and 12 deletions
yudao-module-ai/src/main/java/cn/iocoder/yudao/module/ai/service/chat/AiChatMessageServiceImpl.java
  View file @9a3bfe8
... ... @@ -59,6 +59,8 @@ import reactor.core.publisher.Flux;
59 59  
60 60 import java.time.LocalDateTime;
61 61 import java.util.*;
  62 +import java.util.concurrent.atomic.AtomicBoolean;
  63 +import java.util.concurrent.atomic.AtomicReference;
62 64 import java.util.stream.Collectors;
63 65  
64 66 import static cn.iocoder.yudao.framework.common.exception.util.ServiceExceptionUtil.exception;
... ... @@ -231,20 +233,24 @@ public class AiChatMessageServiceImpl implements AiChatMessageService {
231 233 // 4.3 流式返回
232 234 StringBuffer contentBuffer = new StringBuffer();
233 235 StringBuffer reasoningContentBuffer = new StringBuffer();
  236 +
  237 + // 防止执行多次知识库和联网搜索
  238 + AtomicBoolean firstExecuteFlag = new AtomicBoolean(true);
  239 + AtomicReference<List<AiChatMessageRespVO.KnowledgeSegment>> cacheSegments = new AtomicReference<>();
  240 + AtomicReference<List<AiWebSearchResponse.WebPage>> cacheWebSearchPages = new AtomicReference<>();
234 241 return streamResponse.map(chunk -> {
235 242 // 仅首次：返回知识库、联网搜索
236   - List<AiChatMessageRespVO.KnowledgeSegment> segments = null;
237   - List<AiWebSearchResponse.WebPage> webSearchPages = null;
238 243 if (StrUtil.isEmpty(contentBuffer)) {
239   - Map<Long, AiKnowledgeDocumentDO> documentMap = TenantUtils.executeIgnore(() ->
240   - knowledgeDocumentService.getKnowledgeDocumentMap(
241   - convertSet(knowledgeSegments, AiKnowledgeSegmentSearchRespBO::getDocumentId)));
242   - segments = BeanUtils.toBean(knowledgeSegments, AiChatMessageRespVO.KnowledgeSegment.class, segment -> {
243   - AiKnowledgeDocumentDO document = documentMap.get(segment.getDocumentId());
244   - segment.setDocumentName(document != null ? document.getName() : null);
245   - });
246   - if (webSearchResponse != null) {
247   - webSearchPages = webSearchResponse.getLists();
  244 + if (firstExecuteFlag.compareAndSet(true, false)) { // CAS 操作，确保仅执行一次
  245 + Map<Long, AiKnowledgeDocumentDO> documentMap = TenantUtils.executeIgnore(() -> knowledgeDocumentService.getKnowledgeDocumentMap(
  246 + convertSet(knowledgeSegments, AiKnowledgeSegmentSearchRespBO::getDocumentId)));
  247 + cacheSegments.set(BeanUtils.toBean(knowledgeSegments, AiChatMessageRespVO.KnowledgeSegment.class, segment -> {
  248 + AiKnowledgeDocumentDO document = documentMap.get(segment.getDocumentId());
  249 + segment.setDocumentName(document != null ? document.getName() : null);
  250 + }));
  251 + if (webSearchResponse != null) {
  252 + cacheWebSearchPages.set(webSearchResponse.getLists());
  253 + }
248 254 }
249 255 }
250 256 // 响应结果
... ... @@ -261,7 +267,7 @@ public class AiChatMessageServiceImpl implements AiChatMessageService {
261 267 .setReceive(BeanUtils.toBean(assistantMessage, AiChatMessageSendRespVO.Message.class)
262 268 .setContent(StrUtil.nullToDefault(newContent, "")) // 避免 null 的 情况
263 269 .setReasoningContent(StrUtil.nullToDefault(newReasoningContent, "")) // 避免 null 的 情况
264   - .setSegments(segments).setWebSearchPages(webSearchPages))); // 知识库 + 联网搜索
  270 + .setSegments(cacheSegments.get()).setWebSearchPages(cacheWebSearchPages.get()))); // 知识库 + 联网搜索
265 271 }).doOnComplete(() -> {
266 272 // 忽略租户，因为 Flux 异步无法透传租户
267 273 TenantUtils.executeIgnore(() -> chatMessageMapper.updateById(
... ...