From 698a531daed7d23b179b05e7edd8cbc993083eab Mon Sep 17 00:00:00 2001 From: Seol_JY Date: Tue, 19 Dec 2023 16:21:28 +0900 Subject: [PATCH] Refactor ContentFormatTransformer for improved readability - Extract Method - Replace explicit instanceof checks and casting with pattern matching - Change var to explicit type --- .../transformer/MetadataTransformerIT.java | 2 +- .../transformer/ContentFormatTransformer.java | 94 ++++++++++++------- 2 files changed, 59 insertions(+), 37 deletions(-) diff --git a/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/transformer/MetadataTransformerIT.java b/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/transformer/MetadataTransformerIT.java index aca29f27b..d984cb73d 100644 --- a/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/transformer/MetadataTransformerIT.java +++ b/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/transformer/MetadataTransformerIT.java @@ -189,7 +189,7 @@ public class MetadataTransformerIT { @Bean public ContentFormatTransformer contentFormatTransformer(DefaultContentFormatter defaultContentFormatter) { - return new ContentFormatTransformer(defaultContentFormatter, false); + return new ContentFormatTransformer(defaultContentFormatter); } } diff --git a/spring-ai-core/src/main/java/org/springframework/ai/transformer/ContentFormatTransformer.java b/spring-ai-core/src/main/java/org/springframework/ai/transformer/ContentFormatTransformer.java index 9465766d8..880abc735 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/transformer/ContentFormatTransformer.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/transformer/ContentFormatTransformer.java @@ -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 apply(List 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 updatedEmbedExcludeKeys = new ArrayList<>(docFormatter.getExcludedEmbedMetadataKeys()); + updatedEmbedExcludeKeys.addAll(toUpdateFormatter.getExcludedEmbedMetadataKeys()); + + List 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); + } + }