Skip to content

feat: add Amazon Bedrock Knowledge Base retriever component #14014

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
PVidyadhar wants to merge 4 commits into
base: main
Choose a base branch
from
Open

feat: add Amazon Bedrock Knowledge Base retriever component #14014

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
852219b
feat: add Amazon Bedrock Knowledge Base retriever component
PVidyadhar Jul 10, 2026
46c6967
[autofix.ci] apply automated fixes
autofix-ci[bot] Jul 16, 2026
2784d64
[autofix.ci] apply automated fixes (attempt 2/3)
autofix-ci[bot] Jul 16, 2026
ce28878
[autofix.ci] apply automated fixes (attempt 3/3)
autofix-ci[bot] Jul 16, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
254 changes: 251 additions & 3 deletions src/lfx/src/lfx/_assets/component_index.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7446,6 +7446,254 @@
}
}
],
[
"aws",
{
"BedrockKnowledgeBaseRetriever": {
"base_classes": [
"JSON"
],
"beta": false,
"conditional_paths": [],
"custom_fields": {},
"description": "Retrieve documents using langchain-aws AmazonKnowledgeBasesRetriever with managed or vector search.",
"display_name": "Amazon Bedrock Knowledge Base",
"documentation": "",
"edited": false,
"field_order": [
"knowledge_base_id",
"query",
"use_agentic_retrieval",
"region_name",
"number_of_results",
"aws_access_key_id",
"aws_secret_access_key"
],
"frozen": false,
"icon": "Amazon",
"legacy": false,
"metadata": {
"code_hash": "ee317a64d2ac",
"dependencies": {
"dependencies": [
{
"name": "langflow",
"version": null
},
{
"name": "boto3",
"version": "1.40.61"
},
{
"name": "botocore",
"version": "1.40.61"
},
{
"name": "langchain_aws",
"version": "1.1.0"
}
],
"total_dependencies": 4
},
"module": "lfx.components.aws.bedrock_knowledge_base_retriever.BedrockKnowledgeBaseRetrieverComponent"
},
"minimized": false,
"output_types": [],
"outputs": [
{
"allows_loop": false,
"cache": true,
"display_name": "Retrieved Documents",
"group_outputs": false,
"method": "retrieve",
"name": "documents",
"selected": "JSON",
"tool_mode": true,
"types": [
"JSON"
],
"value": "__UNDEFINED__"
}
],
"pinned": false,
"template": {
"_type": "Component",
"aws_access_key_id": {
"_input_type": "SecretStrInput",
"advanced": false,
"display_name": "AWS Access Key ID",
"dynamic": false,
"info": "AWS access key. Optional if using IAM role or environment credentials.",
"input_types": [],
"load_from_db": true,
"name": "aws_access_key_id",
"override_skip": false,
"password": true,
"placeholder": "",
"required": false,
"show": true,
"title_case": false,
"track_in_telemetry": false,
"type": "str",
"value": ""
},
"aws_secret_access_key": {
"_input_type": "SecretStrInput",
"advanced": false,
"display_name": "AWS Secret Access Key",
"dynamic": false,
"info": "AWS secret key. Optional if using IAM role or environment credentials.",
"input_types": [],
"load_from_db": true,
"name": "aws_secret_access_key",
"override_skip": false,
"password": true,
"placeholder": "",
"required": false,
"show": true,
"title_case": false,
"track_in_telemetry": false,
"type": "str",
"value": ""
},
"code": {
"advanced": true,
"dynamic": true,
"fileTypes": [],
"file_path": "",
"info": "",
"list": false,
"load_from_db": false,
"multiline": true,
"name": "code",
"password": false,
"placeholder": "",
"required": true,
"show": true,
"title_case": false,
"type": "code",
"value": "\"\"\"Amazon Bedrock Knowledge Base Retriever component for Langflow.\n\nExposes the langchain-aws AmazonKnowledgeBasesRetriever as a dedicated Langflow node.\n\"\"\"\n\nfrom langflow.custom import Component\nfrom langflow.io import BoolInput, IntInput, MessageTextInput, Output, SecretStrInput\nfrom langflow.schema import Data\n\n\ndef _get_source_uri(result: dict) -> str:\n \"\"\"Extract source URI from a retrieval result, handling all location types.\"\"\"\n location = result.get(\"location\", {})\n loc_type = location.get(\"type\", \"\")\n if loc_type == \"S3\" or \"s3Location\" in location:\n return location.get(\"s3Location\", {}).get(\"uri\", \"\")\n if loc_type == \"WEB\" or \"webLocation\" in location:\n return location.get(\"webLocation\", {}).get(\"url\", \"\")\n if \"confluenceLocation\" in location:\n return location.get(\"confluenceLocation\", {}).get(\"url\", \"\")\n if \"salesforceLocation\" in location:\n return location.get(\"salesforceLocation\", {}).get(\"url\", \"\")\n if \"sharePointLocation\" in location:\n return location.get(\"sharePointLocation\", {}).get(\"url\", \"\")\n if \"customDocumentLocation\" in location:\n return location.get(\"customDocumentLocation\", {}).get(\"id\", \"\")\n # Fallback to metadata._source_uri (for agentic results)\n return result.get(\"metadata\", {}).get(\"_source_uri\", \"\")\n\n\nclass BedrockKnowledgeBaseRetrieverComponent(Component):\n \"\"\"Retrieves documents from an Amazon Bedrock Knowledge Base.\n\n Wraps langchain-aws's AmazonKnowledgeBasesRetriever as a dedicated Langflow node.\n Supports both Managed Knowledge Bases (recommended) and Vector search.\n \"\"\"\n\n display_name = \"Amazon Bedrock Knowledge Base\"\n description = \"Retrieve documents using langchain-aws AmazonKnowledgeBasesRetriever with managed or vector search.\"\n icon = \"Amazon\"\n name = \"BedrockKnowledgeBaseRetriever\"\n\n inputs = [\n MessageTextInput(\n name=\"knowledge_base_id\",\n display_name=\"Knowledge Base ID\",\n info=\"The ID of the Amazon Bedrock Knowledge Base (10 alphanumeric characters).\",\n required=True,\n ),\n MessageTextInput(\n name=\"query\",\n display_name=\"Query\",\n info=\"The search query to retrieve relevant documents.\",\n required=True,\n ),\n BoolInput(\n name=\"use_agentic_retrieval\",\n display_name=\"Use Agentic Retrieval\",\n value=True,\n info=(\n \"If enabled, tries AgenticRetrieveStream first (query decomposition + managed reranking)\"\n \" with fallback to standard Retrieve.\"\n ),\n ),\n MessageTextInput(\n name=\"region_name\",\n display_name=\"AWS Region\",\n value=\"us-east-1\",\n info=\"AWS region where the Knowledge Base is located.\",\n ),\n IntInput(\n name=\"number_of_results\",\n display_name=\"Number of Results\",\n value=5,\n info=\"Maximum number of results to return.\",\n ),\n SecretStrInput(\n name=\"aws_access_key_id\",\n display_name=\"AWS Access Key ID\",\n info=\"AWS access key. Optional if using IAM role or environment credentials.\",\n required=False,\n ),\n SecretStrInput(\n name=\"aws_secret_access_key\",\n display_name=\"AWS Secret Access Key\",\n info=\"AWS secret key. Optional if using IAM role or environment credentials.\",\n required=False,\n ),\n ]\n\n outputs = [\n Output(display_name=\"Retrieved Documents\", name=\"documents\", method=\"retrieve\"),\n ]\n\n def retrieve(self) -> list[Data]:\n \"\"\"Retrieve documents using langchain-aws AmazonKnowledgeBasesRetriever.\"\"\"\n use_agentic_retrieval = self.use_agentic_retrieval\n\n # Try agentic retrieval first\n if use_agentic_retrieval:\n try:\n import boto3\n from botocore.config import Config\n\n # Build credentials kwargs for boto3\n boto_kwargs: dict = {\"region_name\": self.region_name}\n if self.aws_access_key_id and self.aws_secret_access_key:\n boto_kwargs[\"aws_access_key_id\"] = self.aws_access_key_id\n boto_kwargs[\"aws_secret_access_key\"] = self.aws_secret_access_key\n\n client = boto3.client(\n \"bedrock-agent-runtime\",\n config=Config(user_agent_extra=\"langflow/bedrock-kb\"),\n **boto_kwargs,\n )\n response = client.agentic_retrieve_stream(\n messages=[{\"content\": {\"text\": self.query}, \"role\": \"user\"}],\n retrievers=[\n {\n \"configuration\": {\n \"knowledgeBase\": {\n \"knowledgeBaseId\": self.knowledge_base_id,\n \"retrievalOverrides\": {\"maxNumberOfResults\": self.number_of_results},\n }\n }\n }\n ],\n agenticRetrieveConfiguration={\n \"foundationModelType\": \"MANAGED\",\n \"rerankingModelType\": \"MANAGED\",\n },\n )\n results = []\n for event in response.get(\"stream\", []):\n if \"result\" in event and \"results\" in event[\"result\"]:\n results.extend(\n Data(\n text=result.get(\"content\", {}).get(\"text\", \"\"),\n data={\n \"source\": _get_source_uri(result),\n \"score\": result.get(\"score\", 0.0),\n \"knowledge_base_id\": self.knowledge_base_id,\n },\n )\n for result in event[\"result\"][\"results\"]\n )\n if results:\n return results\n except Exception as e: # noqa: BLE001\n import logging\n\n logging.getLogger(__name__).debug(\"Agentic retrieval unavailable, falling back: %s\", e)\n\n try:\n from langchain_aws.retrievers import AmazonKnowledgeBasesRetriever\n except ImportError as err:\n msg = \"langchain-aws is required. Install with: pip install langchain-aws>=0.2.0\"\n raise ImportError(msg) from err\n\n # Build retrieval config\n retrieval_config = {\"managedSearchConfiguration\": {\"numberOfResults\": self.number_of_results}}\n # Build credentials kwargs\n credentials_kwargs = {}\n if self.aws_access_key_id and self.aws_secret_access_key:\n credentials_kwargs[\"credentials_profile_name\"] = None\n credentials_kwargs[\"aws_access_key_id\"] = self.aws_access_key_id\n credentials_kwargs[\"aws_secret_access_key\"] = self.aws_secret_access_key\n\n retriever = AmazonKnowledgeBasesRetriever(\n knowledge_base_id=self.knowledge_base_id,\n region_name=self.region_name,\n retrieval_config=retrieval_config,\n **credentials_kwargs,\n )\n\n docs = retriever.invoke(self.query)\n\n results = []\n for doc in docs:\n results.append(\n Data(\n text=doc.page_content,\n data={\n \"source\": doc.metadata.get(\"source\", \"\"),\n \"score\": doc.metadata.get(\"score\", 0.0),\n \"knowledge_base_id\": self.knowledge_base_id,\n },\n )\n )\n\n return results\n"
},
"knowledge_base_id": {
"_input_type": "MessageTextInput",
"advanced": false,
"display_name": "Knowledge Base ID",
"dynamic": false,
"info": "The ID of the Amazon Bedrock Knowledge Base (10 alphanumeric characters).",
"input_types": [
"Message"
],
"list": false,
"list_add_label": "Add More",
"load_from_db": false,
"name": "knowledge_base_id",
"override_skip": false,
"placeholder": "",
"required": true,
"show": true,
"title_case": false,
"tool_mode": false,
"trace_as_input": true,
"trace_as_metadata": true,
"track_in_telemetry": false,
"type": "str",
"value": ""
},
"number_of_results": {
"_input_type": "IntInput",
"advanced": false,
"display_name": "Number of Results",
"dynamic": false,
"info": "Maximum number of results to return.",
"list": false,
"list_add_label": "Add More",
"name": "number_of_results",
"override_skip": false,
"placeholder": "",
"required": false,
"show": true,
"title_case": false,
"tool_mode": false,
"trace_as_metadata": true,
"track_in_telemetry": true,
"type": "int",
"value": 5
},
"query": {
"_input_type": "MessageTextInput",
"advanced": false,
"display_name": "Query",
"dynamic": false,
"info": "The search query to retrieve relevant documents.",
"input_types": [
"Message"
],
"list": false,
"list_add_label": "Add More",
"load_from_db": false,
"name": "query",
"override_skip": false,
"placeholder": "",
"required": true,
"show": true,
"title_case": false,
"tool_mode": false,
"trace_as_input": true,
"trace_as_metadata": true,
"track_in_telemetry": false,
"type": "str",
"value": ""
},
"region_name": {
"_input_type": "MessageTextInput",
"advanced": false,
"display_name": "AWS Region",
"dynamic": false,
"info": "AWS region where the Knowledge Base is located.",
"input_types": [
"Message"
],
"list": false,
"list_add_label": "Add More",
"load_from_db": false,
"name": "region_name",
"override_skip": false,
"placeholder": "",
"required": false,
"show": true,
"title_case": false,
"tool_mode": false,
"trace_as_input": true,
"trace_as_metadata": true,
"track_in_telemetry": false,
"type": "str",
"value": "us-east-1"
},
"use_agentic_retrieval": {
"_input_type": "BoolInput",
"advanced": false,
"display_name": "Use Agentic Retrieval",
"dynamic": false,
"info": "If enabled, tries AgenticRetrieveStream first (query decomposition + managed reranking) with fallback to standard Retrieve.",
"list": false,
"list_add_label": "Add More",
"name": "use_agentic_retrieval",
"override_skip": false,
"placeholder": "",
"required": false,
"show": true,
"title_case": false,
"tool_mode": false,
"trace_as_metadata": true,
"track_in_telemetry": true,
"type": "bool",
"value": true
}
},
"tool_mode": false
}
}
],
[
"azure",
{
Expand Down Expand Up @@ -118554,9 +118802,9 @@
]
],
"metadata": {
"num_components": 354,
"num_modules": 95
"num_components": 355,
"num_modules": 96
},
"sha256": "d050e1e1e8c2fc3c8ae9cfddc650bb4a8ef08b3d95d814fb40dc9d7dd646859f",
"sha256": "0915a784a51d9e85c40a146bfcba5f8c0f8fb8cea4a5bf55d9bb6e096f8af893",
"version": "1.10.2"
}
62 changes: 62 additions & 0 deletions src/lfx/src/lfx/components/aws/BEDROCK_MANAGED_KB.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,62 @@
# Bedrock Managed Knowledge Base Support

