From 7775a76f51267de6b102faf2d73556b730385600 Mon Sep 17 00:00:00 2001 From: Mark Pollack Date: Wed, 24 Jul 2024 17:37:34 -0400 Subject: [PATCH] Refactor SummaryMetadataEnricher for readability --- .../transformer/SummaryMetadataEnricher.java | 31 +++++++++++-------- 1 file changed, 18 insertions(+), 13 deletions(-) diff --git a/spring-ai-core/src/main/java/org/springframework/ai/transformer/SummaryMetadataEnricher.java b/spring-ai-core/src/main/java/org/springframework/ai/transformer/SummaryMetadataEnricher.java index 40aee27df..60ecf450b 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/transformer/SummaryMetadataEnricher.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/transformer/SummaryMetadataEnricher.java @@ -34,6 +34,7 @@ import org.springframework.util.CollectionUtils; * 'section_summary', 'prev_section_summary', 'next_section_summary' metadata fields. * * @author Christian Tzolov + * @since 1.0.0 */ public class SummaryMetadataEnricher implements DocumentTransformer { @@ -51,11 +52,11 @@ public class SummaryMetadataEnricher implements DocumentTransformer { Summarize the key topics and entities of the section. - Summary: """; + Summary:"""; public enum SummaryType { - PREVIOUS, CURRENT, NEXT; + PREVIOUS, CURRENT, NEXT } @@ -105,21 +106,25 @@ public class SummaryMetadataEnricher implements DocumentTransformer { } for (int i = 0; i < documentSummaries.size(); i++) { - Map summaryMetadata = new HashMap<>(); - if (i > 0 && this.summaryTypes.contains(SummaryType.PREVIOUS)) { - summaryMetadata.put(PREV_SECTION_SUMMARY_METADATA_KEY, documentSummaries.get(i - 1)); - } - if (i < (documentSummaries.size() - 1) && this.summaryTypes.contains(SummaryType.NEXT)) { - summaryMetadata.put(NEXT_SECTION_SUMMARY_METADATA_KEY, documentSummaries.get(i + 1)); - } - if (this.summaryTypes.contains(SummaryType.CURRENT)) { - summaryMetadata.put(SECTION_SUMMARY_METADATA_KEY, documentSummaries.get(i)); - } - + Map summaryMetadata = getSummaryMetadata(i, documentSummaries); documents.get(i).getMetadata().putAll(summaryMetadata); } return documents; } + private Map getSummaryMetadata(int i, List documentSummaries) { + Map summaryMetadata = new HashMap<>(); + if (i > 0 && this.summaryTypes.contains(SummaryType.PREVIOUS)) { + summaryMetadata.put(PREV_SECTION_SUMMARY_METADATA_KEY, documentSummaries.get(i - 1)); + } + if (i < (documentSummaries.size() - 1) && this.summaryTypes.contains(SummaryType.NEXT)) { + summaryMetadata.put(NEXT_SECTION_SUMMARY_METADATA_KEY, documentSummaries.get(i + 1)); + } + if (this.summaryTypes.contains(SummaryType.CURRENT)) { + summaryMetadata.put(SECTION_SUMMARY_METADATA_KEY, documentSummaries.get(i)); + } + return summaryMetadata; + } + }