Feat(rerank): Add Cohere Reranker support with topN result filtering

- Implemented rerank logic using Cohere Rerank API with WebClient

Signed-off-by: KoreaNirsa <[email protected]>

Author: KoreaNirsa

spring-ai-rag/src/main/java/org/springframework/ai/rag/postretrieval/rerank/CohereApi.java

@@ -0,0 +1,34 @@
+package org.springframework.ai.rag.postretrieval.rerank;
+
+/**
+ * Represents the API key holder for Cohere API authentication.
+ *
+ * @author KoreaNirsa
+ */
+public class CohereApi {
+	private String apiKey;
+
+	public static Builder builder() {
+		return new Builder();
+	}
+
+	public String getApiKey() {
+		return apiKey;
+	}
+
+	public static class Builder {
+		private final CohereApi instance = new CohereApi();
+
+		public Builder apiKey(String key) {
+			instance.apiKey = key;
+			return this;
+		}
+
+		public CohereApi build() {
+			if (instance.apiKey == null || instance.apiKey.isBlank()) {
+				throw new IllegalArgumentException("API key must be provided.");
+			}
+			return instance;
+		}
+	}
+}

spring-ai-rag/src/main/java/org/springframework/ai/rag/postretrieval/rerank/CohereReranker.java

@@ -0,0 +1,113 @@
+package org.springframework.ai.rag.postretrieval.rerank;
+
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.springframework.ai.document.Document;
+import org.springframework.http.HttpHeaders;
+import org.springframework.web.reactive.function.client.WebClient;
+
+/**
+ * A Reranker implementation that integrates with Cohere's Rerank API.
+ * This component reorders retrieved documents based on semantic relevance to the input query.
+ *
+ * @author KoreaNirsa
+ * @see <a href="https://docs.cohere.com/reference/rerank">Cohere Rerank API Documentation</a>
+ */
+public class CohereReranker {
+	private static final String COHERE_RERANK_ENDPOINT = "https://api.cohere.ai/v1/rerank";
+
+	private static final Logger logger = LoggerFactory.getLogger(CohereReranker.class);
+	
+	private static final int MAX_DOCUMENTS = 1000;
+
+	private final WebClient webClient;
+
+	/**
+	 * Constructs a CohereReranker that communicates with the Cohere Rerank API.
+	 * Initializes the internal WebClient with the provided API key for authorization.
+	 *
+	 * @param cohereApi the API configuration object containing the required API key (must not be null)
+	 * @throws IllegalArgumentException if cohereApi is null
+	 */
+    CohereReranker(CohereApi cohereApi) {
+        if (cohereApi == null) {
+            throw new IllegalArgumentException("CohereApi must not be null");
+        }
+
+        this.webClient = WebClient.builder()
+                .baseUrl(COHERE_RERANK_ENDPOINT)
+                .defaultHeader(HttpHeaders.AUTHORIZATION, "Bearer " + cohereApi.getApiKey())
+                .build();
+    }
+
+    /**
+     * Reranks a list of documents based on the provided query using the Cohere API.
+     *
+     * @param query The user input query.
+     * @param documents The list of documents to rerank.
+     * @param topN The number of top results to return (at most).
+     * @return A reranked list of documents. If the API fails, returns the original list.
+     */
+    public List<Document> rerank(String query, List<Document> documents, int topN) {
+        if (topN < 1) {
+            throw new IllegalArgumentException("topN must be ≥ 1. Provided: " + topN);
+        }
+
+        if (documents == null || documents.isEmpty()) {
+            logger.warn("Empty document list provided. Skipping rerank.");
+            return Collections.emptyList();
+        }
+
+        if (documents.size() > MAX_DOCUMENTS) {
+            logger.warn("Cohere recommends ≤ {} documents per rerank request. Larger sizes may cause errors.", MAX_DOCUMENTS);
+            return documents;
+        }
+
+        int adjustedTopN = Math.min(topN, documents.size());
+
+        Map<String, Object> payload = Map.of(
+            "query", query,
+            "documents", documents.stream().map(Document::getText).toList(),
+            "top_n", adjustedTopN
+        );
+
+        // Call the API and process the result
+        return sendRerankRequest(payload)
+                .map(results -> results.stream()
+                    .sorted(Comparator.comparingDouble(RerankResponse.Result::getRelevanceScore).reversed())
+                    .map(r -> documents.get(r.getIndex()))
+                    .toList())
+                .orElseGet(() -> {
+                    logger.warn("Cohere response is null or invalid");
+                    return documents;
+                });
+    }
+
+    /**
+     * Sends a rerank request to the Cohere API and returns the result list.
+     *
+     * @param payload The request body including query, documents, and top_n.
+     * @return An Optional list of reranked results, or empty if failed.
+     */
+    private Optional<List<RerankResponse.Result>> sendRerankRequest(Map<String, Object> payload) {
+        try {
+            RerankResponse response = webClient.post()
+                    .bodyValue(payload)
+                    .retrieve()
+                    .bodyToMono(RerankResponse.class)
+                    .block();
+
+            return Optional.ofNullable(response)
+                           .map(RerankResponse::getResults);
+        } catch (Exception e) {
+            logger.error("Cohere rerank failed, fallback to original order: {}", e.getMessage(), e);
+            return Optional.empty();
+        }
+    }
+}

