Spring/spring-ai
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Refactor ContentFormatTransformer for improved readability

Browse Source 
- Extract Method
- Replace explicit instanceof checks and casting with pattern matching
- Change var to explicit type
This commit is contained in:
Seol_JY
2023-12-19 16:21:28 +09:00
committed by Mark Pollack
parent fc7ee87113
commit 698a531dae
2 changed files with 59 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
models/spring-ai-openai/src/test/java/org/springframework/ai/openai/transformer/MetadataTransformerIT.java
View File

@@ -189,7 +189,7 @@ public class MetadataTransformerIT {
@Bean
public ContentFormatTransformer contentFormatTransformer(DefaultContentFormatter defaultContentFormatter) {
return new ContentFormatTransformer(defaultContentFormatter, false);
return new ContentFormatTransformer(defaultContentFormatter);
}
}

94
spring-ai-core/src/main/java/org/springframework/ai/transformer/ContentFormatTransformer.java
View File

@@ -24,21 +24,38 @@ import org.springframework.ai.document.Document;
import org.springframework.ai.document.DocumentTransformer;
/**
* ContentFormatTransformer processes a list of documents by applying a content formatter
* to each document.
*
* @author Christian Tzolov
* @since 1.0.0
*/
public class ContentFormatTransformer implements DocumentTransformer {
/**
* Disable the content-formatter template rewrite.
*/
private boolean disableTemplateRewrite = false;
private final boolean disableTemplateRewrite;
private ContentFormatter contentFormatter;
private final ContentFormatter contentFormatter;
/**
* Creates a ContentFormatTransformer object with the given ContentFormatter.
* @param contentFormatter the ContentFormatter to be used for transforming the
* documents
*/
public ContentFormatTransformer(ContentFormatter contentFormatter) {
this(contentFormatter, false);
}
/**
* The ContentFormatTransformer class is responsible for processing a list of
* documents by applying a content formatter to each document.
* @param contentFormatter The ContentFormatter to be used for transforming the
* documents
* @param disableTemplateRewrite Flag indicating whether to disable the
* content-formatter template rewrite
*/
public ContentFormatTransformer(ContentFormatter contentFormatter, boolean disableTemplateRewrite) {
this.contentFormatter = contentFormatter;
this.disableTemplateRewrite = disableTemplateRewrite;
@@ -47,45 +64,50 @@ public class ContentFormatTransformer implements DocumentTransformer {
/**
* Post process documents chunked from loader. Allows extractors to be chained.
* @param documents to post process.
* @return
* @return processed documents
*/
public List<Document> apply(List<Document> documents) {
if (this.contentFormatter != null) {
documents.forEach(document -> {
// Update formatter
if (document.getContentFormatter() instanceof DefaultContentFormatter
&& this.contentFormatter instanceof DefaultContentFormatter) {
DefaultContentFormatter docFormatter = (DefaultContentFormatter) document.getContentFormatter();
DefaultContentFormatter toUpdateFormatter = (DefaultContentFormatter) this.contentFormatter;
var updatedEmbedExcludeKeys = new ArrayList<>(docFormatter.getExcludedEmbedMetadataKeys());
updatedEmbedExcludeKeys.addAll(toUpdateFormatter.getExcludedEmbedMetadataKeys());
var updatedInterfaceExcludeKeys = new ArrayList<>(docFormatter.getExcludedInferenceMetadataKeys());
updatedInterfaceExcludeKeys.addAll(toUpdateFormatter.getExcludedInferenceMetadataKeys());
var builder = DefaultContentFormatter.builder()
.withExcludedEmbedMetadataKeys(updatedEmbedExcludeKeys)
.withExcludedInferenceMetadataKeys(updatedInterfaceExcludeKeys)
.withMetadataTemplate(docFormatter.getMetadataTemplate())
.withMetadataSeparator(docFormatter.getMetadataSeparator());
if (!this.disableTemplateRewrite) {
builder.withTextTemplate(docFormatter.getTextTemplate());
}
document.setContentFormatter(builder.build());
}
else {
// Override formatter
document.setContentFormatter(this.contentFormatter);
}
});
if (contentFormatter != null) {
documents.forEach(this::processDocument);
}
return documents;
}
private void processDocument(Document document) {
if (document.getContentFormatter() instanceof DefaultContentFormatter docFormatter
&& contentFormatter instanceof DefaultContentFormatter toUpdateFormatter) {
updateFormatter(document, docFormatter, toUpdateFormatter);
}
else {
overrideFormatter(document);
}
}
private void updateFormatter(Document document, DefaultContentFormatter docFormatter,
DefaultContentFormatter toUpdateFormatter) {
List<String> updatedEmbedExcludeKeys = new ArrayList<>(docFormatter.getExcludedEmbedMetadataKeys());
updatedEmbedExcludeKeys.addAll(toUpdateFormatter.getExcludedEmbedMetadataKeys());
List<String> updatedInterfaceExcludeKeys = new ArrayList<>(docFormatter.getExcludedInferenceMetadataKeys());
updatedInterfaceExcludeKeys.addAll(toUpdateFormatter.getExcludedInferenceMetadataKeys());
DefaultContentFormatter.Builder builder = DefaultContentFormatter.builder()
.withExcludedEmbedMetadataKeys(updatedEmbedExcludeKeys)
.withExcludedInferenceMetadataKeys(updatedInterfaceExcludeKeys)
.withMetadataTemplate(docFormatter.getMetadataTemplate())
.withMetadataSeparator(docFormatter.getMetadataSeparator());
if (!disableTemplateRewrite) {
builder.withTextTemplate(docFormatter.getTextTemplate());
}
document.setContentFormatter(builder.build());
}
private void overrideFormatter(Document document) {
document.setContentFormatter(contentFormatter);
}
}