diff --git a/packages/components/nodes/retrievers/AWSBedrockKBRetriever/AWSBedrockKBRetriever.test.ts b/packages/components/nodes/retrievers/AWSBedrockKBRetriever/AWSBedrockKBRetriever.test.ts new file mode 100644 index 00000000000..e3ade340065 --- /dev/null +++ b/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'); + }); +}); diff --git a/packages/components/nodes/retrievers/AWSBedrockKBRetriever/AWSBedrockKBRetriever.ts b/packages/components/nodes/retrievers/AWSBedrockKBRetriever/AWSBedrockKBRetriever.ts index c83001d0b61..b7092f36ae9 100644 --- a/packages/components/nodes/retrievers/AWSBedrockKBRetriever/AWSBedrockKBRetriever.ts +++ b/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,114 @@ 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 { + // 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.result?.results) { + for (const item of event.result.results) { + results.push(new Document({ + pageContent: item.content?.text ?? '', + metadata: { + source: getSourceUri(item), + score: item.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, diff --git a/packages/components/nodes/retrievers/AWSBedrockKBRetriever/BEDROCK_MANAGED_KB.md b/packages/components/nodes/retrievers/AWSBedrockKBRetriever/BEDROCK_MANAGED_KB.md new file mode 100644 index 00000000000..d218295e00f --- /dev/null +++ b/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:::knowledge-base/" +} +``` + +## 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)