SpringBoot LangChain4j RAG知识库(三)
嵌入模型工厂
public class EmbeddingModelFactory {
//获取向量模型
public EmbeddingModel getEmbeddingModel(String type, RagEmbeddingUseConfig ragEmbeddingUseConfig) {
return switch (type) {
case "In-process (ONNX)" -> getAllMiniLmL6V2EmbeddingModel();
case "Amazon Bedrock" -> getBaBedrockTitanEmbeddingModel(ragEmbeddingUseConfig);
case "Azure OpenAI" -> getAzureOpenAiEmbeddingModel(ragEmbeddingUseConfig);
case "Cohere" -> getAzureCohereEmbeddingModel(ragEmbeddingUseConfig);
case "DashScope" -> getQwenEmbeddingModel(ragEmbeddingUseConfig);
case "GitHub" -> getGitHubModelsEmbeddingModel(ragEmbeddingUseConfig);
case "Google" -> getVertexAiEmbeddingModel(ragEmbeddingUseConfig);
case "Hugging" -> getHuggingFaceEmbeddingModel(ragEmbeddingUseConfig);
case "Jina" -> getJinaEmbeddingModel(ragEmbeddingUseConfig);
// java 20+ case "Jlama" -> getJlamaEmbeddingModel(ragEmbeddingUseConfig);
case "LocalAI" -> getLocalAiEmbeddingModel(ragEmbeddingUseConfig);
case "Mistral" -> getMistralAiEmbeddingModel(ragEmbeddingUseConfig);
case "Nomic" -> getNomicEmbeddingModel(ragEmbeddingUseConfig);
case "Ollama" -> getOllamaEmbeddingModel(ragEmbeddingUseConfig);
case "OpenAI" -> getOpenAiEmbeddingModel(ragEmbeddingUseConfig);
case "OpenAI Official" -> getOpenAiOfficialEmbeddingModel(ragEmbeddingUseConfig);
case "Qianfan" -> getQianfanEmbeddingModel(ragEmbeddingUseConfig);
case "Voyage" -> getVoyageAiEmbeddingModel(ragEmbeddingUseConfig);
case "Cloudflare" -> getWorkersAiEmbeddingModel(ragEmbeddingUseConfig);
case "Zhipu" -> getZhipuAiEmbeddingModel(ragEmbeddingUseConfig);
case "Xinference" -> getXinferenceEmbeddingModel(ragEmbeddingUseConfig);
case "Google AI Gemini Embeddings" -> getGoogleAiEmbeddingModel(ragEmbeddingUseConfig);
default -> null;
};
}
/**
* In-process (ONNX) 向量模型
* @return
*/
public EmbeddingModel getAllMiniLmL6V2EmbeddingModel() {
return new AllMiniLmL6V2EmbeddingModel();
}
/**
* Amazon Bedrock AWS 嵌入模型 需要配置AWS的凭据
* V1 和 V2 两个版本
* 向量维度 V1:1536 v2:256
* @return
*/
public BedrockTitanEmbeddingModel getBaBedrockTitanEmbeddingModel(RagEmbeddingUseConfig ragEmbeddingUseConfigVo) {
Map<String,String> map = JSONObject.parseObject(ragEmbeddingUseConfigVo.getEmbeddingConfigParams(),Map.class);
BedrockTitanEmbeddingModel embeddingModel = BedrockTitanEmbeddingModel.builder()
//地区
.region(Region.of(map.get("region")) )
//模型
.model(map.get("modelId"))
//向量维度
.dimensions(Integer.valueOf(map.get("dimensions")))
.normalize(Boolean.valueOf(map.get("normalize")))
.build();
return embeddingModel;
}
/**
* Azure OpenAI 嵌入模型
*
* @return
*/
public EmbeddingModel getAzureOpenAiEmbeddingModel (RagEmbeddingUseConfig ragEmbeddingUseConfigVo) {
Map<String,String> map = JSONObject.parseObject(ragEmbeddingUseConfigVo.getEmbeddingConfigParams(),Map.class);
AzureOpenAiEmbeddingModel model = AzureOpenAiEmbeddingModel.builder()
.apiKey(map.get("apiKey"))
.endpoint(map.get("endpoint"))
.deploymentName(map.get("deploymentName"))
.logRequestsAndResponses(Boolean.parseBoolean(map.get("deploymentName")))
.build();
// Response<Embedding> response = model.embed("文本向量转换 ");
return model;
}
/**
* Cohere 嵌入模型
* embed-english-light-v3.0
* @return
*/
public EmbeddingModel getAzureCohereEmbeddingModel (RagEmbeddingUseConfig ragEmbeddingUseConfigVo) {
Map<String,String> map = JSONObject.parseObject(ragEmbeddingUseConfigVo.getEmbeddingConfigParams(),Map.class);
CohereEmbeddingModel model = CohereEmbeddingModel.builder()
.baseUrl(map.get("baseUrl"))
.apiKey(map.get("apiKey"))
.modelName(map.get("modelName"))
.inputType(map.get("inputType"))
.timeout(ofSeconds(Long.parseLong(map.get("timeout"))))
.logRequests(Boolean.valueOf(map.get("logRequests")))
.logResponses(Boolean.valueOf(map.get("logResponses"))) // embeddings are huge in logs
.build();
// Response<List<Embedding>> response = model.embedAll(asList(segment1, segment2));
return model;
}
/**
* DashScope 嵌入模型
* @return
*/
private QwenEmbeddingModel getQwenEmbeddingModel(RagEmbeddingUseConfig ragEmbeddingUseConfigVo) {
Map<String,String> map = JSONObject.parseObject(ragEmbeddingUseConfigVo.getEmbeddingConfigParams(),Map.class);
return QwenEmbeddingModel.builder()
.apiKey(map.get("apiKey"))
.modelName(map.get("modelName"))
.dimension(Integer.valueOf(map.get("dimension")))
.build();
}
/**
* GitHub Models 嵌入模型
* ??废弃
* @return
*/
private EmbeddingModel getGitHubModelsEmbeddingModel(RagEmbeddingUseConfig ragEmbeddingUseConfigVo) {
Map<String,String> map = JSONObject.parseObject(ragEmbeddingUseConfigVo.getEmbeddingConfigParams(),Map.class);
// GitHubModelsEmbeddingModel model = GitHubModelsEmbeddingModel.builder()
// .gitHubToken(gitHubToken)
// .modelName(TEXT_EMBEDDING_3_SMALL)
// .logRequestsAndResponses(true)
// .build();
//
// Response<Embedding> response = model.embed("Please embed this sentence.");
//
// System.out.println(response);
// return model;
return null;
}
/**
* Google Vertex AI 嵌入模型
* @return
*/
private EmbeddingModel getVertexAiEmbeddingModel(RagEmbeddingUseConfig ragEmbeddingUseConfigVo) {
Map<String,String> map = JSONObject.parseObject(ragEmbeddingUseConfigVo.getEmbeddingConfigParams(),Map.class);
EmbeddingModel embeddingModel = VertexAiEmbeddingModel.builder()
.project(map.get("project"))
.location(map.get("location"))
.endpoint(map.get("endpoint")) // optional
.publisher(map.get("publisher"))
.modelName(map.get("modelName"))
.maxRetries(Integer.valueOf(map.get("maxRetries"))) // 2 by default
.maxSegmentsPerBatch(Integer.valueOf(map.get("maxSegmentsPerBatch"))) // up to 250 segments per batch
.maxTokensPerBatch(Integer.valueOf(map.get("maxTokensPerBatch"))) // up to 2048 tokens per segment
.taskType(VertexAiEmbeddingModel.TaskType.valueOf(map.get("taskType"))) // see below for the different task types
.titleMetadataKey(map.get("titleMetadataKey")) // for the RETRIEVAL_DOCUMENT task, you can specify a title
// for the text segment to identify its document origin
.autoTruncate(Boolean.valueOf(map.get("autoTruncate"))) // false by default: truncates segments longer than 2,048 input tokens
.outputDimensionality(Integer.valueOf(map.get("outputDimensionality"))) // for models that support different output vector dimensions
// .credentials(map.get("credentials")) // custom Google Cloud credentials
.build();
return embeddingModel;
}
/**
* Hugging Face 嵌入模型
* @return
*/
private EmbeddingModel getHuggingFaceEmbeddingModel(RagEmbeddingUseConfig ragEmbeddingUseConfigVo) {
Map<String,String> map = JSONObject.parseObject(ragEmbeddingUseConfigVo.getEmbeddingConfigParams(),Map.class);
EmbeddingModel embeddingModel = HuggingFaceEmbeddingModel.builder()
.accessToken(map.get("accessToken"))
.modelId(map.get("modelId") )//"sentence-transformers/all-MiniLM-L6-v2"
.waitForModel(Boolean.valueOf(map.get("waitForModel")))
.timeout(ofSeconds(Long.parseLong(map.get("timeout"))))
.build();
return embeddingModel;
}
/**
* Jina AI 嵌入模型
* @return
*/
private EmbeddingModel getJinaEmbeddingModel(RagEmbeddingUseConfig ragEmbeddingUseConfigVo) {
Map<String,String> map = JSONObject.parseObject(ragEmbeddingUseConfigVo.getEmbeddingConfigParams(),Map.class);
EmbeddingModel model = JinaEmbeddingModel.builder()
.apiKey(map.get("apiKey"))
.modelName(map.get("modelName"))
.logRequests(Boolean.valueOf(map.get("logRequests")))
.logResponses(Boolean.valueOf(map.get("logResponses"))) // embeddings are huge in logs
.build();
return model;
}
/**
* Jlama 嵌入模型 java 20+
* @return
*/
// private EmbeddingModel getJlamaEmbeddingModel(RagEmbeddingUseConfigVo ragEmbeddingUseConfigVo) {
// Map<String,String> map = JSONObject.parseObject(ragEmbeddingUseConfigVo.getEmbeddingConfigParams(),Map.class);
// EmbeddingModel embeddingModel = JlamaEmbeddingModel.builder()
// .modelName(map.get("modelName"))
// .build();
// return embeddingModel;
// }
/**
* LocalAI 嵌入模型
* @return
*/
private EmbeddingModel getLocalAiEmbeddingModel(RagEmbeddingUseConfig ragEmbeddingUseConfigVo) {
Map<String,String> map = JSONObject.parseObject(ragEmbeddingUseConfigVo.getEmbeddingConfigParams(),Map.class);
EmbeddingModel model = LocalAiEmbeddingModel.builder()
.baseUrl(map.get("baseUrl"))
.modelName(map.get("modelName"))
.logRequests(Boolean.valueOf(map.get("logRequests")))
.logResponses(Boolean.valueOf(map.get("logResponses")))
.build();
return model;
}
/**
* Mistral AI 嵌入模型
* @return
*/
private EmbeddingModel getMistralAiEmbeddingModel(RagEmbeddingUseConfig ragEmbeddingUseConfigVo) {
Map<String,String> map = JSONObject.parseObject(ragEmbeddingUseConfigVo.getEmbeddingConfigParams(),Map.class);
EmbeddingModel embeddingModel = MistralAiEmbeddingModel.builder()
.apiKey(map.get("apiKey"))
.modelName(map.get("modelName"))
.build();
return embeddingModel;
}
/**
* Nomic 嵌入模型
* @return
*/
private EmbeddingModel getNomicEmbeddingModel(RagEmbeddingUseConfig ragEmbeddingUseConfigVo) {
Map<String,String> map = JSONObject.parseObject(ragEmbeddingUseConfigVo.getEmbeddingConfigParams(),Map.class);
EmbeddingModel model = NomicEmbeddingModel.builder()
.baseUrl(map.get("baseUrl"))
.apiKey(map.get("apiKey"))
.modelName(map.get("modelName"))
.taskType(map.get("taskType"))
.maxSegmentsPerBatch(Integer.valueOf(map.get("maxSegmentsPerBatch")))
.timeout(ofSeconds(Long.parseLong(map.get("timeout"))))
.logRequests(Boolean.valueOf(map.get("logRequests")))
.logResponses(Boolean.valueOf(map.get("logResponses"))) // embeddings are huge in logs
.build();
return model;
}
/**
* Ollama 嵌入模型
* @return
*/
private EmbeddingModel getOllamaEmbeddingModel(RagEmbeddingUseConfig ragEmbeddingUseConfigVo) {
HttpClientBuilder httpClientBuilder = new SpringRestClientBuilder();
Map<String,String> map = JSONObject.parseObject(ragEmbeddingUseConfigVo.getEmbeddingConfigParams(),Map.class);
EmbeddingModel model = OllamaEmbeddingModel.builder()
.baseUrl(map.get("baseUrl"))
.modelName(map.get("modelName"))
.httpClientBuilder(httpClientBuilder)
.build();
return model;
}
/**
* OpenAI 嵌入模型
* @return
*/
private EmbeddingModel getOpenAiEmbeddingModel(RagEmbeddingUseConfig ragEmbeddingUseConfigVo) {
HttpClientBuilder httpClientBuilder = new SpringRestClientBuilder();
Map<String,String> map = JSONObject.parseObject(ragEmbeddingUseConfigVo.getEmbeddingConfigParams(),Map.class);
EmbeddingModel model = OpenAiEmbeddingModel.builder()
.apiKey(map.get("apiKey"))
// .baseUrl(map.get("baseUrl"))
.modelName(map.get("modelName"))
.httpClientBuilder(httpClientBuilder)
.build();
return model;
}
/**
* OpenAI Official SDK 嵌入模型
* @return
*/
private EmbeddingModel getOpenAiOfficialEmbeddingModel(RagEmbeddingUseConfig ragEmbeddingUseConfigVo) {
Map<String,String> map = JSONObject.parseObject(ragEmbeddingUseConfigVo.getEmbeddingConfigParams(),Map.class);
EmbeddingModel model = OpenAiOfficialEmbeddingModel.builder()
.baseUrl(map.get("baseUrl"))
.apiKey(map.get("apiKey"))
.modelName(map.get("modelName"))
.isAzure(Boolean.parseBoolean(map.get("isAzure"))) // Azure 使用Azure 嵌入模型
.isGitHubModels(Boolean.parseBoolean(map.get("isGitHubModels")))//github 提供
.build();
return model;
}
/**
* Qianfan 嵌入模型
* @return
*/
private EmbeddingModel getQianfanEmbeddingModel(RagEmbeddingUseConfig ragEmbeddingUseConfigVo) {
Map<String,String> map = JSONObject.parseObject(ragEmbeddingUseConfigVo.getEmbeddingConfigParams(),Map.class);
QianfanEmbeddingModel model = QianfanEmbeddingModel.builder()
.user(map.get("user"))
.apiKey(map.get("apiKey"))
.secretKey(map.get("secretKey"))
.endpoint(map.get("endpoint"))
.logRequests(Boolean.valueOf(map.get("logRequests")))
.logResponses(Boolean.valueOf(map.get("logResponses")))
.build();
return model;
}
/**
* Google Vertex AI 嵌入模型
* @return
*/
private EmbeddingModel getVoyageAiEmbeddingModel(RagEmbeddingUseConfig ragEmbeddingUseConfigVo) {
Map<String,String> map = JSONObject.parseObject(ragEmbeddingUseConfigVo.getEmbeddingConfigParams(),Map.class);
EmbeddingModel model = VoyageAiEmbeddingModel.builder()
.apiKey(map.get("apiKey"))
.modelName(map.get("modelName"))
.timeout(Duration.ofSeconds(Long.parseLong(map.get("timeout"))))
.logRequests(Boolean.valueOf(map.get("logRequests")))
.logResponses(Boolean.valueOf(map.get("logResponses"))) // embeddings are huge in logs
.build();
return model;
}
/**
* Cloudflare Workers AI 嵌入模型
* @return
*/
private EmbeddingModel getWorkersAiEmbeddingModel(RagEmbeddingUseConfig ragEmbeddingUseConfigVo) {
Map<String,String> map = JSONObject.parseObject(ragEmbeddingUseConfigVo.getEmbeddingConfigParams(),Map.class);
EmbeddingModel embeddingModel = WorkersAiEmbeddingModel.builder()
.modelName(map.get("modelName"))
.accountId(map.get("accountId"))
.apiToken(map.get("apiToken"))
.build();
return embeddingModel;
}
/**
* Zhipu AI 嵌入模型
* @return
*/
private EmbeddingModel getZhipuAiEmbeddingModel(RagEmbeddingUseConfig ragEmbeddingUseConfigVo) {
Map<String,String> map = JSONObject.parseObject(ragEmbeddingUseConfigVo.getEmbeddingConfigParams(),Map.class);
ZhipuAiEmbeddingModel embeddingModel = ZhipuAiEmbeddingModel.builder()
.model(map.get("model"))
.apiKey(map.get("apiKey"))
.logRequests(Boolean.valueOf(map.get("logRequests")))
.logResponses(Boolean.valueOf(map.get("logResponses")))
.maxRetries(Integer.valueOf(map.get("maxRetries")))
.build();
return embeddingModel;
}
/**
* Xinference 嵌入模型
* @return
*/
private EmbeddingModel getXinferenceEmbeddingModel(RagEmbeddingUseConfig ragEmbeddingUseConfigVo) {
Map<String,String> map = JSONObject.parseObject(ragEmbeddingUseConfigVo.getEmbeddingConfigParams(),Map.class);
EmbeddingModel model = XinferenceEmbeddingModel.builder()
.baseUrl(map.get("baseUrl"))
.apiKey(map.get("apiKey"))
.modelName(map.get("modelName"))
.logRequests(Boolean.valueOf(map.get("logRequests")))
.logResponses(Boolean.valueOf(map.get("logResponses")))
.build();
return model;
}
/**
* Google AI Gemini Embeddings 嵌入模型
* @return
*/
private EmbeddingModel getGoogleAiEmbeddingModel(RagEmbeddingUseConfig ragEmbeddingUseConfigVo) {
Map<String,String> map = JSONObject.parseObject(ragEmbeddingUseConfigVo.getEmbeddingConfigParams(),Map.class);
EmbeddingModel embeddingModel = GoogleAiEmbeddingModel.builder()
.apiKey(map.get("apiKey"))
.modelName(map.get("modelName"))
.taskType(GoogleAiEmbeddingModel.TaskType.valueOf(map.get("taskType")))
.outputDimensionality(Integer.valueOf(map.get("outputDimensionality")))
.titleMetadataKey(map.get("titleMetadataKey"))
.maxRetries(Integer.valueOf(map.get("maxRetries")))
.timeout(Duration.ofSeconds(Long.parseLong(map.get("timeout"))))
.logRequestsAndResponses(Boolean.valueOf(map.get("logRequestsAndResponses")))
.build();
return embeddingModel;
}
}3、文档解析
Langchain4j 支持的解析文件很多,目前发现excel不是逐行解析,所以需要自定义excel解析方式
pom依赖添加
<dependency>
<groupId>com.alibaba</groupId>
<artifactId>easyexcel</artifactId>
<version>3.3.4</version>
<exclusions>
<exclusion>
<groupId>org.apache.poi</groupId>
<artifactId>poi-ooxml-schemas</artifactId>
</exclusion>
</exclusions>
</dependency>
<dependency>
<groupId>org.apache.poi</groupId>
<artifactId>poi-ooxml</artifactId>
<version>4.1.2</version>
</dependency>解析方式
//解析器接口
public interface IAgentDocumentParser extends DocumentParser{
/**
* 解析成多 Document
* @param inputStream
* @return
*/
List<Document> parseMore(InputStream inputStream);
}//默认解析器
public class DefaultDocumentParser implements IAgentDocumentParser{
private DocumentParser parser;
public DefaultDocumentParser(DocumentParser parser) {
super();
this.parser = parser;
}
@Override
public Document parse(InputStream inputStream) {
return parser.parse(inputStream);
}
@Override
public List<Document> parseMore(InputStream inputStream) {
return Lists.newArrayList(parse(inputStream));
}
}
//自定义excel解析器
@Slf4j
public class ExcelFileDocumentParser implements IAgentDocumentParser {
@Override
public Document parse(InputStream inputStream) {
throw new BaseException("ExcelFileDocumentParser 不支持解析成单 Document");
}
@Override
public List<Document> parseMore(InputStream inputStream) {
List<Map<String, String>> dicts = readExcelWithHeader(inputStream);
List<Document> documents = Lists.newArrayListWithExpectedSize(dicts.size());
dicts.forEach(dict -> {
Metadata metadata = new Metadata();
StringBuilder sb = new StringBuilder("");
dict.forEach((k, v) -> {
String value = v == null ? "" : v;
sb.append(value).append("\t");
metadata.put(k, value);
});
Document document = Document.from(sb.toString(), metadata);
documents.add(document);
});
return documents;
}
/**
* 读取 Excel,使用表头作为 Map 的 key
*
* @return List<Map<String, String>> 每行数据作为一个Map
*/
private List<Map<String, String>> readExcelWithHeader(InputStream inputStream) {
List<Map<String, String>> result = new ArrayList<>();
try {
// 方式1:使用监听器
EasyExcel.read(inputStream, new ReadListener<Map<Integer, String>>() {
private List<ReadCellData<?>> headList;
@Override
public void invokeHead(Map<Integer, ReadCellData<?>> headMap, AnalysisContext context) {
// 获取表头
headList = new ArrayList<>(headMap.values());
}
@Override
public void invoke(Map<Integer, String> data, AnalysisContext context) {
// 将数据转换为以表头为key的Map
Map<String, String> rowData = new LinkedHashMap<>();
for (Map.Entry<Integer, String> entry : data.entrySet()) {
int index = entry.getKey();
if (index < headList.size()) {
String headerName = headList.get(index).getStringValue();
rowData.put(headerName, entry.getValue());
}
}
result.add(rowData);
}
@Override
public void doAfterAllAnalysed(AnalysisContext context) {
log.info("Excel读取完成,共读取" + result.size() + "行数据");
}
}).sheet().doRead();
} catch (Exception e) {
log.error("读取Excel失败", e);
throw new BaseException("读取Excel失败");
}
return result;
}
}文件上传
public class FileUploadDocument {
/**
* 处理上传的文件并解析为文档
*/
public List<Document> processUploadedFile(MultipartFile file) throws IOException {
// 验证文件
validateFile(file);
// 根据文件类型选择解析器
IAgentDocumentParser parser = getDocumentParser(file);
// 创建加载器
MultipartFileDocumentLoader loader = new MultipartFileDocumentLoader(parser);
return loader.load(file);
}
/**
* 批量处理上传的文件
*/
public List<Document> processUploadedFiles(List<MultipartFile> files) throws IOException {
List<Document> documents = new ArrayList<>();
for (MultipartFile file : files) {
try {
validateFile(file);
List<Document> docs = processUploadedFile(file);
documents.addAll(docs);
} catch (Exception e) {
// 记录错误但继续处理
System.err.println("Failed to process file: " +
file.getOriginalFilename());
}
}
return documents;
}
/**
* 验证文件
*/
private void validateFile(MultipartFile file) {
if (file.isEmpty()) {
throw new IllegalArgumentException("File is empty");
}
String fileName = file.getOriginalFilename();
if (fileName == null || fileName.trim().isEmpty()) {
throw new IllegalArgumentException("Invalid file name");
}
// 检查文件大小(例如限制为 50MB)
long maxFileSize = 50 * 1024 * 1024; // 50MB
if (file.getSize() > maxFileSize) {
throw new IllegalArgumentException("File size exceeds limit");
}
// 检查文件扩展名
String extension = getFileExtension(fileName);
if (!isSupportedExtension(extension)) {
throw new IllegalArgumentException("Unsupported file type: " + extension);
}
}
/**
* 根据文件类型获取解析器
*/
private IAgentDocumentParser getDocumentParser(MultipartFile file) {
String fileName = file.getOriginalFilename();
String extension = getFileExtension(fileName).toLowerCase();
switch (extension) {
case "pdf" -> {
return new DefaultDocumentParser(new ApachePdfBoxDocumentParser());
}
case "txt", "html" -> {
return new DefaultDocumentParser(new TextDocumentParser());
}
case "md" -> {
return new DefaultDocumentParser(new MarkdownDocumentParser());
}
case "yml" -> {
return new DefaultDocumentParser(new YamlDocumentParser());
}
case "yaml" -> {
return new DefaultDocumentParser(new YamlDocumentParser());
//DOC, DOCX, PPT, PPTX, XLS, XLSX
}
case "docx", "doc", "ppt", "pptx" -> {
return new DefaultDocumentParser(new ApachePoiDocumentParser());
}
case "xls", "xlsx", ".csv" -> {
return new ExcelFileDocumentParser();
}
default -> {
return new DefaultDocumentParser(new ApacheTikaDocumentParser());
}
}
}
private String getFileExtension(String fileName) {
int lastDot = fileName.lastIndexOf('.');
return lastDot == -1 ? "" : fileName.substring(lastDot + 1);
}
private boolean isSupportedExtension(String extension) {
Set<String> supported = Set.of(
"pdf", "doc", "docx", "xls", "xlsx", "ppt", "pptx",
"txt", "md", "html", "htm", "csv", "json"
);
return supported.contains(extension.toLowerCase());
}
@Data
public static class DocumentChunk {
private String content;
private int startIndex;
private int endIndex;
private Map<String, Object> metadata;
// getters and setters
}保存文件,文件留存记录
public class MultipartFileDocumentLoader {
private final IAgentDocumentParser documentParser;
public MultipartFileDocumentLoader() {
this.documentParser = new DefaultDocumentParser(new ApacheTikaDocumentParser());
}
public MultipartFileDocumentLoader(IAgentDocumentParser documentParser) {
this.documentParser = documentParser;
}
/**
* 加载单个 MultipartFile 并添加元数据
*/
public List<Document> load(MultipartFile file) throws IOException {
try (InputStream inputStream = file.getInputStream()) {
List<Document> documents = documentParser.parseMore(inputStream);
documents.forEach(document -> {
// 添加文件元数据
document.metadata().put("source", file.getOriginalFilename());
document.metadata().put("file_name", file.getOriginalFilename());
document.metadata().put("file_size", String.valueOf(file.getSize()));
document.metadata().put("content_type", file.getContentType());
});
return documents;
}
}
/**
* 批量加载多个 MultipartFile 并添加元数据
*/
public List<Document> loadAll(List<MultipartFile> files)
throws IOException {
List<Document> documents = new ArrayList<>();
for (MultipartFile file : files) {
try {
List<Document> docs = load(file);
documents.addAll(docs);
} catch (IOException e) {
// 记录错误但继续处理其他文件
System.err.println("Failed to load file: " + file.getOriginalFilename() +
" - " + e.getMessage());
throw e; // 或者可以选择继续处理
}
}
return documents;
}
/**
* 从字节数组加载文档
*/
public List<Document> loadFromBytes(byte[] bytes, String fileName, String contentType) {
return loadFromBytes(bytes, fileName, contentType, Collections.emptyMap());
}
public List<Document> loadFromBytes(byte[] bytes, String fileName, String contentType,
Map<String, Object> metadata) {
try (ByteArrayInputStream inputStream = new ByteArrayInputStream(bytes)) {
List<Document> documents = documentParser.parseMore(inputStream);
documents.forEach(document->{
document.metadata().put("source", fileName);
document.metadata().put("file_name", fileName);
document.metadata().put("file_size", String.valueOf(bytes.length));
document.metadata().put("content_type", contentType);
if (metadata != null) {
document.metadata().putAll(metadata);
}
});
return documents;
} catch (IOException e) {
throw new RuntimeException("Failed to parse document from bytes", e);
}
}
}