spring-ai-rag/src/main/java/org/springframework/ai/rag/postretrieval/rerank/RerankConfig.java

@@ -0,0 +1,32 @@
+package org.springframework.ai.rag.postretrieval.rerank;
+
+import org.springframework.ai.rag.postretrieval.document.DocumentPostProcessor;
+import org.springframework.beans.factory.annotation.Value;
+import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+
+/**
+ * Rerank configuration that conditionally registers a DocumentPostProcessor
+ * when rerank is enabled via application properties.
+ *
+ * This configuration is activated only when the following properties are set
+ * 
+ * <ul>
+ *   <li>spring.ai.rerank.enabled=true</li>
+ *   <li>spring.ai.rerank.cohere.api-key=your-api-key</li>
+ * </ul>
+ *
+ * @author KoreaNirsa
+ */
+@Configuration
+public class RerankConfig {
+    @Value("${spring.ai.rerank.cohere.api-key}")
+    private String apiKey;
+    
+    @Bean
+    @ConditionalOnProperty(name = "spring.ai.rerank.enabled", havingValue = "true")
+    public DocumentPostProcessor rerankerPostProcessor() {
+        return new RerankerPostProcessor(CohereApi.builder().apiKey(apiKey).build());
+    }
+}

spring-ai-rag/src/main/java/org/springframework/ai/rag/postretrieval/rerank/RerankResponse.java

@@ -0,0 +1,51 @@
+package org.springframework.ai.rag.postretrieval.rerank;
+
+import java.util.List;
+
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+/**
+ * Represents the response returned from Cohere's Rerank API.
+ * The response includes a list of result objects that specify document indices
+ * and their semantic relevance scores.
+ *
+ * @author KoreaNirsa
+ */
+public class RerankResponse {
+    private List<Result> results;
+
+    public List<Result> getResults() {
+        return results;
+    }
+
+    public void setResults(List<Result> results) {
+        this.results = results;
+    }
+
+    /**
+     * Represents a single reranked document result returned by the Cohere API.
+     * Contains the original index and the computed relevance score.
+     */
+    public static class Result {
+        private int index;
+
+        @JsonProperty("relevance_score")
+        private int relevanceScore;
+
+        public int getIndex() {
+            return index;
+        }
+
+        public void setIndex(int index) {
+            this.index = index;
+        }
+
+        public int getRelevanceScore() {
+            return relevanceScore;
+        }
+
+        public void setRelevanceScore(int relevanceScore) {
+            this.relevanceScore = relevanceScore;
+        }
+    }
+}

spring-ai-rag/src/main/java/org/springframework/ai/rag/postretrieval/rerank/RerankerPostProcessor.java

@@ -0,0 +1,49 @@
+package org.springframework.ai.rag.postretrieval.rerank;
+
+import java.util.List;
+
+import org.springframework.ai.document.Document;
+import org.springframework.ai.rag.Query;
+import org.springframework.ai.rag.postretrieval.document.DocumentPostProcessor;
+
+/**
+ * The only supported entrypoint for rerank functionality in Spring AI RAG.
+ * This component delegates reranking logic to CohereReranker, using the provided API key.
+ * 
+ * This class is registered as a DocumentPostProcessor bean only if
+ * spring.ai.rerank.enabled=true is set in the application properties.
+ * 
+ * @author KoreaNirsa
+ */
+public class RerankerPostProcessor implements DocumentPostProcessor {
+    private final CohereReranker reranker;
+
+    RerankerPostProcessor(CohereApi cohereApi) {
+        this.reranker = new CohereReranker(cohereApi);
+    }
+
+    /**
+     * Processes the retrieved documents by applying semantic reranking using the Cohere API
+     * 
+     * @param query     the user's input query
+     * @param documents the list of documents to be reranked 
+     * @return a list of documents sorted by relevance score 
+     */
+    @Override
+    public List<Document> process(Query query, List<Document> documents) {
+        int topN = extractTopN(query);
+        return reranker.rerank(query.text(), documents, topN);
+    }
+
+    /**
+     * Extracts the top-N value from the query context. 
+     * If not present or invalid, it defaults to 3
+     * 
+     * @param query the query containing optional context parameters
+     * @return the number of top documents to return
+     */
+    private int extractTopN(Query query) {
+        Object value = query.context().get("topN");
+        return (value instanceof Number num) ? num.intValue() : 3;
+    }
+}