From 112813f06a63f2fd1f495cb9ec02970fd64f1cfc Mon Sep 17 00:00:00 2001 From: ghdcksgml1 Date: Fri, 21 Mar 2025 21:55:04 +0900 Subject: [PATCH 1/2] refact: sort by descending score In ConcatenationDocumentJoiner Signed-off-by: ghdcksgml1 --- .../join/ConcatenationDocumentJoiner.java | 8 ++++++-- .../ConcatenationDocumentJoinerTests.java | 20 ++++++++++++++++++- 2 files changed, 25 insertions(+), 3 deletions(-) diff --git a/spring-ai-core/src/main/java/org/springframework/ai/rag/retrieval/join/ConcatenationDocumentJoiner.java b/spring-ai-core/src/main/java/org/springframework/ai/rag/retrieval/join/ConcatenationDocumentJoiner.java index 56038587fa6..48e0845da4a 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/rag/retrieval/join/ConcatenationDocumentJoiner.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/rag/retrieval/join/ConcatenationDocumentJoiner.java @@ -17,6 +17,7 @@ package org.springframework.ai.rag.retrieval.join; import java.util.ArrayList; +import java.util.Comparator; import java.util.List; import java.util.Map; import java.util.function.Function; @@ -34,7 +35,7 @@ * by concatenating them into a single collection of documents. In case of duplicate * documents, the first occurrence is kept. The score of each document is kept as is. * - * @author Thomas Vitale + * @author Thomas Vitale, ghdcksgml1 * @since 1.0.0 */ public class ConcatenationDocumentJoiner implements DocumentJoiner { @@ -54,7 +55,10 @@ public List join(Map>> documentsForQuery) { .flatMap(List::stream) .flatMap(List::stream) .collect(Collectors.toMap(Document::getId, Function.identity(), (existing, duplicate) -> existing)) - .values()); + .values() + .stream() + .sorted(Comparator.comparing((Document d1) -> (d1.getScore() != null) ? d1.getScore() : 0.0).reversed()) + .toList()); } } diff --git a/spring-ai-core/src/test/java/org/springframework/ai/rag/retrieval/join/ConcatenationDocumentJoinerTests.java b/spring-ai-core/src/test/java/org/springframework/ai/rag/retrieval/join/ConcatenationDocumentJoinerTests.java index 39a588555e6..1e5eebb6b32 100644 --- a/spring-ai-core/src/test/java/org/springframework/ai/rag/retrieval/join/ConcatenationDocumentJoinerTests.java +++ b/spring-ai-core/src/test/java/org/springframework/ai/rag/retrieval/join/ConcatenationDocumentJoinerTests.java @@ -31,7 +31,7 @@ /** * Unit tests for {@link ConcatenationDocumentJoiner}. * - * @author Thomas Vitale + * @author Thomas Vitale, ghdcksgml1 */ class ConcatenationDocumentJoinerTests { @@ -92,4 +92,22 @@ void whenDuplicatedDocumentsThenOnlyFirstOccurrenceIsKept() { assertThat(result).extracting(Document::getText).containsOnlyOnce("Content 2"); } + @Test + void whenSeveralQueryExistsInMapThenDocumentsAreJoinedInDescendingScoreOrder() { + DocumentJoiner documentJoiner = new ConcatenationDocumentJoiner(); + var documentsForQuery = new HashMap>>(); + documentsForQuery.put(new Query("query1"), + List.of(List.of(Document.builder().id("1").text("Content 1").score(0.9).build(), + Document.builder().id("4").text("Content 4").score(0.6).build()), + List.of(Document.builder().id("2").text("Content 2").score(0.8).build()))); + documentsForQuery.put(new Query("query2"), + List.of(List.of(Document.builder().id("3").text("Content 3").score(0.7).build()))); + + List result = documentJoiner.join(documentsForQuery); + + assertThat(result).hasSize(4); + assertThat(result).extracting(Document::getId).containsExactly("1", "2", "3", "4"); + assertThat(result).extracting(Document::getScore).containsExactly(0.9, 0.8, 0.7, 0.6); + } + } From dcba3ba6a5aaaeb91f066648f422784ef889f890 Mon Sep 17 00:00:00 2001 From: ghdcksgml1 Date: Sat, 22 Mar 2025 12:57:45 +0900 Subject: [PATCH 2/2] refact: javadocs Signed-off-by: ghdcksgml1 --- .../ai/rag/retrieval/join/ConcatenationDocumentJoiner.java | 3 ++- .../rag/retrieval/join/ConcatenationDocumentJoinerTests.java | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/spring-ai-core/src/main/java/org/springframework/ai/rag/retrieval/join/ConcatenationDocumentJoiner.java b/spring-ai-core/src/main/java/org/springframework/ai/rag/retrieval/join/ConcatenationDocumentJoiner.java index 48e0845da4a..801de92d1fb 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/rag/retrieval/join/ConcatenationDocumentJoiner.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/rag/retrieval/join/ConcatenationDocumentJoiner.java @@ -35,7 +35,8 @@ * by concatenating them into a single collection of documents. In case of duplicate * documents, the first occurrence is kept. The score of each document is kept as is. * - * @author Thomas Vitale, ghdcksgml1 + * @author Thomas Vitale + * @author ghdcksgml1 * @since 1.0.0 */ public class ConcatenationDocumentJoiner implements DocumentJoiner { diff --git a/spring-ai-core/src/test/java/org/springframework/ai/rag/retrieval/join/ConcatenationDocumentJoinerTests.java b/spring-ai-core/src/test/java/org/springframework/ai/rag/retrieval/join/ConcatenationDocumentJoinerTests.java index 1e5eebb6b32..f955ea68e21 100644 --- a/spring-ai-core/src/test/java/org/springframework/ai/rag/retrieval/join/ConcatenationDocumentJoinerTests.java +++ b/spring-ai-core/src/test/java/org/springframework/ai/rag/retrieval/join/ConcatenationDocumentJoinerTests.java @@ -31,7 +31,8 @@ /** * Unit tests for {@link ConcatenationDocumentJoiner}. * - * @author Thomas Vitale, ghdcksgml1 + * @author Thomas Vitale + * @author ghdcksgml1 */ class ConcatenationDocumentJoinerTests {