feat: add Managed Knowledge Base support to AWS Bedrock KB Retriever

packages/components/nodes/retrievers/AWSBedrockKBRetriever/AWSBedrockKBRetriever.test.ts

@@ -0,0 +1,54 @@
+/**
+ * Tests for AWSBedrockKBRetriever node configuration.
+ */
+
+jest.mock('@langchain/aws', () => ({
+    AmazonKnowledgeBaseRetriever: jest.fn(),
+}));
+
+jest.mock('../../../src/Interface', () => ({}));
+jest.mock('../../../src/awsToolsUtils', () => ({
+    getAWSCredentialConfig: jest.fn(),
+}));
+jest.mock('../../../src/modelLoader', () => ({
+    MODEL_TYPE: { EMBEDDING: 'EMBEDDING' },
+    getRegions: jest.fn(),
+}));
+jest.mock('@aws-sdk/client-bedrock-agent-runtime', () => ({}));
+
+const { nodeClass: NodeClass } = require('./AWSBedrockKBRetriever');
+
+describe('AWSBedrockKBRetriever', () => {
+    const nodeInstance = new NodeClass();
+
+    it('should have Knowledge Base Type input', () => {
+        const kbTypeInput = nodeInstance.inputs.find(
+            (input: any) => input.name === 'knowledgeBaseType',
+        );
+        expect(kbTypeInput).toBeDefined();
+        expect(kbTypeInput.type).toBe('options');
+        expect(kbTypeInput.default).toBe('MANAGED');
+    });
+
+    it('should have MANAGED and VECTOR options', () => {
+        const kbTypeInput = nodeInstance.inputs.find(
+            (input: any) => input.name === 'knowledgeBaseType',
+        );
+        const optionNames = kbTypeInput.options.map((o: any) => o.name);
+        expect(optionNames).toContain('MANAGED');
+        expect(optionNames).toContain('VECTOR');
+    });
+
+    it('should have Knowledge Base ID input', () => {
+        const kbIdInput = nodeInstance.inputs.find(
+            (input: any) => input.name === 'knoledgeBaseID',
+        );
+        expect(kbIdInput).toBeDefined();
+        expect(kbIdInput.type).toBe('string');
+    });
+
+    it('should be categorized as a Retriever', () => {
+        expect(nodeInstance.category).toBe('Retrievers');
+        expect(nodeInstance.baseClasses).toContain('BaseRetriever');
+    });
+});

packages/components/nodes/retrievers/AWSBedrockKBRetriever/AWSBedrockKBRetriever.ts

@@ -4,6 +4,19 @@ import { getAWSCredentialConfig } from '../../../src/awsToolsUtils'
 import { RetrievalFilter } from '@aws-sdk/client-bedrock-agent-runtime'
 import { MODEL_TYPE, getRegions } from '../../../src/modelLoader'
 
