Skip to content

Feat(rerank): Add Cohere Reranker with topN filtering and fallback support #3791

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 4 commits into
base: main
Choose a base branch
from

Conversation

KoreaNirsa
Copy link
Contributor

Summary

This PR introduces a new DocumentPostProcessor implementation that leverages the Cohere Rerank API to reorder retrieved documents based on semantic relevance.
It enables post-retrieval reranking using the Query input and relevance scores returned by the Cohere API.

Key Features

  1. RerankerPostProcessor Implementation : Integrates with Cohere Rerank API using WebClient and implements the DocumentPostProcessor interface.
  2. Top-N Filtering Support : Uses topN value from Query.context() to determine how many top-ranked documents to return.
  3. Score-Based Sorting : Reranked documents are sorted by descending relevance score.
  4. Conditional Activation : Rerank is only activated if spring.ai.rerank.enabled=true is explicitly set.
  5. External API Key Configuration : The Cohere API key is injected via spring.ai.rerank.cohere.api-key.
  6. Fallback Behavior : If rerank is disabled or not configured, a default no-op DocumentPostProcessor is registered instead.

Configuration & Behavior Notes

The rerank feature is only activated when the following properties are present in your configuration

spring.ai.rerank.enabled=true
spring.ai.rerank.cohere.api-key=your-api-key

If omitted or disabled, rerank is skipped, and no HTTP calls to Cohere will be made.

Controlling topN Results

If "topN" is not provided in the Query.context() or is invalid, it defaults to 3

Default behavior without topN

Query query = new Query("Rank documents based on AI ethics and NLP.");

Custom topN usage

Map<String, Object> context = Map.of("topN", 5);
Query query = new Query("Rank documents based on AI ethics and NLP.", List.of(), context);

This ensures the rerank feature remains optional, configurable, and backward-compatible.


Usage Example

@RestController
@RequestMapping("/rerank")
public class RerankController {
    private final DocumentPostProcessor reranker;

    public RerankController(DocumentPostProcessor reranker) {
        this.reranker = reranker;
    }

    @GetMapping
    public void rerank() {
    	Map<String, Object> context = Map.of("topN", 5);
    	Query query = new Query("Rank documents based on AI ethics and NLP.", List.of(), context);
    	
        List<Document> documents = List.of(
            new Document("AI is a system that thinks and learns like humans."),
            new Document("Machine learning is a subfield of AI."),
            new Document("NLP enables machines to understand human language."),
            new Document("Deep learning is based on neural networks."),
            new Document("Computer vision processes images and videos."),
            new Document("AI is also related to robotics."),
            new Document("Reinforcement learning is reward-based."),
            new Document("AI is used in various industries."),
            new Document("Machine translation is a key NLP application."),
            new Document("AI ethics deals with responsibility and trust.")
        );

        List<Document> reranked = reranker.apply(query, documents);
        reranked.forEach(doc -> System.out.println("- " + doc.getText() + " | score=" + doc.getMetadata().get("score")));
    }
}

Result

image

Notes

  • I'm not fully confident whether this PR adheres to Spring AI’s architectural principles. If it diverges from the expected structure, I am open to revising the implementation or closing this PR and resubmitting a more aligned version.
  • No test code is included at this stage. Once the overall structure is approved, I’m happy to follow up with tests.


@KoreaNirsa KoreaNirsa force-pushed the feature/cohere-rerank branch from 767a1cb to b2aa03a Compare July 11, 2025 06:32
- Implemented rerank logic using Cohere Rerank API with WebClient

Signed-off-by: nirsa <[email protected]>
- Register a default DocumentPostProcessor that returns documents as-is

Signed-off-by: nirsa <[email protected]>
exception
- Changed setRelevanceScore parameter from int to double
- Prevented IllegalFormatConversionException when formatting with %.4f

Signed-off-by: nirsa <[email protected]>
@KoreaNirsa KoreaNirsa force-pushed the feature/cohere-rerank branch from b2aa03a to 1906c1f Compare July 11, 2025 06:42
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant