Add numberOfRerankedResults to Bedrock KB store#6570
Open
rkargul wants to merge 1 commit into
Open
Conversation
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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Expose the
numberOfRerankedResultsfield of the Bedrock Agent Runtime Retrieve API inBedrockKnowledgeBaseVectorStore, 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-resultsauto-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
BedrockKnowledgeBaseVectorStoreshould allow configuring how many results the Bedrock reranking model returns, by exposing thenumberOfRerankedResultsfield ofVectorSearchBedrockRerankingConfigurationfrom the Bedrock Agent Runtime Retrieve API.And via Spring Boot auto-configuration:
Current Behavior
The reranking support introduced with the store only exposes the reranking model ARN.
BedrockKnowledgeBaseVectorStore#buildRerankingConfig()setsmodelConfigurationbut nevernumberOfRerankedResults, 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 theVectorStoreabstraction 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
numberOfRerankedResultscontrols, 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 usinggetNativeClient()gives up the VectorStore abstraction and filter expression support.Changes
numberOfRerankedResultsbuilder option onBedrockKnowledgeBaseVectorStore, applied to the reranking configuration only whenrerankingModelArnis setspring.ai.vectorstore.bedrock-knowledge-base.number-of-reranked-resultsproperty (validated with @min(1)) wired in the auto-configuration