Skip to content

Add numberOfRerankedResults to Bedrock KB store#6570

Open
rkargul wants to merge 1 commit into
spring-projects:mainfrom
rkargul:add-number-of-reranked-results-bedrock-kb
Open

Add numberOfRerankedResults to Bedrock KB store#6570
rkargul wants to merge 1 commit into
spring-projects:mainfrom
rkargul:add-number-of-reranked-results-bedrock-kb

Conversation

@rkargul

@rkargul rkargul commented Jul 7, 2026

Copy link
Copy Markdown

Expose the numberOfRerankedResults field of the Bedrock Agent Runtime Retrieve API in BedrockKnowledgeBaseVectorStore, so users can control how many results the Bedrock reranking model returns.

The value is configurable via the store builder and the new spring.ai.vectorstore.bedrock-knowledge-base.number-of-reranked-results auto-configuration property. It is only applied when a reranking model ARN is configured, since the reranking configuration is only attached to the Retrieve request in that case.

Closes #6569

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.

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

And via Spring Boot auto-configuration:

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

Current Behavior

The reranking support introduced with the store only exposes the reranking model ARN. BedrockKnowledgeBaseVectorStore#buildRerankingConfig() sets modelConfiguration but never numberOfRerankedResults, so the Retrieve call always uses the service default. The class is final and builds the request internally, so there is no way to set this value without forking the store or bypassing the VectorStore abstraction and calling the AWS SDK directly.

Context

We use this store in a RAG pipeline with a Bedrock reranking model: retrieve a larger candidate set (top-k), then have the reranker narrow it down to a small high-relevance set before it reaches the prompt. That is exactly what numberOfRerankedResults controls, and the field is already available in the bedrockagentruntime SDK version pinned by Spring AI. Trimming client-side after retrieval does not change what the reranker returns, and using getNativeClient() gives up the VectorStore abstraction and filter expression support.

Changes

  • numberOfRerankedResults builder option on BedrockKnowledgeBaseVectorStore, applied to the reranking configuration only when rerankingModelArn is set
  • New spring.ai.vectorstore.bedrock-knowledge-base.number-of-reranked-results property (validated with @min(1)) wired in the auto-configuration
  • Unit tests covering: value passed to the API, omitted when unset, ignored without a reranking model ARN, rejection of non-positive values
  • Reference documentation updates for the new property

Expose the numberOfRerankedResults field of the Bedrock Agent Runtime
Retrieve API in BedrockKnowledgeBaseVectorStore, so users can control
how many results the Bedrock reranking model returns.

The value is configurable via the store builder and the new
spring.ai.vectorstore.bedrock-knowledge-base.number-of-reranked-results
auto-configuration property. It is only applied when a reranking model
ARN is configured, since the reranking configuration is only attached
to the Retrieve request in that case.

Closes spring-projects#6569

Signed-off-by: Radek Kargul <radoslaw.kargul@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Support numberOfRerankedResults in BedrockKnowledgeBaseVectorStore

2 participants