package com.ard.agent.service;
|
|
import lombok.extern.slf4j.Slf4j;
|
import org.springframework.ai.document.Document;
|
import org.springframework.ai.reader.tika.TikaDocumentReader;
|
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
|
import org.springframework.ai.vectorstore.VectorStore;
|
import org.springframework.core.io.FileSystemResource;
|
import org.springframework.stereotype.Service;
|
import org.springframework.web.multipart.MultipartFile;
|
|
import java.io.IOException;
|
import java.nio.file.Files;
|
import java.nio.file.Path;
|
import java.time.LocalDateTime;
|
import java.util.Arrays;
|
import java.util.HashMap;
|
import java.util.List;
|
import java.util.Map;
|
|
@Slf4j
|
@Service
|
public class KnowledgeBaseService {
|
|
private final VectorStore vectorStore;
|
private final TokenTextSplitter textSplitter;
|
|
public KnowledgeBaseService(VectorStore vectorStore) {
|
this.vectorStore = vectorStore;
|
this.textSplitter = new TokenTextSplitter(800, 150, 5, 10000, true);
|
}
|
|
public Map<String, Object> uploadDocument(MultipartFile file) {
|
Map<String, Object> result = new HashMap<>();
|
String fileName = file.getOriginalFilename();
|
result.put("fileName", fileName);
|
|
Path tempFile = null;
|
try {
|
log.info("开始处理文档: {}", fileName);
|
|
String contentType = file.getContentType();
|
if (!isSupportedContentType(contentType)) {
|
result.put("success", false);
|
result.put("message", "不支持的文件类型: " + contentType);
|
return result;
|
}
|
|
// 创建临时文件
|
tempFile = Files.createTempFile("upload_", "_" + fileName);
|
file.transferTo(tempFile.toFile());
|
log.info("临时文件创建成功: {}", tempFile.toString());
|
|
// 关键修改：使用 FileSystemResource 而不是字符串路径
|
FileSystemResource resource = new FileSystemResource(tempFile.toFile());
|
TikaDocumentReader reader = new TikaDocumentReader(resource);
|
List<Document> documents = reader.get();
|
|
if (documents == null || documents.isEmpty()) {
|
result.put("success", false);
|
result.put("message", "文档解析结果为空");
|
return result;
|
}
|
|
List<Document> splitDocuments = textSplitter.apply(documents);
|
log.info("文本分割完成，共 {} 个分片", splitDocuments.size());
|
|
for (Document doc : splitDocuments) {
|
doc.getMetadata().put("fileName", fileName);
|
doc.getMetadata().put("contentType", contentType);
|
doc.getMetadata().put("uploadTime", LocalDateTime.now().toString());
|
doc.getMetadata().put("fileSize", file.getSize());
|
}
|
|
vectorStore.add(splitDocuments);
|
|
log.info("文档入库完成: {}，共 {} 个向量分片", fileName, splitDocuments.size());
|
result.put("success", true);
|
result.put("message", "入库成功");
|
result.put("chunksCount", splitDocuments.size());
|
|
} catch (Exception e) {
|
log.error("文档入库失败: {}", fileName, e);
|
result.put("success", false);
|
result.put("message", "入库失败：" + e.getMessage());
|
} finally {
|
if (tempFile != null) {
|
try {
|
Files.deleteIfExists(tempFile);
|
} catch (IOException e) {
|
log.warn("删除临时文件失败: {}", tempFile, e);
|
}
|
}
|
}
|
return result;
|
}
|
|
private boolean isSupportedContentType(String contentType) {
|
if (contentType == null) return false;
|
String[] supportedTypes = {
|
"text/plain",
|
"application/pdf",
|
"application/msword",
|
"application/vnd.openxmlformats-officedocument.wordprocessingml.document",
|
"application/vnd.ms-excel",
|
"application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
|
"text/markdown"
|
};
|
return Arrays.asList(supportedTypes).contains(contentType);
|
}
|
|
|
}
|
