diff --git a/src/main/java/com/knowledge/base/application/service/LLMAppService.java b/src/main/java/com/knowledge/base/application/service/LLMAppService.java index 55dd490..37efc5d 100644 --- a/src/main/java/com/knowledge/base/application/service/LLMAppService.java +++ b/src/main/java/com/knowledge/base/application/service/LLMAppService.java @@ -2,6 +2,7 @@ package com.knowledge.base.application.service; import org.springframework.web.servlet.mvc.method.annotation.SseEmitter; +import java.util.List; import java.util.Map; public interface LLMAppService { @@ -22,6 +23,16 @@ public interface LLMAppService { * @throws Exception */ SseEmitter ask(String llmToken, String question, Map params) throws Exception; + + /** + * 按关键词pin文档（其它文档取消pin） + * + * @param llmToken LLM鉴权token + * @param workspaceSlug 当前工作区slug + * @param keywords 搜索关键词 + * @return 本次被pin的docPath集合 + */ + List pinDocsByKeywords(String llmToken, String workspaceSlug, List keywords); } 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 4807359..957b1e1 100644 --- a/src/main/java/com/knowledge/base/application/service/LLMAppServiceImpl.java +++ b/src/main/java/com/knowledge/base/application/service/LLMAppServiceImpl.java @@ -1,6 +1,15 @@ package com.knowledge.base.application.service; +import cn.hutool.core.collection.CollectionUtil; +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.ThreadPoolConfig; +import com.knowledge.base.infrastructure.north.dto.llm.WorkspaceAttachment; +import com.knowledge.base.infrastructure.south.es.FileElasticsearchGateway; +import com.knowledge.base.infrastructure.south.llm.AnythingLLMService; import com.knowledge.base.infrastructure.south.llm.LLMServiceFactory; import com.knowledge.base.infrastructure.util.RateLimiterManager; import com.knowledge.base.infrastructure.util.ThreadPoolUtil; @@ -12,7 +21,8 @@ import org.springframework.stereotype.Service; import org.springframework.web.servlet.mvc.method.annotation.SseEmitter; import java.io.IOException; -import java.util.Map; +import java.util.*; +import java.util.stream.Collectors; @Service @RequiredArgsConstructor @@ -23,6 +33,11 @@ public class LLMAppServiceImpl implements LLMAppService { private final RateLimiterManager rateLimiterManager; + private final AnythingLLMService anythingLLMService; + + private final FileElasticsearchGateway esGateway; + + @Override public String getToken(String password) throws Exception { return llmServiceFactory.current().fetchToken(password); @@ -31,6 +46,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分钟 + String keyword = (String) params.get("keyword"); + pinDocsByKeywords(llmToken, ConstantConfig.DEFAULT_SLUG_ID, Lists.newArrayList(keyword)); ThreadPoolUtil.execute(() -> { try { @@ -51,4 +68,59 @@ public class LLMAppServiceImpl implements LLMAppService { return emitter; } + + @Override + public List pinDocsByKeywords(String llmToken, String workspaceSlug, List keywords) { + if(CollectionUtil.isEmpty(keywords)) { + return Collections.emptyList(); + } + + List fileEsModels = esGateway.searchByKeywords(keywords); + Set esFilepaths = fileEsModels.stream() + .map(FileEsModel::getFilepath) + .filter(Objects::nonNull) + .collect(Collectors.toSet()); + + if (esFilepaths.isEmpty()) { + log.info("[pinDocsByKeywords] 未查到匹配ES文档，跳过pin操作。"); + return Lists.newArrayList(); + } + + // 2. 获取当前工作区所有附件（docPath -> url） + List attachments = anythingLLMService.fetchAttachments(llmToken, workspaceSlug); + if (attachments == null || attachments.isEmpty()) { + log.info("[pinDocsByKeywords] 当前工作区无已嵌入附件。"); + return Collections.emptyList(); + } + + // 3. 找到需要pin的docPath集合 + Set toPinDocs = attachments.stream() + .filter(att -> esFilepaths.stream().anyMatch(f -> att.getUrl().endsWith(f))) + .collect(Collectors.toSet()); + + // 4. 只unpin之前已pin的文档 + attachments.stream() + .filter(WorkspaceAttachment::isPinned) + .forEach(att -> { + try { + log.info("unpin doc {}", att.getUrl()); + anythingLLMService.updatePin(llmToken, workspaceSlug, att.getDocpath(), false); + } catch (Exception e) { + log.warn("取消pin失败: {}", att.getDocpath(), e); + } + }); + + // 5. pin目标文档 + toPinDocs.forEach(doc -> { + try { + anythingLLMService.updatePin(llmToken, workspaceSlug, doc.getDocpath(), true); + } catch (Exception e) { + log.warn("pin失败: {}", doc, e); + } + }); + + log.info("[pinDocsByKeywords] 共pin住文档{}条: {}", toPinDocs.size(), JSONUtil.toJsonStr(toPinDocs)); + + return toPinDocs.stream().map(WorkspaceAttachment::getUrl).collect(Collectors.toList()); + } } diff --git a/src/main/java/com/knowledge/base/infrastructure/north/dto/llm/WorkspaceAttachment.java b/src/main/java/com/knowledge/base/infrastructure/north/dto/llm/WorkspaceAttachment.java new file mode 100644 index 0000000..9aae841 --- /dev/null +++ b/src/main/java/com/knowledge/base/infrastructure/north/dto/llm/WorkspaceAttachment.java @@ -0,0 +1,10 @@ +package com.knowledge.base.infrastructure.north.dto.llm; + +import lombok.Data; + +@Data +public class WorkspaceAttachment { + private String docpath; + private String url; + private boolean pinned; +} 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 30f1a4d..5f523d9 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 @@ -1,6 +1,9 @@ package com.knowledge.base.infrastructure.south.es; import cn.hutool.core.bean.BeanUtil; +import cn.hutool.core.bean.copier.CopyOptions; +import cn.hutool.core.collection.CollectionUtil; +import com.google.common.collect.Lists; import com.knowledge.base.domain.doc.model.FileEsField; import com.knowledge.base.domain.doc.model.FileEsModel; import com.knowledge.base.infrastructure.util.SafeIdUtil; @@ -32,6 +35,7 @@ import org.springframework.stereotype.Component; import java.io.IOException; import java.util.List; import java.util.Map; +import java.util.stream.Collectors; /** * @author Luke.ye @@ -46,6 +50,47 @@ public class FileElasticsearchGateway { @Autowired private RestHighLevelClient esClient; + + public List searchByKeywords(List keywords) { + if (CollectionUtil.isEmpty(keywords)) { + return Lists.newArrayList(); + } + + BoolQueryBuilder boolQuery = QueryBuilders.boolQuery(); + for (String kw : keywords) { + boolQuery.must(QueryBuilders.multiMatchQuery(kw) + .field("filename", 10.0f) + .field("content", 1.0f) + .type(org.elasticsearch.index.query.MultiMatchQueryBuilder.Type.BEST_FIELDS)); + } + boolQuery.filter(QueryBuilders.rangeQuery(FileEsField.EXPIRE_TIME) + .gte(System.currentTimeMillis()) + .timeZone("+08:00")); + + SearchSourceBuilder builder = new SearchSourceBuilder() + .query(boolQuery) + .from(0) + .size(6) + .sort("_score", SortOrder.DESC); + + SearchRequest request = new SearchRequest(INDEX_NAME).source(builder); + SearchResponse response; + try { + response = esClient.search(request, RequestOptions.DEFAULT); + } catch (IOException e) { + log.error("[searchByKeywords] error.", e); + throw new RuntimeException(e); + } + + List result = Lists.newArrayList(); + for (SearchHit hit : response.getHits().getHits()) { + FileEsModel fileEsModel = BeanUtil.mapToBean(hit.getSourceAsMap(), FileEsModel.class, true, CopyOptions.create()); + result.add(fileEsModel); + } + return result; + } + + public SearchResponse searchFileDocuments(List> keywordGroups, int page, int size) { try { int from = (page - 1) * size; diff --git a/src/main/java/com/knowledge/base/infrastructure/south/llm/AnythingLLMService.java b/src/main/java/com/knowledge/base/infrastructure/south/llm/AnythingLLMService.java new file mode 100644 index 0000000..f5522cf --- /dev/null +++ b/src/main/java/com/knowledge/base/infrastructure/south/llm/AnythingLLMService.java @@ -0,0 +1,66 @@ +package com.knowledge.base.infrastructure.south.llm; + +import com.knowledge.base.infrastructure.north.dto.llm.WorkspaceAttachment; + +import java.io.File; +import java.util.List; +import java.util.Map; + +/** + * 具体的使用顺序可参考 Readme.md + */ +public interface AnythingLLMService extends LLMService{ + /** + * 将OB仓库中的文档导入LLM本地文件 + * @param llmToken + * @param vaultDir + * @return + */ + Map importObsidianToLLM(String llmToken, File vaultDir); + + /** + * 从本地文件中移除文件夹 + * @param llmToken + * @param folderName + * @return + */ + String removeFolderFromLLM(String llmToken, String folderName); + + /** + * 获取本地文件列表 + * @param llmToken + * @param workspaceSlug + * @param keyword + * @return + */ + Map getLocalFileItems(String llmToken, String workspaceSlug, String keyword); + + /** + * 将本地文件嵌入工作区 + * @param llmToken + * @param workspaceSlug + * @param adds + * @param deletes + * @return + */ + Map updateEmbeddings(String llmToken, String workspaceSlug, List adds, List deletes); + + + /** + * 获取工作区下所有的附件 + * @param llmToken + * @param workspaceSlug + * @return + */ + List fetchAttachments(String llmToken, String workspaceSlug); + + /** + * 更新工作区中文件的pin状态 + * @param llmToken + * @param workspaceSlug + * @param docPath + * @param pinStatus + * @return + */ + String updatePin(String llmToken, String workspaceSlug, String docPath, boolean pinStatus); +} diff --git a/src/main/java/com/knowledge/base/infrastructure/south/llm/AnythingLLMServiceImpl.java b/src/main/java/com/knowledge/base/infrastructure/south/llm/AnythingLLMServiceImpl.java index 9377f66..ba703a1 100644 --- a/src/main/java/com/knowledge/base/infrastructure/south/llm/AnythingLLMServiceImpl.java +++ b/src/main/java/com/knowledge/base/infrastructure/south/llm/AnythingLLMServiceImpl.java @@ -8,6 +8,7 @@ import com.fasterxml.jackson.databind.ObjectMapper; import com.google.common.collect.Maps; import com.knowledge.base.infrastructure.cache.iface.UserCacheService; import com.knowledge.base.infrastructure.config.ConstantConfig; +import com.knowledge.base.infrastructure.north.dto.llm.WorkspaceAttachment; import com.knowledge.base.infrastructure.util.HttpHelper; import com.knowledge.base.infrastructure.util.http.WriterAdapter; import lombok.RequiredArgsConstructor; @@ -30,7 +31,7 @@ import java.util.stream.Stream; @Slf4j @Service @RequiredArgsConstructor -public class AnythingLLMServiceImpl implements LLMService { +public class AnythingLLMServiceImpl implements AnythingLLMService { @Value("${llm.remote.base-url}") private String baseUrl; @@ -98,6 +99,7 @@ public class AnythingLLMServiceImpl implements LLMService { } + @Override public Map importObsidianToLLM(String llmToken, File vaultDir) { if (!vaultDir.exists() || !vaultDir.isDirectory()) { throw new IllegalArgumentException("无效的 Obsidian vault 目录: " + vaultDir.getAbsolutePath()); @@ -151,6 +153,7 @@ public class AnythingLLMServiceImpl implements LLMService { } } + @Override public String removeFolderFromLLM(String llmToken, String folderName) { String url = baseUrl + "/api/system/remove-folder"; @@ -173,10 +176,8 @@ public class AnythingLLMServiceImpl implements LLMService { } } - - @Override - public Map fetchAttachments(String llmToken, String workspaceSlug) { + public List fetchAttachments(String llmToken, String workspaceSlug) { String url = baseUrl + "/api/workspace/" + workspaceSlug; Map headers = new HashMap<>(); headers.put("Authorization", "Bearer " + llmToken); @@ -190,7 +191,7 @@ public class AnythingLLMServiceImpl implements LLMService { new TypeReference>() {} ); - Map result = new HashMap<>(); + List result = new ArrayList<>(); if (response.containsKey("workspace")) { Map workspace = (Map) response.get("workspace"); Object documents = workspace.get("documents"); @@ -200,15 +201,21 @@ public class AnythingLLMServiceImpl implements LLMService { Map doc = (Map) docObj; String docpath = (String) doc.get("docpath"); String metadataStr = (String) doc.get("metadata"); + boolean pinned = Boolean.TRUE.equals(doc.get("pinned")); + String urlVal = ""; if (docpath != null && metadataStr != null) { try { Map metadata = new ObjectMapper().readValue(metadataStr, new TypeReference>() {}); - String urlVal = (String) metadata.getOrDefault("url", ""); - result.put(docpath, urlVal); + urlVal = (String) metadata.getOrDefault("url", ""); } catch (Exception e) { log.warn("解析 metadata 失败: {}", metadataStr, e); } } + WorkspaceAttachment attachment = new WorkspaceAttachment(); + attachment.setDocpath(docpath); + attachment.setUrl(urlVal); + attachment.setPinned(pinned); + result.add(attachment); } } } @@ -220,6 +227,7 @@ public class AnythingLLMServiceImpl implements LLMService { } + @Override public String updatePin(String llmToken, String workspaceSlug, String docPath, boolean pinStatus) { String url = baseUrl + "/api/workspace/" + workspaceSlug + "/update-pin"; Map headers = new HashMap<>(); @@ -242,6 +250,7 @@ public class AnythingLLMServiceImpl implements LLMService { } } + @Override public Map getLocalFileItems(String llmToken, String workspaceSlug, String keyword) { String url = baseUrl + "/api/system/local-files"; @@ -279,7 +288,7 @@ public class AnythingLLMServiceImpl implements LLMService { } } - + @Override public Map updateEmbeddings(String llmToken, String workspaceSlug, List adds, List deletes) { String url = baseUrl + "/api/workspace/" + workspaceSlug + "/update-embeddings"; diff --git a/src/main/java/com/knowledge/base/infrastructure/south/llm/LLMService.java b/src/main/java/com/knowledge/base/infrastructure/south/llm/LLMService.java index 8706f71..49e19cf 100644 --- a/src/main/java/com/knowledge/base/infrastructure/south/llm/LLMService.java +++ b/src/main/java/com/knowledge/base/infrastructure/south/llm/LLMService.java @@ -10,7 +10,5 @@ public interface LLMService { void streamAnswer(String llmToken, String question, Map params, WriterAdapter writer) throws Exception; boolean supports(String type); - - Map fetchAttachments(String llmToken, String workspaceSlug); } diff --git a/src/test/java/com/knowledge/base/application/service/LLMAppServiceImplTest.java b/src/test/java/com/knowledge/base/application/service/LLMAppServiceImplTest.java new file mode 100644 index 0000000..e0f920a --- /dev/null +++ b/src/test/java/com/knowledge/base/application/service/LLMAppServiceImplTest.java @@ -0,0 +1,55 @@ +package com.knowledge.base.application.service; + +import cn.hutool.json.JSONUtil; +import com.google.common.collect.Lists; +import com.google.common.collect.Maps; +import com.knowledge.base.infrastructure.config.ConstantConfig; +import com.knowledge.base.infrastructure.south.llm.AnythingLLMServiceImpl; +import lombok.extern.slf4j.Slf4j; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.boot.test.context.SpringBootTest; + +import java.util.List; +import java.util.Map; + +@SpringBootTest(properties = { + "spring.profiles.active=dev-windows" +}) +@Slf4j +public class LLMAppServiceImplTest { + + @Autowired + private AnythingLLMServiceImpl llmService; + + @Autowired + private LLMAppService llmAppService; + + private String token; + + private Map llmRepo = Maps.newHashMap(); + + public static String getSlug() { + return ConstantConfig.DEFAULT_SLUG_ID; + } + + @BeforeEach + public void setUp() { + try { + this.token = llmService.fetchToken("lukeye@6"); + } catch (Exception e) { + log.warn("获取 token 失败: {}", e.getMessage(), e); + } + } + + @Test + public void testPinDocsByKeywords() throws Exception { + try { + List result = llmAppService.pinDocsByKeywords(token, getSlug(), Lists.newArrayList("商业")); + log.info("result: {}", JSONUtil.toJsonStr(result)); + } catch (Exception e) { + log.warn("测试 pinDocsByKeywords 失败！: {}", e.getMessage(), e); + } + } +} diff --git a/src/test/java/com/knowledge/base/infrastructure/south/llm/AnythingLLMServiceTest.java b/src/test/java/com/knowledge/base/infrastructure/south/llm/AnythingLLMServiceTest.java index d878700..66131fd 100644 --- a/src/test/java/com/knowledge/base/infrastructure/south/llm/AnythingLLMServiceTest.java +++ b/src/test/java/com/knowledge/base/infrastructure/south/llm/AnythingLLMServiceTest.java @@ -6,6 +6,7 @@ import cn.hutool.json.JSONUtil; import com.google.common.collect.Lists; import com.google.common.collect.Maps; import com.knowledge.base.infrastructure.config.ConstantConfig; +import com.knowledge.base.infrastructure.north.dto.llm.WorkspaceAttachment; import com.knowledge.base.infrastructure.util.BeanTool; import com.knowledge.base.infrastructure.util.http.WriterAdapter; import lombok.extern.slf4j.Slf4j; @@ -133,8 +134,8 @@ public class AnythingLLMServiceTest { public void testFetchAttachments() { try { String slug = getSlug(); - Map attachments = llmService.fetchAttachments(token, slug); - attachments.forEach((k, v) -> log.info("Attachment: {} -> {}", k, v)); + List attachments = llmService.fetchAttachments(token, slug); + attachments.forEach(e -> log.info("Attachment: {} -> {} -> {} ", e.getDocpath(), e.getUrl(), e.isPinned())); } catch (Exception e) { log.warn("测试 fetchAttachments 异常", e); }