lukeye/doc-parser-server
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Luke.Ye 2025-07-08 01:42:07 +08:00
parent da85233841
commit 54a9241dbb
9 changed files with 280 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
src/main/java/com/knowledge/base/application/service/LLMAppService.java
View File

@ -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);
}

74
src/main/java/com/knowledge/base/application/service/LLMAppServiceImpl.java
View File

@ -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());
}
}

10
src/main/java/com/knowledge/base/infrastructure/north/dto/llm/WorkspaceAttachment.java Normal file
View File

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

45
src/main/java/com/knowledge/base/infrastructure/south/es/FileElasticsearchGateway.java
View File

@ -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;

66
src/main/java/com/knowledge/base/infrastructure/south/llm/AnythingLLMService.java Normal file
View File

@ -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);
}

25
src/main/java/com/knowledge/base/infrastructure/south/llm/AnythingLLMServiceImpl.java
View File

@ -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", "");
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<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";

2
src/main/java/com/knowledge/base/infrastructure/south/llm/LLMService.java
View File

@ -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);
}

55
src/test/java/com/knowledge/base/application/service/LLMAppServiceImplTest.java Normal file
View File

@ -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);
}
}
}

5
src/test/java/com/knowledge/base/infrastructure/south/llm/AnythingLLMServiceTest.java
View File

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