+function getSourceUri(result: any): string {
+    if (result == null) return ''
+    const location = result.location ?? {}
+    if (location.s3Location) return location.s3Location.uri ?? ''
+    if (location.webLocation) return location.webLocation.url ?? ''
+    if (location.confluenceLocation) return location.confluenceLocation.url ?? ''
+    if (location.salesforceLocation) return location.salesforceLocation.url ?? ''
+    if (location.sharePointLocation) return location.sharePointLocation.url ?? ''
+    if (location.customDocumentLocation) return location.customDocumentLocation.id ?? ''
+    // Fallback for agentic results
+    return result.metadata?._source_uri ?? ''
+}
+
 class AWSBedrockKBRetriever_Retrievers implements INode {
     label: string
     name: string
@@ -84,6 +97,27 @@ class AWSBedrockKBRetriever_Retrievers implements INode {
                 additionalParams: true,
                 default: undefined
             },
+            {
+                label: 'Knowledge Base Type',
+                name: 'knowledgeBaseType',
+                type: 'options',
+                description: 'Type of knowledge base. MANAGED (recommended) uses managed search. VECTOR uses traditional vector search.',
+                options: [
+                    {
+                        label: 'MANAGED',
+                        name: 'MANAGED',
+                        description: 'Managed knowledge base - Bedrock handles embedding, storage, and retrieval automatically'
+                    },
+                    {
+                        label: 'VECTOR',
+                        name: 'VECTOR',
+                        description: 'Traditional vector knowledge base with external vector store'
+                    }
+                ],
+                optional: true,
+                additionalParams: true,
+                default: 'VECTOR'
+            },
             {
                 label: 'Filter',
                 name: 'filter',
@@ -108,7 +142,112 @@ class AWSBedrockKBRetriever_Retrievers implements INode {
         const topK = nodeData.inputs?.topK as number
         const overrideSearchType = (nodeData.inputs?.searchType != '' ? nodeData.inputs?.searchType : undefined) as 'HYBRID' | 'SEMANTIC'
         const filter = (nodeData.inputs?.filter != '' ? JSON.parse(nodeData.inputs?.filter) : undefined) as RetrievalFilter
+        const knowledgeBaseType = (nodeData.inputs?.knowledgeBaseType != '' ? nodeData.inputs?.knowledgeBaseType : undefined) as
+            | 'MANAGED'
+            | 'VECTOR'
+            | undefined
         const credentialConfig = await getAWSCredentialConfig(nodeData, options, region)
+
+        if (knowledgeBaseType === 'MANAGED') {
+            // Try langchain retriever first (works once @langchain/aws supports managedSearchConfiguration).
+            // If it fails with managed config error, fall back to direct Bedrock API call.
+            const langchainRetriever = new AmazonKnowledgeBaseRetriever({
+                topK,
+                knowledgeBaseId: knoledgeBaseID,
+                region,
+                clientOptions: {
+                    ...(credentialConfig.credentials && { credentials: credentialConfig.credentials })
+                }
+            })
+
+            const { BedrockAgentRuntimeClient, RetrieveCommand, AgenticRetrieveStreamCommand } = await import('@aws-sdk/client-bedrock-agent-runtime')
+            const { Document } = await import('@langchain/core/documents')
+            const { BaseRetriever } = await import('@langchain/core/retrievers')
+            const client = new BedrockAgentRuntimeClient({
+                region,
+                customUserAgent: [["flowise", "bedrock-kb"]],
+                ...(credentialConfig.credentials && { credentials: credentialConfig.credentials })
+            })
+
+            // Create a proper BaseRetriever subclass for managed KB
+            class ManagedKBRetriever extends BaseRetriever {
+                lc_namespace = ['custom']
+
+                async _getRelevantDocuments(query: string): Promise<any[]> {
+                    // Try agentic retrieval first if enabled
+                    const useAgenticRetrieval = process.env.USE_AGENTIC_RETRIEVAL !== 'false'
+                    if (useAgenticRetrieval) {
+                        try {
+                            const agenticCmd = new AgenticRetrieveStreamCommand({
+                                messages: [{ content: { text: query }, role: 'user' }],
+                                retrievers: [{
+                                    configuration: {
+                                        knowledgeBase: {
+                                            knowledgeBaseId: knoledgeBaseID,
+                                            retrievalOverrides: { maxNumberOfResults: topK },
+                                        },
+                                    },
+                                }],
+                                agenticRetrieveConfiguration: {
+                                    foundationModelType: 'MANAGED',
+                                    rerankingModelType: 'MANAGED',
+                                },
+                            } as any)
+                            const agenticResponse = await client.send(agenticCmd)
+                            const results: any[] = []
+                            const stream = (agenticResponse as any).stream
+                            if (stream) {
+                                for await (const event of stream) {
+                                    if (event.retrievalResult) {
+                                        results.push(new Document({
+                                            pageContent: event.retrievalResult.content?.text ?? '',
+                                            metadata: {
+                                                source: getSourceUri(event.retrievalResult),
+                                                score: event.retrievalResult.score ?? 0
+                                            }
+                                        }))
+                                    }
+                                }
+                            }
+                            if (results.length > 0) {
+                                return results
+                            }
+                        } catch {
+                            // Fall through to plain retrieve
+                        }
+                    }
+
+                    try {
+                        // Try langchain path (will work once @langchain/aws adds managed support)
+                        return await langchainRetriever.getRelevantDocuments(query)
+                    } catch (e: any) {
+                        if (e?.message?.includes('managedSearchConfiguration') || e?.message?.includes('vectorSearchConfiguration is not supported')) {
+                            // Fallback: direct API call with managedSearchConfiguration
+                            const command = new RetrieveCommand({
+                                knowledgeBaseId: knoledgeBaseID,
+                                retrievalQuery: { text: query },
+                                retrievalConfiguration: {
+                                    managedSearchConfiguration: { numberOfResults: topK }
+                                } as any
+                            })
+                            const response = await client.send(command)
+                            return (response.retrievalResults ?? []).map((result: any) => new Document({
+                                pageContent: result.content?.text ?? '',
+                                metadata: {
+                                    source: getSourceUri(result),
+                                    score: result.score ?? 0
+                                }
+                            }))
+                        }
+                        throw e
+                    }
+                }
+            }
+
+            return new ManagedKBRetriever()
+        }
+
+        // For VECTOR KBs, use the langchain retriever (existing behavior)
         const retriever = new AmazonKnowledgeBaseRetriever({
             topK,
             knowledgeBaseId: knoledgeBaseID,

packages/components/nodes/retrievers/AWSBedrockKBRetriever/BEDROCK_MANAGED_KB.md

@@ -0,0 +1,58 @@
+# Bedrock Managed Knowledge Base Support
+
+## Overview
+Adds a Flowise retriever node that queries Amazon Bedrock Knowledge Bases for managed retrieval in chatflow applications.
+
+## Usage
+In Flowise canvas:
+1. Drag the **AWS Bedrock KB Retriever** node onto the canvas
+2. Configure AWS credentials and Knowledge Base ID in the node properties
+3. Connect as a retriever input to a Conversational Retrieval Chain or Agent node
+
+```json
+{
+  "node": "AWSBedrockKBRetriever",
+  "inputs": {
+    "knowledgeBaseId": "YOUR_KB_ID",
+    "region": "us-east-1",
+    "useAgenticRetrieval": true,
+    "maxResults": 5
+  }
+}
+```
+
+## Configuration
+| Variable | Description | Default |
+|---|---|---|
+| KNOWLEDGE_BASE_ID | Bedrock Knowledge Base ID | None |
+| AWS_REGION | AWS region for the KB | us-east-1 |
+| USE_AGENTIC_RETRIEVAL | Enable agentic retrieval | true |
+| MAX_RESULTS | Maximum retrieval results | 5 |
+
+## Features
+- Managed search (no vector store needed)
+- Agentic retrieval with query decomposition + reranking
+- Automatic fallback to plain Retrieve if agentic fails
+- Multi-source support (S3, Web, Confluence, SharePoint)
+- Visual drag-and-drop configuration in Flowise UI
+
+## SDK Requirements
+- @aws-sdk/client-bedrock-agent-runtime >= 3.700
+- flowise >= 2.0
+
+## Required IAM Permissions
+```json
+{
+  "Effect": "Allow",
+  "Action": [
+    "bedrock:Retrieve",
+    "bedrock:AgenticRetrieveStream"
+  ],
+  "Resource": "arn:aws:bedrock:<region>:<account-id>:knowledge-base/<kb-id>"
+}
+```
+
+## References
+- [Build a Managed Knowledge Base](https://docs.aws.amazon.com/bedrock/latest/userguide/kb-build-managed.html)
+- [Retrieve API](https://docs.aws.amazon.com/bedrock/latest/userguide/kb-test-retrieve.html)
+- [Agentic Retrieval](https://docs.aws.amazon.com/bedrock/latest/userguide/kb-test-agentic.html)