spring-projects / spring-projects/spring-ai

Support numberOfRerankedResults in BedrockKnowledgeBaseVectorStore

Open Beginner friendly
#6,569 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

status: waiting-for-triage
Dominant language
Java
Stars
9.5k
Forks
2.9k
Avg merge
1d 7h
Merged PRs (30d)
6

Description

Expected Behavior
BedrockKnowledgeBaseVectorStore should allow configuring how many results the Bedrock reranking model returns, by exposing the numberOfRerankedResults field of VectorSearchBedrockRerankingConfiguration from the Bedrock Agent Runtime Retrieve API.

Builder API:

BedrockKnowledgeBaseVectorStore store = BedrockKnowledgeBaseVectorStore
      .builder(bedrockAgentRuntimeClient, "kb-id")
      .rerankingModelArn("arn:aws:bedrock:us-west-2::foundation-model/amazon.rerank-v1:0")
      .numberOfRerankedResults(3)
      .build();

And the matching Spring Boot auto-configuration property:

spring.ai.vectorstore.bedrock-knowledge-base.reranking-model-arn=arn:aws:bedrock:us-west-2::foundation-model/amazon.rerank-v1:0
spring.ai.vectorstore.bedrock-knowledge-base.number-of-reranked-results=3

The value would only be applied when reranking-model-arn is set, since the reranking configuration is only attached to the request in that case.

Current Behavior

The reranking support added in #5251 only exposes the reranking model ARN. When BedrockKnowledgeBaseVectorStore#buildRerankingConfig() builds the VectorSearchBedrockRerankingConfiguration, it sets modelConfiguration but never numberOfRerankedResults, so the Retrieve call always falls back to the service default.

There is no workaround within Spring AI: the class is final with no extension point, and the request is built internally, so users who need to control the reranked result count have to fork the store or bypass the VectorStore abstraction and call the AWS SDK directly.

Context
I use BedrockKnowledgeBaseVectorStore in a RAG pipeline with a Bedrock reranking model. I need to retrieve a larger candidate set (top-k) and want the reranker to narrow it down to a small, high-relevance set before it reaches the prompt, which is exactly what numberOfRerankedResults controls in the Retrieve API (supported by the bedrockagentruntime SDK version already pinned by Spring AI).

Alternatives considered: trimming results client-side after retrieval (does not change what the reranker returns and wastes reranked results), and calling BedrockAgentRuntimeClient directly via getNativeClient() (gives up the VectorStore abstraction and filter expression support).

The change is small and follows the existing pattern used for searchType and rerankingModelArn (store builder field plus auto-configuration property). I have a working implementation with unit tests and documentation updates and will submit a PR for this issue.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start at BedrockKnowledgeBaseVectorStore#buildRerankingConfig and follow the existing rerankingModelArn builder field into the Spring Boot auto-configuration. Check the Bedrock Retrieve request configuration and the existing unit-test and documentation patterns for searchType and rerankingModelArn. Done means the builder and number-of-reranked-results property set the value when reranking-model-arn is configured, with tests and documentation updated.

Written by the indexing model from the issue text.

Assessment

Tech stack
aws, java
Domain
ai, backend, cloud
Issue type
Feature
Difficulty
2/5
Estimated time
Half a day
Activity status
Quiet
Clarity
Clearly specified
Newbie friendliness
68/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.