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

RAG - Document joiner should sort by score

Browse Source 
The current implementation of ConcatenationDocumentJoiner should sort the final document list by score in descending order, so to keep the most relevant documents at the front of the list. This PR fixes that.

Signed-off-by: Thomas Vitale <ThomasVitale@users.noreply.github.com>
This commit is contained in:
Thomas Vitale
2025-05-18 17:56:09 +02:00
committed by Mark Pollack
parent 46491c0426
commit f9f71a805d
2 changed files with 32 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
spring-ai-rag/src/main/java/org/springframework/ai/rag/retrieval/join/ConcatenationDocumentJoiner.java
View File

@@ -1,5 +1,5 @@
/*
* Copyright 2023-2024 the original author or authors.
* Copyright 2023-2025 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
@@ -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;
@@ -32,7 +33,8 @@ import org.springframework.util.Assert;
/**
* Combines documents retrieved based on multiple queries and from multiple data sources
* 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.
* documents, the first occurrence is kept. The score of each document is kept as is. The
* result is a list of unique documents sorted by their score in descending order.
*
* @author Thomas Vitale
* @since 1.0.0
@@ -54,7 +56,11 @@ public class ConcatenationDocumentJoiner implements DocumentJoiner {
.flatMap(List::stream)
.flatMap(List::stream)
.collect(Collectors.toMap(Document::getId, Function.identity(), (existing, duplicate) -> existing))
.values());
.values()
.stream()
.sorted(Comparator.comparingDouble((Document doc) -> doc.getScore() != null ? doc.getScore() : 0.0)
.reversed())
.toList());
}
}

24
spring-ai-rag/src/test/java/org/springframework/ai/rag/retrieval/join/ConcatenationDocumentJoinerTests.java
View File

@@ -21,7 +21,6 @@ import java.util.List;
import java.util.Map;
import org.junit.jupiter.api.Test;
import org.springframework.ai.document.Document;
import org.springframework.ai.rag.Query;
@@ -92,4 +91,27 @@ class ConcatenationDocumentJoinerTests {
assertThat(result).extracting(Document::getText).containsOnlyOnce("Content 2");
}
@Test
void shouldSortDocumentsByDescendingScore() {
//@formatter:off
DocumentJoiner documentJoiner = new ConcatenationDocumentJoiner();
var documentsForQuery = new HashMap<Query, List<List<Document>>>();
documentsForQuery.put(new Query("query1"), List.of(
List.of(
Document.builder().id("1").text("Content 1").score(0.81).build(),
Document.builder().id("2").text("Content 2").score(0.83).build()),
List.of(
Document.builder().id("3").text("Content 3").score(null).build())));
documentsForQuery.put(new Query("query2"), List.of(
List.of(
Document.builder().id("4").text("Content 4").score(0.85).build(),
Document.builder().id("5").text("Content 5").score(0.77).build())));
List<Document> result = documentJoiner.join(documentsForQuery);
assertThat(result).hasSize(5);
assertThat(result).extracting(Document::getId).containsExactly("4", "2", "1", "5", "3");
//@formatter:on
}
}