diff --git a/src/main/java/com/knowledge/base/application/service/LLMAppServiceImpl.java b/src/main/java/com/knowledge/base/application/service/LLMAppServiceImpl.java index fabb15c..e07b820 100644 --- a/src/main/java/com/knowledge/base/application/service/LLMAppServiceImpl.java +++ b/src/main/java/com/knowledge/base/application/service/LLMAppServiceImpl.java @@ -1,11 +1,13 @@ package com.knowledge.base.application.service; import cn.hutool.core.collection.CollectionUtil; +import cn.hutool.core.util.BooleanUtil; import cn.hutool.core.util.StrUtil; import cn.hutool.json.JSONUtil; import com.google.common.collect.Lists; import com.knowledge.base.domain.doc.model.FileEsModel; import com.knowledge.base.infrastructure.config.ConstantConfig; +import com.knowledge.base.infrastructure.config.DynamicConfig; import com.knowledge.base.infrastructure.config.ThreadPoolConfig; import com.knowledge.base.infrastructure.north.dto.llm.WorkspaceAttachment; import com.knowledge.base.infrastructure.south.es.FileElasticsearchGateway; @@ -39,6 +41,8 @@ public class LLMAppServiceImpl implements LLMAppService { private final FileElasticsearchGateway esGateway; + private final DynamicConfig dynamicConfig; + /** * pin和unpin的并发控制 */ @@ -54,9 +58,8 @@ public class LLMAppServiceImpl implements LLMAppService { @Override public SseEmitter ask(String llmToken, String question, Map params) throws Exception { SseEmitter emitter = new SseEmitter(300 * 1000L); // 超时时间设为5分钟 - List keywords = extractKeywords(question); - log.info("正在回答问题: question: {}, keywords: {}", question, JSONUtil.toJsonStr(keywords)); - List pinnedDocs = pinDocsByKeywords(llmToken, ConstantConfig.DEFAULT_SLUG_ID, keywords, question); + + List finalPinnedDocs = getWorkspaceAttachments(llmToken, question); ThreadPoolUtil.execute(() -> { try { @@ -73,22 +76,44 @@ public class LLMAppServiceImpl implements LLMAppService { log.warn("SSE发送错误信息失败", ioException); } } finally { - // 问题结束后,需要unpin掉 - log.info("问题已回答完成: question: {}, pinnedDocs: {}", question, JSONUtil.toJsonStr(pinnedDocs)); - pinnedDocs.forEach(doc -> { - try { - anythingLLMService.updatePin(llmToken, ConstantConfig.DEFAULT_SLUG_ID, doc.getDocpath(), false); - } catch (Exception e) { - log.warn("unpin失败: {}", doc, e); - } - }); - PROCESSING = false; + unpinLlmAttachments(llmToken, question, finalPinnedDocs); } }, ThreadPoolConfig.SSE_POOL); return emitter; } + private void unpinLlmAttachments(String llmToken, String question, List finalPinnedDocs) { + if(BooleanUtil.toBoolean(dynamicConfig.getEnableAutoPin())) { + // 问题结束后,需要unpin掉 + log.info("问题已回答完成: question: {}, pinnedDocs: {}", question, JSONUtil.toJsonStr(finalPinnedDocs)); + finalPinnedDocs.forEach(doc -> { + try { + anythingLLMService.updatePin(llmToken, ConstantConfig.DEFAULT_SLUG_ID, doc.getDocpath(), false); + } catch (Exception e) { + log.warn("unpin失败: {}", doc, e); + } + }); + PROCESSING = false; + }else { + log.info("[unpinLlmAttachments] 自动pin文档功能未开启. {}", dynamicConfig.getEnableAutoPin()); + } + } + + private List getWorkspaceAttachments(String llmToken, String question) { + List finalPinnedDocs; + if(BooleanUtil.toBoolean(dynamicConfig.getEnableAutoPin())) { + List keywords = extractKeywords(question); + log.info("正在回答问题: question: {}, keywords: {}", question, JSONUtil.toJsonStr(keywords)); + List pinnedDocs = pinDocsByKeywords(llmToken, ConstantConfig.DEFAULT_SLUG_ID, keywords, question); + finalPinnedDocs = pinnedDocs; + } else { + finalPinnedDocs = Collections.emptyList(); + log.info("[getWorkspaceAttachments] 自动pin文档功能未开启. {}", dynamicConfig.getEnableAutoPin()); + } + return finalPinnedDocs; + } + @Override public List pinDocsByKeywords(String llmToken, String workspaceSlug, List keywords, String question) { if(CollectionUtil.isEmpty(keywords)) { diff --git a/src/main/java/com/knowledge/base/infrastructure/config/DynamicConfig.java b/src/main/java/com/knowledge/base/infrastructure/config/DynamicConfig.java index fdafd1a..ef8ccf2 100644 --- a/src/main/java/com/knowledge/base/infrastructure/config/DynamicConfig.java +++ b/src/main/java/com/knowledge/base/infrastructure/config/DynamicConfig.java @@ -36,6 +36,9 @@ public class DynamicConfig { @Value("${llm.sse.rate.limit: 3}") private String llmSseRateLimit; + @Value("${llm.enable.auto.pin: false}") + private String enableAutoPin; + @Value("${markdown.path}") private String mdDirectoryPath; diff --git a/src/main/java/com/knowledge/base/infrastructure/south/es/FileElasticsearchGateway.java b/src/main/java/com/knowledge/base/infrastructure/south/es/FileElasticsearchGateway.java index ef3de00..5406dc1 100644 --- a/src/main/java/com/knowledge/base/infrastructure/south/es/FileElasticsearchGateway.java +++ b/src/main/java/com/knowledge/base/infrastructure/south/es/FileElasticsearchGateway.java @@ -66,7 +66,7 @@ public class FileElasticsearchGateway { SearchSourceBuilder builder = new SearchSourceBuilder() .query(query) .from(0) - .size(6) + .size(3) .sort("_score", SortOrder.DESC); SearchRequest request = new SearchRequest(INDEX_NAME).source(builder); diff --git a/src/test/java/com/knowledge/base/application/service/LLMAppServiceImplTest.java b/src/test/java/com/knowledge/base/application/service/LLMAppServiceImplTest.java index 05fded0..d860ab2 100644 --- a/src/test/java/com/knowledge/base/application/service/LLMAppServiceImplTest.java +++ b/src/test/java/com/knowledge/base/application/service/LLMAppServiceImplTest.java @@ -49,6 +49,14 @@ public class LLMAppServiceImplTest { try { List result = llmAppService.pinDocsByKeywords(token, getSlug(), Lists.newArrayList("商业"), "[[商业]] 的本质是什么?"); log.info("result: {}", JSONUtil.toJsonStr(result)); + + result.forEach(doc -> { + try { + llmService.updatePin(token, ConstantConfig.DEFAULT_SLUG_ID, doc.getDocpath(), false); + } catch (Exception e) { + log.warn("unpin失败: {}", doc, e); + } + }); } catch (Exception e) { log.warn("测试 pinDocsByKeywords 失败!: {}", e.getMessage(), e); }