feat: add Managed Knowledge Base support to AWS Bedrock KB Retriever#6612
feat: add Managed Knowledge Base support to AWS Bedrock KB Retriever#6612PVidyadhar wants to merge 1 commit into
Conversation
PVidyadhar
commented
Jul 11, 2026
- Added Knowledge Base Type dropdown (MANAGED/VECTOR)
- MANAGED path tries langchain first, falls back to direct Bedrock API
- Direct API uses managedSearchConfiguration
- Returns proper langchain Document objects via BaseRetriever subclass
- Unit tests included
- Added BEDROCK_MANAGED_KB.md design doc
There was a problem hiding this comment.
Code Review
This pull request introduces support for Amazon Bedrock Managed Knowledge Bases to the AWSBedrockKBRetriever node, allowing users to select between managed and traditional vector search. It implements a new ManagedKBRetriever class that supports agentic retrieval with query decomposition and automatic fallback mechanisms, adds corresponding unit tests, and provides documentation. The code review feedback suggests several key improvements: adding defensive nullish checks to prevent runtime errors in getSourceUri, optimizing performance by avoiding dynamic imports inside the query execution loop, and using the public getRelevantDocuments method instead of the protected _getRelevantDocuments to ensure LangChain lifecycle hooks and tracing are preserved.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| function getSourceUri(result: any): string { | ||
| const location = result.location ?? {} |
There was a problem hiding this comment.
To prevent potential runtime errors, add a defensive nullish check for the result parameter using loose equality (== null). Since this handles potentially invalid data from an external source, prefer throwing an error rather than silently returning an empty value.
| function getSourceUri(result: any): string { | |
| const location = result.location ?? {} | |
| function getSourceUri(result: any): string { | |
| if (result == null) { | |
| throw new Error('Invalid result received from AWS Bedrock') | |
| } | |
| const location = result.location ?? {} |
References
- In JavaScript/TypeScript, use loose equality (
== null) as a standard idiom for a 'nullish' check that covers bothnullandundefined. - When handling potentially invalid data from external sources (like an API response), prefer throwing an error for invalid input types rather than silently returning a default or empty value. This promotes fail-fast behavior.
| } | ||
| }) | ||
|
|
||
| const { BedrockAgentRuntimeClient, RetrieveCommand } = await import('@aws-sdk/client-bedrock-agent-runtime') |
There was a problem hiding this comment.
Import AgenticRetrieveStreamCommand here along with the other Bedrock client classes. This allows us to avoid re-importing the module dynamically on every single query execution inside _getRelevantDocuments.
| const { BedrockAgentRuntimeClient, RetrieveCommand } = await import('@aws-sdk/client-bedrock-agent-runtime') | |
| const { BedrockAgentRuntimeClient, RetrieveCommand, AgenticRetrieveStreamCommand } = await import('@aws-sdk/client-bedrock-agent-runtime') |
| // Try agentic retrieval first if enabled | ||
| const useAgenticRetrieval = process.env.USE_AGENTIC_RETRIEVAL !== 'false' | ||
| if (useAgenticRetrieval) { | ||
| try { | ||
| const { AgenticRetrieveStreamCommand } = await import('@aws-sdk/client-bedrock-agent-runtime') | ||
| const agenticCmd = new AgenticRetrieveStreamCommand({ |
There was a problem hiding this comment.
Use the pre-imported AgenticRetrieveStreamCommand and perform a feature-detection check. This avoids the overhead of dynamic imports on every query execution.
// Try agentic retrieval first if enabled
const useAgenticRetrieval = process.env.USE_AGENTIC_RETRIEVAL !== 'false'
if (useAgenticRetrieval && AgenticRetrieveStreamCommand) {
try {
const agenticCmd = new AgenticRetrieveStreamCommand({|
|
||
| try { | ||
| // Try langchain path (will work once @langchain/aws adds managed support) | ||
| return await langchainRetriever._getRelevantDocuments(query) |
There was a problem hiding this comment.
Avoid calling the protected method _getRelevantDocuments directly on langchainRetriever. Bypassing the public API skips standard LangChain lifecycle hooks, callbacks, and tracing (e.g., LangSmith). Use the public getRelevantDocuments method instead.
| return await langchainRetriever._getRelevantDocuments(query) | |
| return await langchainRetriever.getRelevantDocuments(query) |
b8ae9f8 to
ca93f4f
Compare
- Added Knowledge Base Type dropdown (MANAGED/VECTOR) - MANAGED path tries langchain first, falls back to direct Bedrock API - Direct API uses managedSearchConfiguration - Returns proper langchain Document objects via BaseRetriever subclass - Unit tests included - Added BEDROCK_MANAGED_KB.md design doc
ca93f4f to
e165f6a
Compare