## Overview
Adds a Langflow component that queries Amazon Bedrock Knowledge Bases for managed retrieval in visual workflows.

## Usage
```python
from lfx.components.aws.bedrock_kb import BedrockKnowledgeBaseComponent

component = BedrockKnowledgeBaseComponent()
component.set(
knowledge_base_id="YOUR_KB_ID",
region="us-east-1",
use_agentic_retrieval=True,
)
results = component.retrieve("What are our data retention policies?")
```

In Langflow UI: drag the **Bedrock Knowledge Base** component onto the canvas and configure via the properties panel.

## Configuration
| Variable | Description | Default |
|---|---|---|
| KNOWLEDGE_BASE_ID | Bedrock Knowledge Base ID | None |
| AWS_REGION | AWS region for the KB | us-east-1 |
| AWS_PROFILE | AWS credentials profile | None |
| 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 configuration in Langflow canvas

## SDK Requirements
- boto3 >= 1.43
- langflow >= 1.0

## Reranking Options
For managed search, these reranking modes are available:
- `MANAGED` (default) — automatic reranking by Bedrock
- `NONE` — disable reranking
- `CUSTOM` — your own Bedrock reranking model (e.g., Cohere Rerank v3.5)

## 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)

## Required IAM Permissions
```json
{
"Effect": "Allow",
"Action": [
"bedrock:Retrieve",
"bedrock:AgenticRetrieveStream"
],
"Resource": "arn:aws:bedrock:<region>:<account-id>:knowledge-base/<kb-id>"
}
```
Empty file added src/lfx/src/lfx/components/aws/__init__.py
View file
Open in desktop
Empty file.
Loading
Loading