llm逻辑优化，pin住相关文档

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<String, Object> params) throws Exception;
+
+    /**
+     * 按关键词pin文档（其它文档取消pin）
+     *
+     * @param llmToken       LLM鉴权token
+     * @param workspaceSlug  当前工作区slug
+     * @param keywords       搜索关键词
+     * @return  本次被pin的docPath集合
+     */
+    List<String> pinDocsByKeywords(String llmToken, String workspaceSlug, List<String> keywords);
 }
 
 

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<String, Object> 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<String> pinDocsByKeywords(String llmToken, String workspaceSlug, List<String> keywords) {
+        if(CollectionUtil.isEmpty(keywords)) {
+            return Collections.emptyList();
+        }
+
+        List<FileEsModel> fileEsModels = esGateway.searchByKeywords(keywords);
+        Set<String> 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<WorkspaceAttachment> attachments = anythingLLMService.fetchAttachments(llmToken, workspaceSlug);
+        if (attachments == null || attachments.isEmpty()) {
+            log.info("[pinDocsByKeywords] 当前工作区无已嵌入附件。");
+            return Collections.emptyList();
+        }
+
+        // 3. 找到需要pin的docPath集合
+        Set<WorkspaceAttachment> 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());
+    }
 }

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;
+}

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<FileEsModel> searchByKeywords(List<String> 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<FileEsModel> 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<List<String>> keywordGroups, int page, int size) {
         try {
             int from = (page - 1) * size;

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<String, Object> importObsidianToLLM(String llmToken, File vaultDir);
+
+    /**
+     * 从本地文件中移除文件夹
+     * @param llmToken
+     * @param folderName
+     * @return
+     */
+    String removeFolderFromLLM(String llmToken, String folderName);
+
+    /**
+     * 获取本地文件列表
+     * @param llmToken
+     * @param workspaceSlug
+     * @param keyword
+     * @return
+     */
+    Map<String, Object> getLocalFileItems(String llmToken, String workspaceSlug, String keyword);
+
+    /**
+     * 将本地文件嵌入工作区
+     * @param llmToken
+     * @param workspaceSlug
+     * @param adds
+     * @param deletes
+     * @return
+     */
+    Map<String, Object> updateEmbeddings(String llmToken, String workspaceSlug, List<String> adds, List<String> deletes);
+
+
+    /**
+     * 获取工作区下所有的附件
+     * @param llmToken
+     * @param workspaceSlug
+     * @return
+     */
+    List<WorkspaceAttachment> 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);
+}

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<String, Object> 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<String, String> fetchAttachments(String llmToken, String workspaceSlug) {
+    public List<WorkspaceAttachment> fetchAttachments(String llmToken, String workspaceSlug) {
         String url = baseUrl + "/api/workspace/" + workspaceSlug;
         Map<String, String> headers = new HashMap<>();
         headers.put("Authorization", "Bearer " + llmToken);
@@ -190,7 +191,7 @@ public class AnythingLLMServiceImpl implements LLMService {
                     new TypeReference<Map<String, Object>>() {}
             );
 
-            Map<String, String> result = new HashMap<>();
+            List<WorkspaceAttachment> 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<String, Object> metadata = new ObjectMapper().readValue(metadataStr, new TypeReference<Map<String, Object>>() {});
-                                    String urlVal = (String) metadata.getOrDefault("url", "");
+                                    urlVal = (String) metadata.getOrDefault("url", "");
-                                    result.put(docpath, urlVal);
                                 } 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<String, String> headers = new HashMap<>();
@@ -242,6 +250,7 @@ public class AnythingLLMServiceImpl implements LLMService {
         }
     }
 
+    @Override
     public Map<String, Object> 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<String, Object> updateEmbeddings(String llmToken, String workspaceSlug, List<String> adds, List<String> deletes) {
         String url = baseUrl + "/api/workspace/" + workspaceSlug + "/update-embeddings";
 

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<String, Object> params, WriterAdapter writer) throws Exception;
 
     boolean supports(String type);
-
-    Map<String, String> fetchAttachments(String llmToken, String workspaceSlug);
 }
 

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<String, Object> 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<String> result = llmAppService.pinDocsByKeywords(token, getSlug(), Lists.newArrayList("商业"));
+            log.info("result: {}", JSONUtil.toJsonStr(result));
+        } catch (Exception e) {
+            log.warn("测试 pinDocsByKeywords 失败！: {}", e.getMessage(), e);
+        }
+    }
+}

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<String, String> attachments = llmService.fetchAttachments(token, slug);
+            List<WorkspaceAttachment> attachments = llmService.fetchAttachments(token, slug);
-            attachments.forEach((k, v) -> log.info("Attachment: {} -> {}", k, v));
+            attachments.forEach(e -> log.info("Attachment: {} -> {} -> {} ", e.getDocpath(), e.getUrl(), e.isPinned()));
         } catch (Exception e) {
             log.warn("测试 fetchAttachments 异常", e);
         }