package tech.aiflowy.ai.service.impl;
|
|
import com.agentsflex.core.llm.embedding.EmbeddingOptions;
|
import org.springframework.core.io.ClassPathResource;
|
import tech.aiflowy.ai.entity.AiDocument;
|
import tech.aiflowy.ai.entity.AiKnowledge;
|
import tech.aiflowy.ai.entity.AiLlm;
|
import tech.aiflowy.ai.mapper.AiDocumentChunkMapper;
|
import tech.aiflowy.ai.mapper.AiDocumentMapper;
|
import tech.aiflowy.ai.service.AiDocumentService;
|
import tech.aiflowy.ai.service.AiKnowledgeService;
|
import tech.aiflowy.ai.service.AiLlmService;
|
import tech.aiflowy.common.ai.MarkdownParser;
|
import com.agentsflex.core.llm.Llm;
|
import com.agentsflex.core.store.DocumentStore;
|
import com.agentsflex.core.store.StoreOptions;
|
import com.agentsflex.core.store.StoreResult;
|
import com.mybatisflex.core.paginate.Page;
|
import com.mybatisflex.core.query.QueryWrapper;
|
import com.mybatisflex.spring.service.impl.ServiceImpl;
|
import org.apache.poi.xwpf.usermodel.*;
|
import org.springframework.beans.factory.annotation.Value;
|
import org.springframework.core.io.UrlResource;
|
import org.springframework.http.MediaType;
|
import org.springframework.http.ResponseEntity;
|
import org.springframework.stereotype.Service;
|
import tech.aiflowy.common.util.StringUtil;
|
|
import javax.annotation.Resource;
|
import java.io.File;
|
import java.io.FileInputStream;
|
import java.io.IOException;
|
import java.io.InputStream;
|
import java.math.BigInteger;
|
import java.nio.file.*;
|
import java.util.Base64;
|
import java.util.List;
|
|
/**
|
* 服务层实现。
|
*
|
* @author michael
|
* @since 2024-08-23
|
*/
|
@Service("AiService")
|
public class AiDocumentServiceImpl extends ServiceImpl<AiDocumentMapper, AiDocument> implements AiDocumentService {
|
|
@javax.annotation.Resource
|
private AiDocumentMapper aiDocumentMapper;
|
|
@javax.annotation.Resource
|
private AiDocumentChunkMapper aiDocumentChunkMapper;
|
|
|
@Value("${aiflowy.storage.local.root}")
|
private String fileUploadPath;
|
|
@Resource
|
private AiKnowledgeService knowledgeService;
|
|
@Resource
|
private AiLlmService aiLlmService;
|
@Override
|
public Page<AiDocument> getDocumentList(String knowledgeId, int pageSize, int pageNum, String fileName) {
|
// 构建 QueryWrapper
|
QueryWrapper queryWrapper = QueryWrapper.create()
|
.select("dt.*", "COUNT(ck.document_id) AS chunk_count") // 选择字段
|
.from("tb_ai_document")
|
.as("dt")// 主表
|
.leftJoin("tb_ai_document_chunk")
|
.as("ck").on("dt.id = ck.document_id") // 左连接
|
.where("dt.knowledge_id = ?", knowledgeId); // 条件 1
|
// 动态添加 fileName 条件
|
if (fileName != null && !fileName.trim().isEmpty()) {
|
queryWrapper.and("dt.title LIKE CONCAT('%', ?, '%')", fileName); // 条件 2
|
}
|
// 分组
|
queryWrapper.groupBy("dt.id");
|
Page<AiDocument> documentVoPage = aiDocumentMapper.paginateAs(pageNum, pageSize, queryWrapper, AiDocument.class);
|
return documentVoPage;
|
}
|
|
/**
|
* 根据文档id删除文件
|
* @param id 文档id
|
* @return
|
*/
|
@Override
|
public boolean removeDoc(String id) {
|
// 查询该文档对应哪些分割的字段，先删除
|
QueryWrapper where = QueryWrapper.create().where("document_id = ? ", id);
|
QueryWrapper aiDocumentWapper = QueryWrapper.create().where("id = ? ", id);
|
AiDocument oneByQuery = aiDocumentMapper.selectOneByQuery(aiDocumentWapper);
|
AiKnowledge knowledge = knowledgeService.getById(oneByQuery.getKnowledgeId());
|
if (knowledge == null) {
|
return false;
|
}
|
|
// 存储到知识库
|
DocumentStore documentStore = knowledge.toDocumentStore();
|
if (documentStore == null) {
|
return false;
|
}
|
|
AiLlm aiLlm = aiLlmService.getById(knowledge.getVectorEmbedLlmId());
|
if (aiLlm == null) {
|
return false;
|
}
|
// 设置向量模型
|
Llm embeddingModel = aiLlm.toLlm();
|
documentStore.setEmbeddingModel(embeddingModel);
|
StoreOptions options = StoreOptions.ofCollectionName(knowledge.getVectorStoreCollection());
|
EmbeddingOptions embeddingOptions = new EmbeddingOptions();
|
embeddingOptions.setModel(aiLlm.getLlmModel());
|
options.setEmbeddingOptions(embeddingOptions);
|
// 查询文本分割表tb_ai_document_chunk中对应的有哪些数据，找出来删除
|
QueryWrapper queryWrapper = QueryWrapper.create()
|
.select("id").from("tb_ai_document_chunk").where("document_id = ?", id);
|
List<BigInteger> chunkIds = aiDocumentChunkMapper.selectListByQueryAs(queryWrapper, BigInteger.class);
|
StoreResult deleteResult = documentStore.delete(chunkIds, options);
|
// if (!deleteResult.isSuccess()){
|
// return false;
|
// }
|
int ck = aiDocumentChunkMapper.deleteByQuery(where);
|
if (ck < 0){
|
return false;
|
}
|
// 再删除指定路径下的文件
|
QueryWrapper wrapper = QueryWrapper.create().where("id = ?", id);
|
AiDocument aiDocument = aiDocumentMapper.selectOneByQuery(wrapper);
|
String filePath = getRootPath() + aiDocument.getDocumentPath();
|
deleteFile(filePath);
|
return true;
|
}
|
|
/**
|
* 文件预览
|
* @param documentId
|
* @return
|
*/
|
@Override
|
public ResponseEntity<?> previewFile(String documentId) throws IOException {
|
QueryWrapper wrapper = QueryWrapper.create().where("id = ?", documentId);
|
AiDocument aiDocument = aiDocumentMapper.selectOneByQuery(wrapper);
|
String FILE_DIRECTORY = getRootPath() + aiDocument.getDocumentPath();
|
// 构建文件路径
|
Path filePath = Paths.get(FILE_DIRECTORY);
|
// Path filePath = Paths.get(FILE_DIRECTORY).resolve(fileName).normalize();
|
org.springframework.core.io.Resource resource = new UrlResource(filePath.toUri());
|
|
// 检查文件是否存在
|
if (!resource.exists()) {
|
return ResponseEntity.status(404).body("文件不存在或无法读取！");
|
}
|
|
// 获取文件扩展名
|
String fileExtension = aiDocument.getDocumentType();
|
|
// 根据文件类型处理
|
switch (fileExtension.toLowerCase()) {
|
case "txt":
|
// 文本文件：直接返回内容
|
String textContent = new String(Files.readAllBytes(filePath));
|
return ResponseEntity.ok()
|
.contentType(MediaType.TEXT_PLAIN)
|
.body(textContent);
|
case "xlsx":
|
// 文本文件：直接返回内容
|
String excelToMarkdownContent = aiDocument.getContent();
|
return ResponseEntity.ok()
|
.contentType(MediaType.TEXT_PLAIN)
|
.body(excelToMarkdownContent);
|
|
case "jpg":
|
case "jpeg":
|
case "png":
|
// 图片文件：返回 Base64 编码
|
byte[] imageBytes = Files.readAllBytes(filePath);
|
String base64Image = Base64.getEncoder().encodeToString(imageBytes);
|
String mimeType = Files.probeContentType(filePath);
|
return ResponseEntity.ok()
|
.contentType(MediaType.parseMediaType(mimeType))
|
.body(base64Image);
|
|
case "pdf":
|
// PDF 文件：返回 Base64 编码
|
// byte[] pdfBytes = Files.readAllBytes(filePath);
|
// String base64Pdf = Base64.getEncoder().encodeToString(pdfBytes);
|
return ResponseEntity.ok()
|
.contentType(MediaType.APPLICATION_PDF)
|
.body(aiDocument.getDocumentPath());
|
|
case "docx":
|
// DOCX 文件：提取文本内容
|
try (FileInputStream fis = new FileInputStream(filePath.toFile());
|
XWPFDocument document = new XWPFDocument(fis)) {
|
|
StringBuilder content = new StringBuilder();
|
|
// 读取段落
|
for (XWPFParagraph paragraph : document.getParagraphs()) {
|
content.append(paragraph.getText()).append("\n");
|
}
|
|
// 读取表格
|
for (XWPFTable table : document.getTables()) {
|
for (XWPFTableRow row : table.getRows()) {
|
for (XWPFTableCell cell : row.getTableCells()) {
|
content.append(cell.getText()).append("\t");
|
}
|
content.append("\n");
|
}
|
}
|
|
return ResponseEntity.ok()
|
.contentType(MediaType.TEXT_PLAIN)
|
.body(content.toString());
|
}
|
case "md":
|
// Markdown 文件：解析为纯文本
|
try (InputStream inputStream = Files.newInputStream(filePath)) {
|
MarkdownParser markdownParser = new MarkdownParser();
|
String plainText = markdownParser.parseToPlainText(inputStream);
|
return ResponseEntity.ok()
|
.contentType(MediaType.TEXT_PLAIN)
|
.body(plainText);
|
}
|
default:
|
// 其他文件类型：提示用户下载
|
return ResponseEntity.status(400).body("不支持预览的文件类型，请下载查看！");
|
}
|
}
|
|
public static boolean deleteFile(String filePath){
|
Path path = Paths.get(filePath);
|
try {
|
// 删除文件
|
Files.delete(path);
|
return true;
|
} catch (NoSuchFileException e) {
|
System.out.println("文件不存在：" + e.getMessage());
|
} catch (DirectoryNotEmptyException e) {
|
System.out.println("目录不为空，无法删除：" + e.getMessage());
|
} catch (IOException e) {
|
System.out.println("发生 I/O 错误：" + e.getMessage());
|
} finally {
|
return false;
|
}
|
}
|
|
private String getRootPath() {
|
if (StringUtil.hasText(this.fileUploadPath)) {
|
return this.fileUploadPath;
|
}
|
ClassPathResource fileResource = new ClassPathResource("/");
|
try {
|
return new File(fileResource.getFile(), "/public").getAbsolutePath();
|
} catch (IOException e) {
|
throw new RuntimeException(e);
|
}
|
}
|
}
|
