218 lines
8.9 KiB
Java
218 lines
8.9 KiB
Java
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;
|
||
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;
|
||
import com.knowledge.base.infrastructure.util.http.FilteredSseOutputAdapter;
|
||
import com.knowledge.base.infrastructure.util.http.WriterAdapter;
|
||
import lombok.RequiredArgsConstructor;
|
||
import lombok.extern.slf4j.Slf4j;
|
||
import org.springframework.stereotype.Service;
|
||
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
|
||
|
||
import java.io.IOException;
|
||
import java.util.*;
|
||
import java.util.regex.Matcher;
|
||
import java.util.regex.Pattern;
|
||
import java.util.stream.Collectors;
|
||
|
||
@Service
|
||
@RequiredArgsConstructor
|
||
@Slf4j
|
||
public class LLMAppServiceImpl implements LLMAppService {
|
||
|
||
private final LLMServiceFactory llmServiceFactory;
|
||
|
||
private final RateLimiterManager rateLimiterManager;
|
||
|
||
private final AnythingLLMService anythingLLMService;
|
||
|
||
private final FileElasticsearchGateway esGateway;
|
||
|
||
private final DynamicConfig dynamicConfig;
|
||
|
||
/**
|
||
* pin和unpin的并发控制
|
||
*/
|
||
private static boolean PROCESSING = false;
|
||
private static String PROCESSING_QUESTION = StrUtil.EMPTY;
|
||
|
||
|
||
@Override
|
||
public String getToken(String password) throws Exception {
|
||
return llmServiceFactory.current().fetchToken(password);
|
||
}
|
||
|
||
@Override
|
||
public SseEmitter ask(String llmToken, String question, Map<String, Object> params) throws Exception {
|
||
SseEmitter emitter = new SseEmitter(300 * 1000L); // 超时时间设为5分钟
|
||
|
||
List<WorkspaceAttachment> finalPinnedDocs = getWorkspaceAttachments(llmToken, question);
|
||
|
||
ThreadPoolUtil.execute(() -> {
|
||
try {
|
||
rateLimiterManager.getRateLimiter(RateLimiterManager.RATE_LIMIT_SCENE_LLM_ASK).acquire();
|
||
WriterAdapter adapter = new FilteredSseOutputAdapter(emitter);
|
||
llmServiceFactory.current().streamAnswer(llmToken, question, params, adapter);
|
||
emitter.complete();
|
||
} catch (Exception e) {
|
||
log.error("LLM调用异常", e);
|
||
try {
|
||
emitter.send(SseEmitter.event().data("{\"error\": \"LLM异常\"}"));
|
||
emitter.completeWithError(e);
|
||
} catch (IOException ioException) {
|
||
log.warn("SSE发送错误信息失败", ioException);
|
||
}
|
||
} finally {
|
||
unpinLlmAttachments(llmToken, question, finalPinnedDocs);
|
||
}
|
||
}, ThreadPoolConfig.SSE_POOL);
|
||
|
||
return emitter;
|
||
}
|
||
|
||
private void unpinLlmAttachments(String llmToken, String question, List<WorkspaceAttachment> 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<WorkspaceAttachment> getWorkspaceAttachments(String llmToken, String question) {
|
||
List<WorkspaceAttachment> finalPinnedDocs;
|
||
if(BooleanUtil.toBoolean(dynamicConfig.getEnableAutoPin())) {
|
||
List<String> keywords = extractKeywords(question);
|
||
log.info("正在回答问题： question: {}, keywords: {}", question, JSONUtil.toJsonStr(keywords));
|
||
List<WorkspaceAttachment> 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<WorkspaceAttachment> pinDocsByKeywords(String llmToken, String workspaceSlug, List<String> keywords, String question) {
|
||
if(CollectionUtil.isEmpty(keywords)) {
|
||
return Collections.emptyList();
|
||
}
|
||
|
||
if(PROCESSING) {
|
||
log.info("有问题正在处理，请稍等，question: ｛｝", PROCESSING_QUESTION);
|
||
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();
|
||
}
|
||
log.info("[pinDocsByKeywords] 已查询到关键词关联的文档。 {}", JSONUtil.toJsonStr(esFilepaths));
|
||
|
||
// 将当前问题标注为处理中
|
||
markQuestionProcessing(question);
|
||
|
||
// 2. 获取当前工作区所有附件（docPath -> url）
|
||
List<WorkspaceAttachment> attachments = anythingLLMService.fetchAttachments(llmToken, workspaceSlug);
|
||
if (attachments == null || attachments.isEmpty()) {
|
||
log.info("[pinDocsByKeywords] 当前工作区无已嵌入附件。");
|
||
return Collections.emptyList();
|
||
}
|
||
|
||
// 3. 找到需要pin的docPath集合
|
||
List<WorkspaceAttachment> toPinDocs = attachments.stream().filter(att -> {
|
||
String url = att.getUrl();
|
||
if (url == null) return false;
|
||
// .md 文件
|
||
if (url.endsWith(".md")) {
|
||
return esFilepaths.stream().anyMatch(f -> url.endsWith(f));
|
||
}
|
||
// 其它文件（如 .pptx、.pdf、.docx等）
|
||
int lastSlash = url.lastIndexOf('/');
|
||
String fileName = lastSlash != -1 ? url.substring(lastSlash + 1) : url;
|
||
// esFilepaths 里也提取末尾文件名
|
||
return esFilepaths.stream().anyMatch(f -> {
|
||
int esLastSlash = f.lastIndexOf('/');
|
||
String esName = esLastSlash != -1 ? f.substring(esLastSlash + 1) : f;
|
||
return fileName.equals(esName);
|
||
});
|
||
}).collect(Collectors.toList());
|
||
|
||
// 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);
|
||
doc.setPinned(true);
|
||
} catch (Exception e) {
|
||
log.warn("pin失败: {}", doc, e);
|
||
}
|
||
});
|
||
|
||
log.info("[pinDocsByKeywords] 共pin住文档{}条: {}", toPinDocs.size(), JSONUtil.toJsonStr(toPinDocs));
|
||
|
||
return toPinDocs;
|
||
}
|
||
|
||
/**
|
||
* 提取所有被LEFT & RIGHT 包裹的子串，返回数组
|
||
*/
|
||
public static List<String> extractKeywords(String text) {
|
||
List<String> keywords = new ArrayList<>();
|
||
String left = ConstantConfig.KEYWORD_PATTERN_LEFT;
|
||
String right = ConstantConfig.KEYWORD_PATTERN_RIGHT;
|
||
Pattern pattern = Pattern.compile(Pattern.quote(left) + "(.*?)" + Pattern.quote(right));
|
||
Matcher matcher = pattern.matcher(text);
|
||
while (matcher.find()) {
|
||
keywords.add(matcher.group(1));
|
||
}
|
||
return keywords;
|
||
}
|
||
|
||
public static void markQuestionProcessing(String question) {
|
||
PROCESSING = true;
|
||
PROCESSING_QUESTION = question;
|
||
}
|
||
|
||
|
||
}
|