feat: add Amazon Bedrock Knowledge Base retriever component#14014
feat: add Amazon Bedrock Knowledge Base retriever component#14014PVidyadhar wants to merge 2 commits into
Conversation
WalkthroughAdds a Langflow Amazon Bedrock Knowledge Base retriever with agentic retrieval, standard fallback, source URI extraction, component registration, documentation, and unit tests for configuration and credential handling. ChangesBedrock Knowledge Base Retrieval
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Langflow
participant BedrockAgentRuntime
participant AmazonKnowledgeBasesRetriever
participant Data
Langflow->>BedrockAgentRuntime: Request agentic retrieval
BedrockAgentRuntime-->>Langflow: Stream retrieval results
Langflow->>Data: Convert agentic results
Langflow->>AmazonKnowledgeBasesRetriever: Invoke standard retrieval on fallback
AmazonKnowledgeBasesRetriever-->>Langflow: Return documents
Langflow->>Data: Convert standard documents
Important Pre-merge checks failedPlease resolve all errors before merging. Addressing warnings is optional. ❌ Failed checks (1 error, 2 warnings)
✅ Passed checks (6 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
PVidyadhar
force-pushed
the
bmkb-managed-kb-support
branch
from
2ea6a3c to
eca4a57
Compare
github-actions
Bot
added
enhancement
PVidyadhar
force-pushed
the
bmkb-managed-kb-support
branch
from
eca4a57 to
c4048c4
Compare
github-actions
Bot
added
enhancement
PVidyadhar
force-pushed
the
bmkb-managed-kb-support
branch
from
af19d69 to
8ba346b
Compare
github-actions
Bot
added
enhancement
PVidyadhar
force-pushed
the
bmkb-managed-kb-support
branch
from
4507a1d to
0a31519
Compare
github-actions
Bot
added
enhancement
PVidyadhar
force-pushed
the
bmkb-managed-kb-support
branch
from
0a31519 to
f157291
Compare
github-actions
Bot
added
enhancement
- Created Langflow Component wrapping langchain-aws AmazonKnowledgeBasesRetriever - Supports MANAGED and VECTOR knowledge base types - Passes AWS credentials when provided - Unit tests included (3/3 pass) - Added BEDROCK_MANAGED_KB.md design doc
PVidyadhar
force-pushed
the
bmkb-managed-kb-support
branch
from
f157291 to
3fc4688
Compare
github-actions
Bot
added
enhancement
github-actions
Bot
added
enhancement
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 8
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/lfx/src/lfx/_assets/component_index.json`:
- Line 7571: Expose agentic retrieval as a per-component BoolInput named
use_agentic_retrieval with a default of True in
BedrockKnowledgeBaseRetrieverComponent.inputs, and update retrieve() to use this
component value instead of the USE_AGENTIC_RETRIEVAL environment variable.
Import BoolInput and preserve the existing agentic-first behavior when the
option is enabled.
- Line 7571: Update the agentic-retrieval try/except in
BedrockKnowledgeBaseRetrieverComponent.retrieve to catch only expected fallback
errors, such as botocore ClientError and AttributeError for unsupported SDK
methods, instead of Exception. Log a warning with the caught error before
falling back to standard retrieval, while preserving the existing fallback
behavior.
- Line 7571: Update BedrockKnowledgeBaseRetrieverComponent so
knowledge_base_type offers MANAGED and VECTOR options and retrieve() builds the
appropriate retrieval configuration, using managedSearchConfiguration for
MANAGED and vectorSearchConfiguration for VECTOR. Ensure the selected type is
also honored by agentic retrieval configuration, then regenerate the component
index asset from the updated component source.
- Around line 7477-7493: Update the boto3 dependency entry in the component
metadata to a 1.43.x version that includes
bedrock-agent-runtime.agentic_retrieve_stream, replacing the current 1.40.61 pin
while leaving the other dependencies unchanged.
In `@src/lfx/src/lfx/components/aws/bedrock_knowledge_base_retriever.py`:
- Line 95: Update the agentic retrieval toggle in the Bedrock knowledge base
retriever component to read from the component’s configured instance property
instead of exclusively from the USE_AGENTIC_RETRIEVAL environment variable.
Preserve the existing boolean behavior while honoring the value selected in the
Langflow properties panel.
- Around line 56-60: Update the knowledge_base_type DropdownInput in the Bedrock
knowledge-base retriever to provide its supported dropdown options, and add a
BoolInput for configuring agentic retrieval directly in the UI. Ensure the
retriever uses this input rather than relying solely on the environment-variable
fallback, while preserving the existing managed knowledge-base behavior.
- Around line 152-153: Update the retrieval configuration logic around
knowledge_base_type so VECTOR builds vectorSearchConfiguration and MANAGED
retains managedSearchConfiguration, using the existing number_of_results value.
Apply the same knowledge-base-type gate to the agentic retrieval path so it only
runs for MANAGED knowledge bases.
In `@src/lfx/tests/unit/components/aws/test_bedrock_knowledge_base_retriever.py`:
- Around line 95-97: Remove the duplicate `@patch` decorator above
test_retrieve_with_credentials, leaving a single patch for
AmazonKnowledgeBasesRetriever so the test receives only the declared
mock_retriever_class argument.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: c9494b0f-d29f-408f-901c-8d6da8bee899
📒 Files selected for processing (7)
PR_DESCRIPTION.mdsrc/lfx/src/lfx/_assets/component_index.jsonsrc/lfx/src/lfx/components/aws/BEDROCK_MANAGED_KB.mdsrc/lfx/src/lfx/components/aws/__init__.pysrc/lfx/src/lfx/components/aws/bedrock_knowledge_base_retriever.pysrc/lfx/tests/unit/components/aws/__init__.pysrc/lfx/tests/unit/components/aws/test_bedrock_knowledge_base_retriever.py
| "dependencies": { | ||
| "dependencies": [ | ||
| { | ||
| "name": "langflow", | ||
| "version": null | ||
| }, | ||
| { | ||
| "name": "boto3", | ||
| "version": "1.40.61" | ||
| }, | ||
| { | ||
| "name": "langchain_aws", | ||
| "version": "1.1.0" | ||
| } | ||
| ], | ||
| "total_dependencies": 3 | ||
| }, |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
sed -n '7470,7605p' src/lfx/src/lfx/_assets/component_index.jsonRepository: langflow-ai/langflow
Length of output: 12709
🌐 Web query:
boto3 1.43 agentic_retrieve_stream bedrock-agent-runtime
💡 Result:
In Boto3 (version 1.43.x), agentic_retrieve_stream is a method available within the bedrock-agent-runtime client that allows for agentic retrieval when querying knowledge bases [1][2]. This operation is designed to retrieve information from one or more knowledge bases using an agentic approach [3][4]. Unlike standard retrieval, which performs a single hybrid search, agentic retrieval uses a foundation model to intelligently decompose complex or multi-step queries into sub-queries and iterates to retrieve relevant information [3][4]. Key features of agentic_retrieve_stream include: - Streaming Response: It returns results through a stream that can include retrieval results, trace events for visibility into the process, and a synthesized natural-language answer [3][5][4]. - Configurable Response Generation: By default, the service generates a response using the retrieved information [5][4]. You can disable this by setting the generateResponse parameter to false, in which case only the retrieval results are returned [3][4]. - Traceability: The stream provides trace events that help monitor the planning and retrieval process [5][4]. - Citations: When response generation is enabled, the resulting synthesized answer can include citations that map specific spans of the text back to the supporting retrieval results [5]. The operation accepts an agenticRetrieveConfiguration, which allows you to specify the foundation model, set the maximum number of agent iterations, and configure reranking [3]. It is primarily used with managed knowledge bases [3][4].
Citations:
- 1: https://docs.aws.amazon.com/boto3/latest/reference/services/bedrock-agent-runtime.html
- 2: https://docs.aws.amazon.com/boto3/latest/reference/services/bedrock-agent-runtime/client/agentic_retrieve_stream.html
- 3: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent-runtime_AgenticRetrieveStream.html
- 4: https://docs.aws.amazon.com/bedrock-agentcore/latest/devguide/gateway-target-connector-managed-kb.html
- 5: https://docs.aws.amazon.com/bedrock/latest/userguide/kb-test-agentic-retrieve.html
Raise the boto3 pin to support agentic retrieval.
boto3 is pinned to 1.40.61, but bedrock-agent-runtime.agentic_retrieve_stream only appears in the newer 1.43.x line. With the current pin, the agentic path will fail and drop into the fallback, so the advertised feature won’t run in environments that use this metadata.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/lfx/src/lfx/_assets/component_index.json` around lines 7477 - 7493,
Update the boto3 dependency entry in the component metadata to a 1.43.x version
that includes bedrock-agent-runtime.agentic_retrieve_stream, replacing the
current 1.40.61 pin while leaving the other dependencies unchanged.
| "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 DropdownInput, 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 DropdownInput(\n name=\"knowledge_base_type\",\n display_name=\"Knowledge Base Type\",\n value=\"MANAGED\",\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 import os\n\n use_agentic_retrieval = os.environ.get(\"USE_AGENTIC_RETRIEVAL\", \"true\").lower() != \"false\"\n\n # Try agentic retrieval first\n if use_agentic_retrieval:\n try:\n import boto3\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(\"bedrock-agent-runtime\", **boto_kwargs)\n response = client.agentic_retrieve_stream(\n knowledgeBaseId=self.knowledge_base_id,\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: # noqa: BLE001, S110\n pass # Fall through to plain retrieve\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" |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Agentic retrieval toggle isn't exposed as a component input.
The PR objectives describe "agentic retrieval" as one of the configurable options in the properties panel, but it's only controllable via the USE_AGENTIC_RETRIEVAL environment variable inside retrieve() — there's no corresponding entry in inputs, field_order, or the template object. Flow authors have no way to toggle this per-component-instance from the Langflow UI as implied by the PR description.
Consider exposing this as a BoolInput (e.g. use_agentic_retrieval, default True) instead of relying solely on a process-wide environment variable.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/lfx/src/lfx/_assets/component_index.json` at line 7571, Expose agentic
retrieval as a per-component BoolInput named use_agentic_retrieval with a
default of True in BedrockKnowledgeBaseRetrieverComponent.inputs, and update
retrieve() to use this component value instead of the USE_AGENTIC_RETRIEVAL
environment variable. Import BoolInput and preserve the existing agentic-first
behavior when the option is enabled.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Broad, silent exception swallowing around agentic retrieval hides real failures.
except Exception: pass # noqa: BLE001, S110 in the agentic-retrieval branch catches everything — auth failures, throttling, malformed request errors, or a missing agentic_retrieve_stream attribute — and silently falls through to standard retrieval with zero logging. A missing bedrock:AgenticRetrieve permission or misconfiguration would be invisible to the operator, since the flow appears to "succeed" via fallback.
Narrow the caught exceptions (e.g. botocore.exceptions.ClientError, AttributeError for missing SDK support) and log a warning before falling back, so misconfigurations are diagnosable.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/lfx/src/lfx/_assets/component_index.json` at line 7571, Update the
agentic-retrieval try/except in BedrockKnowledgeBaseRetrieverComponent.retrieve
to catch only expected fallback errors, such as botocore ClientError and
AttributeError for unsupported SDK methods, instead of Exception. Log a warning
with the caught error before falling back to standard retrieval, while
preserving the existing fallback behavior.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
set -euo pipefail
# Locate the source component and inspect relevant sections
fd -a 'bedrock_knowledge_base_retriever.py|component_index.json' src/lfx/src/lfx -t f
echo
echo "=== Source outline ==="
ast-grep outline src/lfx/src/lfx/components/aws/bedrock_knowledge_base_retriever.py --view expanded || true
echo
echo "=== Source relevant lines ==="
python3 - <<'PY'
from pathlib import Path
p = Path('src/lfx/src/lfx/components/aws/bedrock_knowledge_base_retriever.py')
print('exists:', p.exists())
if p.exists():
lines = p.read_text().splitlines()
for start,end in [(1,260)]:
for i in range(start-1, min(end, len(lines))):
if 1 <= i+1 <= len(lines):
print(f"{i+1:4d}: {lines[i]}")
PY
echo
echo "=== Generated index snippet around the referenced lines ==="
python3 - <<'PY'
from pathlib import Path
p = Path('src/lfx/src/lfx/_assets/component_index.json')
lines = p.read_text().splitlines()
for start,end in [(7560,7635)]:
for i in range(start-1, min(end, len(lines))):
print(f"{i+1:4d}: {lines[i]}")
PYRepository: langflow-ai/langflow
Length of output: 20632
🏁 Script executed:
set -euo pipefail
echo "== source file exists? =="
ls -l src/lfx/src/lfx/components/aws/bedrock_knowledge_base_retriever.py
echo
echo "== search for knowledge_base_type usage =="
rg -n "knowledge_base_type|managedSearchConfiguration|vectorSearchConfiguration|options=" src/lfx/src/lfx/components/aws/bedrock_knowledge_base_retriever.py src/lfx/src/lfx/_assets/component_index.json
echo
echo "== excerpt from source =="
sed -n '1,240p' src/lfx/src/lfx/components/aws/bedrock_knowledge_base_retriever.pyRepository: langflow-ai/langflow
Length of output: 50378
🏁 Script executed:
python3 - <<'PY'
from pathlib import Path
paths = [
Path('src/lfx/src/lfx/components/aws/bedrock_knowledge_base_retriever.py'),
Path('src/lfx/src/lfx/_assets/component_index.json'),
]
for p in paths:
print(f'== {p} ==')
print('exists:', p.exists())
if p.exists():
text = p.read_text()
for needle in ['knowledge_base_type', 'managedSearchConfiguration', 'vectorSearchConfiguration', '"options": []']:
print(needle, text.find(needle))
PYRepository: langflow-ai/langflow
Length of output: 522
Wire knowledge_base_type through retrieval
- The dropdown has no
options, so MANAGED/VECTOR can’t be selected. retrieve()always usesmanagedSearchConfiguration, so VECTOR knowledge bases will fail.
Fix src/lfx/src/lfx/components/aws/bedrock_knowledge_base_retriever.py and regenerate src/lfx/src/lfx/_assets/component_index.json.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/lfx/src/lfx/_assets/component_index.json` at line 7571, Update
BedrockKnowledgeBaseRetrieverComponent so knowledge_base_type offers MANAGED and
VECTOR options and retrieve() builds the appropriate retrieval configuration,
using managedSearchConfiguration for MANAGED and vectorSearchConfiguration for
VECTOR. Ensure the selected type is also honored by agentic retrieval
configuration, then regenerate the component index asset from the updated
component source.
| DropdownInput( | ||
| name="knowledge_base_type", | ||
| display_name="Knowledge Base Type", | ||
| value="MANAGED", | ||
| ), |
There was a problem hiding this comment.
🎯 Functional Correctness | 🔴 Critical | ⚡ Quick win
Fix missing options and add visual configuration for agentic retrieval.
The DropdownInput for knowledge_base_type is missing its options argument, which will result in an empty, unusable dropdown in the UI.
Additionally, the PR summary explicitly mentions visual configuration for agentic retrieval, but the input is missing here (and currently falls back to an environment variable). Add BoolInput to support this behavior directly in the Langflow UI.
🐛 Proposed fix
DropdownInput(
name="knowledge_base_type",
display_name="Knowledge Base Type",
+ options=["MANAGED", "VECTOR"],
value="MANAGED",
+ info="Type of the Knowledge Base.",
),
+ BoolInput(
+ name="use_agentic_retrieval",
+ display_name="Use Agentic Retrieval",
+ info="Enable agentic retrieval with fallback to standard retrieval.",
+ value=True,
+ advanced=True,
+ ),📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| DropdownInput( | |
| name="knowledge_base_type", | |
| display_name="Knowledge Base Type", | |
| value="MANAGED", | |
| ), | |
| DropdownInput( | |
| name="knowledge_base_type", | |
| display_name="Knowledge Base Type", | |
| options=["MANAGED", "VECTOR"], | |
| value="MANAGED", | |
| info="Type of the Knowledge Base.", | |
| ), | |
| BoolInput( | |
| name="use_agentic_retrieval", | |
| display_name="Use Agentic Retrieval", | |
| info="Enable agentic retrieval with fallback to standard retrieval.", | |
| value=True, | |
| advanced=True, | |
| ), |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/lfx/src/lfx/components/aws/bedrock_knowledge_base_retriever.py` around
lines 56 - 60, Update the knowledge_base_type DropdownInput in the Bedrock
knowledge-base retriever to provide its supported dropdown options, and add a
BoolInput for configuring agentic retrieval directly in the UI. Ensure the
retriever uses this input rather than relying solely on the environment-variable
fallback, while preserving the existing managed knowledge-base behavior.
| """Retrieve documents using langchain-aws AmazonKnowledgeBasesRetriever.""" | ||
| import os | ||
|
|
||
| use_agentic_retrieval = os.environ.get("USE_AGENTIC_RETRIEVAL", "true").lower() != "false" |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Read the agentic retrieval toggle from the component property.
To honor the visual configuration in the Langflow properties panel, use the component instance attribute rather than relying exclusively on an environment variable.
🐛 Proposed fix
- use_agentic_retrieval = os.environ.get("USE_AGENTIC_RETRIEVAL", "true").lower() != "false"
+ use_agentic_retrieval = self.use_agentic_retrieval📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| use_agentic_retrieval = os.environ.get("USE_AGENTIC_RETRIEVAL", "true").lower() != "false" | |
| use_agentic_retrieval = self.use_agentic_retrieval |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/lfx/src/lfx/components/aws/bedrock_knowledge_base_retriever.py` at line
95, Update the agentic retrieval toggle in the Bedrock knowledge base retriever
component to read from the component’s configured instance property instead of
exclusively from the USE_AGENTIC_RETRIEVAL environment variable. Preserve the
existing boolean behavior while honoring the value selected in the Langflow
properties panel.
| # Build retrieval config | ||
| retrieval_config = {"managedSearchConfiguration": {"numberOfResults": self.number_of_results}} |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
file='src/lfx/src/lfx/components/aws/bedrock_knowledge_base_retriever.py'
echo '--- file outline ---'
ast-grep outline "$file" --view expanded || true
echo
echo '--- relevant lines ---'
nl -ba "$file" | sed -n '1,260p'
echo
echo '--- search for knowledge_base_type / retrieval config references ---'
rg -n "knowledge_base_type|managedSearchConfiguration|vectorSearchConfiguration|agenticRetrieveConfiguration|foundationModelType|rerankingModelType" src/lfx/src/lfx/components/aws -SRepository: langflow-ai/langflow
Length of output: 735
🌐 Web query:
AWS Bedrock Retrieve API vectorSearchConfiguration managedSearchConfiguration knowledge base type VECTOR docs
💡 Result:
The Amazon Bedrock Retrieve API allows you to query a knowledge base to retrieve relevant information [1][2]. Within the Retrieve API request body, the retrievalConfiguration field is used to specify how the search is performed [3][4]. This field maps to a KnowledgeBaseRetrievalConfiguration object, which supports two primary configuration types [5][3]: 1. vectorSearchConfiguration: This object is used for traditional vector search queries [5]. It allows you to define parameters such as the number of results (numberOfResults) and the search type (overrideSearchType, e.g., HYBRID or SEMANTIC) [3]. It also supports implicit filter configurations [3]. 2. managedSearchConfiguration: This object provides configurations for managed search [6]. Users are advised to use this configuration in the retrievalConfiguration to leverage managed search capabilities instead of relying solely on vectorSearchConfiguration [5]. When creating a knowledge base, the type is specified as VECTOR in the knowledgeBaseConfiguration field, which requires providing a VectorKnowledgeBaseConfiguration object, including details such as the embedding model ARN [7]. The Retrieve API itself does not require a "type" field in the retrieval request; rather, the configuration objects (vectorSearchConfiguration or managedSearchConfiguration) define how the retrieval operation is executed against the vector-type knowledge base [5][3][4].
Citations:
- 1: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent-runtime_Retrieve.html
- 2: https://docs.aws.amazon.com/bedrock/latest/userguide/kb-how-retrieval.html
- 3: https://docs.aws.amazon.com/bedrock/latest/userguide/kb-test-config.html
- 4: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent-runtime_KnowledgeBaseRetrievalConfiguration.html
- 5: https://docs.aws.amazon.com/bedrock/latest/userguide/kb-test-retrieve.html
- 6: https://docs.aws.amazon.com/boto3/latest/reference/services/bedrock-agent-runtime/client/retrieve.html
- 7: https://docs.aws.amazon.com/bedrock/latest/userguide/knowledge-base-create.html
🏁 Script executed:
#!/bin/bash
set -euo pipefail
file='src/lfx/src/lfx/components/aws/bedrock_knowledge_base_retriever.py'
echo '--- file lines 1-260 ---'
sed -n '1,260p' "$file" | cat -n
echo
echo '--- exact matches in file ---'
grep -nE 'knowledge_base_type|managedSearchConfiguration|vectorSearchConfiguration|agenticRetrieveConfiguration|foundationModelType|rerankingModelType|retrieval_config' "$file" || trueRepository: langflow-ai/langflow
Length of output: 9400
🏁 Script executed:
#!/bin/bash
set -euo pipefail
rg -n "AmazonKnowledgeBasesRetriever|knowledge_base_type|managedSearchConfiguration|vectorSearchConfiguration|agenticRetrieveConfiguration" src/lfx/src/lfx -SRepository: langflow-ai/langflow
Length of output: 9957
🌐 Web query:
site:docs.aws.amazon.com Bedrock managedSearchConfiguration vectorSearchConfiguration VECTOR knowledge base retrieve
💡 Result:
When using the Amazon Bedrock Retrieve API to query a knowledge base, the retrievalConfiguration field allows you to specify how the search is performed via two primary options: managedSearchConfiguration and vectorSearchConfiguration [1][2]. ManagedSearchConfiguration When you use managedSearchConfiguration, Amazon Bedrock automatically determines the optimal search strategy for you [3]. This is a simplified approach where the service handles the underlying search logic based on the knowledge base configuration [3]. VectorSearchConfiguration When you use vectorSearchConfiguration, you have more direct control over the search process [4]. This configuration is typically used to customize how results are retrieved and filtered [1]. Key parameters include: 1. overrideSearchType: Allows you to specify the search method. You can choose HYBRID (using both vector embeddings and raw text) or SEMANTIC (using only vector embeddings) [1][4]. Note that HYBRID search is generally available when using an Amazon OpenSearch Serverless vector store that supports filterable text fields [4]. 2. numberOfResults: Specifies the number of text chunks to retrieve [1][4]. 3. filter: Allows you to define filters to apply to metadata fields in your knowledge base data sources before results are returned [4]. 4. implicitFilterConfiguration: Enables the system to automatically apply filters based on query context [1][4]. 5. rerankingConfiguration: Used to improve search relevance by reordering initial results using sophisticated models [4]. In practice, you include either managedSearchConfiguration or vectorSearchConfiguration within the retrievalConfiguration object in your Retrieve or RetrieveAndGenerate API request, depending on whether you want the system to manage the strategy or if you require specific manual overrides [5][1][2][6].
Citations:
- 1: https://docs.aws.amazon.com/bedrock/latest/userguide/kb-test-config.html
- 2: https://docs.aws.amazon.com/botocore/latest/reference/services/bedrock-agent-runtime/paginator/Retrieve.html
- 3: https://docs.aws.amazon.com/boto3/latest/reference/services/bedrock-agent-runtime/client/retrieve.html
- 4: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_KnowledgeBaseVectorSearchConfiguration.html
- 5: https://docs.aws.amazon.com/bedrock/latest/userguide/kb-test-retrieve.html
- 6: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent-runtime_RetrieveAndGenerate.html
🌐 Web query:
langchain-aws AmazonKnowledgeBasesRetriever retrieval_config managedSearchConfiguration vectorSearchConfiguration
💡 Result:
In the langchain-aws library, the AmazonKnowledgeBasesRetriever uses a retrieval_config parameter to control how queries are performed against Amazon Bedrock Knowledge Bases [1][2]. This configuration is structured to support different search strategies depending on the type of knowledge base being accessed [3]. The retrieval_config accepts either a dictionary or a Pydantic object that typically includes one of two primary search configuration fields: 1. vectorSearchConfiguration: This is used for standard Amazon Bedrock Knowledge Bases that utilize vector stores [4][5]. It allows you to configure parameters such as numberOfResults, overrideSearchType (HYBRID or SEMANTIC), and rerankingConfiguration [6]. 2. managedSearchConfiguration: This was introduced to support Bedrock Managed Knowledge Bases [3]. It is specifically required when interacting with managed knowledge base types, as these do not support the vectorSearchConfiguration parameters [3][7]. If you attempt to use vectorSearchConfiguration with a managed knowledge base, you may encounter errors because the underlying service is incompatible with vector-specific search parameters [7]. When configuring your retriever, ensure that the retrieval_config object matches the capabilities and type of the knowledge base you are targeting [3]. Example configuration pattern: retriever = AmazonKnowledgeBasesRetriever( knowledge_base_id="", retrieval_config={ "vectorSearchConfiguration": { "numberOfResults": 4 } }) For managed knowledge bases, you would substitute vectorSearchConfiguration with managedSearchConfiguration in the retrieval_config dictionary [3].
Citations:
- 1: https://reference.langchain.com/python/langchain-aws/retrievers
- 2: https://github.com/langchain-ai/langchain-aws/blob/a9fd40a5e8ffc9750b6b8f624a1e299fa1cb52fd/libs/aws/langchain_aws/retrievers/bedrock.py
- 3: langchain-ai/langchain-aws@8771c55
- 4: https://docs.aws.amazon.com/bedrock/latest/userguide/kb-test-config.html
- 5: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent-runtime_KnowledgeBaseRetrievalConfiguration.html
- 6: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_KnowledgeBaseVectorSearchConfiguration.html
- 7: https://repost.aws/questions/QU58qnvOYPSTSRwKqo9cgXNQ/bedrock-agent-fails-when-querying-attached-managed-knowledge-base-with-error-vectorsearchconfiguration-is-not-supported-for-managed-knowledge-bases
🌐 Web query:
"AWS Bedrock managedSearchConfiguration VECTOR knowledge base" "vectorSearchConfiguration" "Retrieve"
💡 Result:
In the Amazon Bedrock Retrieve API, the retrievalConfiguration parameter allows you to specify how information is fetched from a knowledge base using either managedSearchConfiguration or vectorSearchConfiguration [1][2][3]. These two configurations serve distinct purposes: 1. Managed Search (managedSearchConfiguration): This is a simplified search interface provided by Amazon Bedrock [1][3]. It includes features such as filtering and reranking without requiring manual vector search tuning [2][3]. You use this when you want Bedrock to manage the underlying search optimization [4]. 2. Vector Search (vectorSearchConfiguration): This provides granular, low-level control over how vector search is performed [1][5]. It is used for custom search implementations where you define specific parameters for the vector retrieval process [4]. You must choose one or the other within the retrievalConfiguration object; they are not used simultaneously for a single retrieval request [4]. The managedSearchConfiguration object includes: - Filter: Allows for complex filtering using logical operators like andAll, orAll, equals, notEquals, in, notIn, greaterThan, lessThan, and string contains/starts with [2][3]. - NumberOfResults: Specifies the number of results to return [2][3]. - RerankingConfiguration: Allows you to configure how results are reranked, typically using Bedrock-provided models to improve relevance [2][3]. The vectorSearchConfiguration object is used to define specific search attributes such as the number of results, custom filter expressions, and override search configurations when you require manual control over the vector retrieval logic [5].
Citations:
- 1: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent-runtime_KnowledgeBaseRetrievalConfiguration.html
- 2: https://docs.aws.amazon.com/boto3/latest/reference/services/bedrock-agent-runtime/client/retrieve.html
- 3: https://docs.aws.amazon.com/cli/latest/reference/bedrock-agent-runtime/retrieve.html
- 4: https://docs.aws.amazon.com/bedrock/latest/userguide/kb-test-retrieve.html
- 5: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent-runtime_KnowledgeBaseVectorSearchConfiguration.html
🌐 Web query:
AWS Bedrock agenticRetrieveConfiguration foundationModelType rerankingModelType MANAGED VECTOR knowledge base
💡 Result:
In Amazon Bedrock, agentic retrieval is a feature that allows agents to perform multi-step query planning and evaluation when retrieving data from managed knowledge bases [1][2]. The agenticRetrieveConfiguration object is used to define how this process behaves [1][3]. Key components of agenticRetrieveConfiguration include: Foundation Model Configuration The foundationModelType parameter determines the model used for query planning and evaluation [4][3]. - MANAGED (Default): The service automatically selects an optimal foundation model [2][3]. - CUSTOM: You must provide a specific model configuration, including the model ARN, via the foundationModelConfiguration field [4][3]. Reranking Configuration The rerankingModelType parameter controls how search results are reordered to improve relevance [4][3]. - MANAGED (Default): A service-managed reranking model is used automatically [4][5]. - CUSTOM: You must provide your own reranking model configuration using the rerankingConfiguration field, which requires the model ARN [4][3]. - NONE: Reranking is disabled [4][3]. For managed knowledge bases, the default setting is MANAGED for both the foundation model and the reranking model [4][2]. When choosing the CUSTOM option for either, you are required to provide the specific model configuration [4][3]. This setup supports advanced retrieval workflows where you can tailor the agent's reasoning (via the foundation model) and the precision of the retrieved data (via the reranking model) to meet specific requirements [1][5].
Citations:
- 1: https://docs.aws.amazon.com/bedrock/latest/userguide/kb-test-agentic-retrieve.html
- 2: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent-runtime_AgenticRetrieveStream.html
- 3: https://docs.aws.amazon.com/bedrock-agentcore/latest/devguide/gateway-target-connector-managed-kb.html
- 4: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent-runtime_AgenticRetrieveConfiguration.html
- 5: https://docs.aws.amazon.com/bedrock/latest/userguide/kb-test-retrieve.html
🌐 Web query:
"AmazonKnowledgeBasesRetriever" "managedSearchConfiguration" "vectorSearchConfiguration" "langchain_aws"
💡 Result:
The AmazonKnowledgeBasesRetriever in the langchain_aws library allows you to interact with Amazon Bedrock Knowledge Bases by passing a retrieval_config object [1][2]. This configuration supports both vectorSearchConfiguration and managedSearchConfiguration, which are used to control how search queries are performed and how results are retrieved [3][4]. VectorSearchConfiguration is primarily used for standard vector-based retrieval, where you can specify parameters such as the number of results to return (numberOfResults) [1][2][5]. ManagedSearchConfiguration enables the use of managed search capabilities, such as advanced reranking [4]. You can instantiate the retriever and provide these configurations as follows [1][2]: from langchain_aws import AmazonKnowledgeBasesRetriever # Example using vectorSearchConfiguration retriever = AmazonKnowledgeBasesRetriever( knowledge_base_id="", retrieval_config={ "vectorSearchConfiguration": { "numberOfResults": 4 } },) # Example using managedSearchConfiguration # Recent updates allow passing managedSearchConfiguration within retrieval_config retriever_managed = AmazonKnowledgeBasesRetriever( knowledge_base_id="", retrieval_config={ "managedSearchConfiguration": { # Managed search options, e.g., reranking configuration } },) The retrieval_config parameter accepts these configurations as Python dictionaries or corresponding Pydantic models (such as ManagedSearchConfig) [1][4]. These settings are passed through to the underlying Bedrock Agent Runtime API [1][3].
Citations:
- 1: https://github.com/langchain-ai/langchain-aws/blob/a9fd40a5e8ffc9750b6b8f624a1e299fa1cb52fd/libs/aws/langchain_aws/retrievers/bedrock.py
- 2: https://pypi.org/project/langchain-aws/0.2.31/
- 3: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent-runtime_KnowledgeBaseRetrievalConfiguration.html
- 4: feat: add Bedrock Managed Knowledge Base support langchain-ai/langchain-aws#1155
- 5: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent-runtime_KnowledgeBaseVectorSearchConfiguration.html
🌐 Web query:
site:docs.aws.amazon.com/bedrock/latest/APIReference/ KnowledgeBaseRetrievalConfiguration managedSearchConfiguration
💡 Result:
In Amazon Bedrock, KnowledgeBaseRetrievalConfiguration is a component used within retrieval requests (such as Retrieve or RetrieveAndGenerate) to define how information is fetched from a knowledge base [1][2][3]. The ManagedSearchConfiguration object is a field within the KnowledgeBaseRetrievalConfiguration structure [1]. Its purpose is to configure managed search, a feature that automatically determines the optimal search strategy based on the characteristics of your data store [4]. When using ManagedSearchConfiguration, you can provide parameters such as: - filter: Specifies criteria to filter the retrieved results [5]. - numberOfResults: Defines the number of results to return from the knowledge base [5]. By utilizing managed search, you offload the complexity of fine-tuning search strategies, allowing the system to handle the retrieval process automatically [4].
Citations:
- 1: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent-runtime_KnowledgeBaseRetrievalConfiguration.html
- 2: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent-runtime_Retrieve.html
- 3: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent-runtime_KnowledgeBaseRetrieveAndGenerateConfiguration.html
- 4: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent-runtime_ManagedSearchConfiguration.html
- 5: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent-runtime_RetrieveAndGenerate.html
🌐 Web query:
site:docs.aws.amazon.com/bedrock/latest/userguide/kb-test-config.html managedSearchConfiguration vectorSearchConfiguration
💡 Result:
In Amazon Bedrock, managedSearchConfiguration and vectorSearchConfiguration are distinct retrieval configurations used when querying knowledge bases [1]. managedSearchConfiguration This configuration is used for managed knowledge bases [1]. It simplifies the retrieval process by managing reranking behavior automatically by default [1]. Key features include: - Reranking control: You can set rerankingModelType to MANAGED (the default), CUSTOM (to specify a custom Bedrock reranking model), or NONE (to disable reranking) [1]. - Metadata filtering: Supports various filters such as equals, andAll, orAll, greaterThan, and more to refine query results [2][3]. - Simplified usage: It is designed to be used in the retrievalConfiguration instead of vectorSearchConfiguration when working with managed knowledge bases [1]. vectorSearchConfiguration This configuration is used for custom knowledge bases [1]. It provides granular control over how the search is performed, including: - Search strategy: Allows setting the overrideSearchType to HYBRID (combining vector and keyword search) or SEMANTIC (vector search only) for supported vector stores [4][5][6]. - Implicit filtering: Includes settings for implicit filtering through the implicitFilterConfiguration field [7][4]. - Reranking: Supports detailed reranking configurations via the VectorSearchRerankingConfiguration object [7][4]. When using the Retrieve or RetrieveAndGenerate API operations, you choose the configuration based on your knowledge base type; managedSearchConfiguration is for managed knowledge bases, while vectorSearchConfiguration is intended for custom knowledge bases [1][4][6].
Citations:
- 1: https://docs.aws.amazon.com/bedrock/latest/userguide/kb-test-retrieve.html
- 2: https://docs.aws.amazon.com/bedrock-agentcore/latest/devguide/gateway-target-connector-managed-kb.html
- 3: https://docs.aws.amazon.com/boto3/latest/reference/services/bedrock-agent-runtime/client/retrieve.html
- 4: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent-runtime_KnowledgeBaseVectorSearchConfiguration.html
- 5: https://docs.aws.amazon.com/java/api/latest/software/amazon/awssdk/services/bedrockagentruntime/model/KnowledgeBaseVectorSearchConfiguration.html
- 6: https://docs.aws.amazon.com/AWSJavaSDK/latest/javadoc/com/amazonaws/services/bedrockagentruntime/model/KnowledgeBaseVectorSearchConfiguration.html
- 7: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent-runtime_Retrieve.html
🌐 Web query:
site:docs.aws.amazon.com agenticRetrieveConfiguration foundationModelType rerankingModelType managed vector knowledge base
💡 Result:
In Amazon Bedrock, agenticRetrieveConfiguration is a structural component used within the AgenticRetrieveStream API to define how an agent interacts with a knowledge base, specifically managing model selection for query planning, evaluation, and result reranking [1][2]. Foundation Model Type The foundationModelType specifies the model used for query planning and evaluation [3]. It supports two primary options [2]: MANAGED: Uses a service-default model optimized by Amazon Bedrock [4][1]. CUSTOM: Requires you to provide a specific Bedrock foundation model ARN via foundationModelConfiguration [4][2]. Reranking Model Type The rerankingModelType controls the model used to rerank retrieved search results to improve relevance [3][2]. It supports three options [4][2]: MANAGED: Uses the service-managed reranker (the default for managed knowledge bases) [4][1]. CUSTOM: Allows the use of a custom reranking model by providing the model ARN in the rerankingConfiguration [4][2]. NONE: Disables reranking entirely (the default for custom embedding knowledge bases) [4][5]. Managed Vector Knowledge Bases When using a Bedrock Managed Knowledge Base, the infrastructure—including data ingestion, storage, and indexing—is fully managed by the service [6]. Agentic retrieval is a feature natively supported for these managed environments, allowing for sophisticated multi-step retrieval processes [6]. When configured, the service automatically defaults to managed models for both foundation tasks and reranking unless you explicitly override them with custom configurations [4][1][5]. Additional parameters, such as maxAgentIteration and policyConfiguration (for guardrails), can also be defined within the agenticRetrieveConfiguration to further refine the agent's retrieval behavior [1][2].
Citations:
- 1: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent-runtime_AgenticRetrieveStream.html
- 2: https://docs.aws.amazon.com/bedrock-agentcore/latest/devguide/gateway-target-connector-managed-kb.html
- 3: https://docs.aws.amazon.com/bedrock/latest/userguide/kb-test-agentic-retrieve.html
- 4: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent-runtime_AgenticRetrieveConfiguration.html
- 5: https://docs.aws.amazon.com/boto3/latest/reference/services/bedrock-agent-runtime/client/agentic_retrieve_stream.html
- 6: https://docs.aws.amazon.com/bedrock/latest/userguide/kb-build-managed.html
🌐 Web query:
site:docs.aws.amazon.com/bedrock latest agentic_retrieve_stream foundationModelType rerankingModelType
💡 Result:
In the context of the Amazon Bedrock AgenticRetrieveStream API, foundationModelType and rerankingModelType are configuration fields used to define the intelligence and ranking components of the agentic retrieval process [1][2]. foundationModelType This field specifies the type of foundation model used by the agent for query planning, reasoning, and evaluation during the retrieval process [2][3][4]. Valid values include: - MANAGED: The service automatically selects the optimal foundation model [1][2]. - CUSTOM: You specify a custom foundation model configuration (e.g., providing a model ARN) [1][2]. rerankingModelType This field specifies the type of reranking model used to assess and reorder the relevance of retrieved documents [1][4]. Valid values include: - MANAGED: The service automatically uses a managed reranking model [1][5]. - CUSTOM: You specify a custom reranking model configuration (e.g., providing a model ARN) [1][5]. - NONE: No reranking is applied [1]. If not specified, the system defaults to MANAGED for managed embedding knowledge bases and NONE for custom embedding knowledge bases [1]. When using either CUSTOM type, you must also provide the corresponding configuration (foundationModelConfiguration or rerankingConfiguration) containing the specific model details, such as the model ARN [1][4][6]. Agentic retrieval is designed specifically for managed Amazon Bedrock knowledge bases [4].
Citations:
- 1: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent-runtime_AgenticRetrieveConfiguration.html
- 2: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent-runtime_AgenticRetrieveStream.html
- 3: https://docs.aws.amazon.com/bedrock/latest/userguide/kb-how-retrieval.html
- 4: https://docs.aws.amazon.com/bedrock/latest/userguide/kb-test-agentic-retrieve.html
- 5: https://docs.aws.amazon.com/bedrock/latest/userguide/kb-test-retrieve.html
- 6: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent-runtime_AgenticRetrieveBedrockRerankingModelConfiguration.html
🌐 Web query:
langchain_aws AmazonKnowledgeBasesRetriever retrieval_config vectorSearchConfiguration managedSearchConfiguration docs
💡 Result:
The AmazonKnowledgeBasesRetriever in the langchain-aws library allows you to interact with Amazon Bedrock Knowledge Bases for retrieval-augmented generation (RAG) [1][2][3]. When configuring this retriever, the retrieval_config parameter is used to define how search is performed [4][5]. This parameter typically accepts a dictionary that matches the structure expected by the Amazon Bedrock Retrieve API [4][6]. The two primary configuration types within the retrieval_config are: 1. vectorSearchConfiguration: This is used for standard vector-based search [7][6]. It allows you to specify parameters such as numberOfResults (the number of source chunks to retrieve) and overrideSearchType (to choose between HYBRID or SEMANTIC search) [7]. 2. managedSearchConfiguration: This is designed for managed knowledge bases [6][8]. It is important to note that managed knowledge bases typically do not support vectorSearchConfiguration [8]. If you are working with a managed knowledge base, you must use managedSearchConfiguration instead, as attempting to use vectorSearchConfiguration may result in a validation error [8]. Example Usage: from langchain_aws import AmazonKnowledgeBasesRetriever # Standard vector search configuration retriever = AmazonKnowledgeBasesRetriever( knowledge_base_id="", retrieval_config={ "vectorSearchConfiguration": { "numberOfResults": 4, "overrideSearchType": "HYBRID" } },) # Example for managed search (if applicable to your knowledge base type) # retriever = AmazonKnowledgeBasesRetriever( # knowledge_base_id="", # retrieval_config={ # "managedSearchConfiguration": { # "someOption": "value" # } # }, #) For more details on specific fields within these configurations, refer to the official Amazon Bedrock API documentation for KnowledgeBaseVectorSearchConfiguration and KnowledgeBaseRetrievalConfiguration [7][6].
Citations:
- 1: https://reference.langchain.com/python/langchain-aws/retrievers/bedrock/AmazonKnowledgeBasesRetriever
- 2: https://docs.langchain.com/oss/python/integrations/retrievers/bedrock
- 3: https://docs.langchain.com/oss/python/integrations/providers/aws
- 4: https://github.com/langchain-ai/langchain-aws/blob/a9fd40a5e8ffc9750b6b8f624a1e299fa1cb52fd/libs/aws/langchain_aws/retrievers/bedrock.py
- 5: https://reference.langchain.com/python/langchain-community/retrievers/bedrock/AmazonKnowledgeBasesRetriever
- 6: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent-runtime_KnowledgeBaseRetrievalConfiguration.html
- 7: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent-runtime_KnowledgeBaseVectorSearchConfiguration.html
- 8: https://repost.aws/questions/QU58qnvOYPSTSRwKqo9cgXNQ/bedrock-agent-fails-when-querying-attached-managed-knowledge-base-with-error-vectorsearchconfiguration-is-not-supported-for-managed-knowledge-bases
Switch retrieval config by knowledge base type
knowledge_base_type is ignored here. Build retrieval_config from it so VECTOR uses vectorSearchConfiguration and MANAGED keeps managedSearchConfiguration. The agentic retrieval path should also be gated the same way, since it’s intended for managed knowledge bases.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/lfx/src/lfx/components/aws/bedrock_knowledge_base_retriever.py` around
lines 152 - 153, Update the retrieval configuration logic around
knowledge_base_type so VECTOR builds vectorSearchConfiguration and MANAGED
retains managedSearchConfiguration, using the existing number_of_results value.
Apply the same knowledge-base-type gate to the agentic retrieval path so it only
runs for MANAGED knowledge bases.
| @patch("langchain_aws.retrievers.AmazonKnowledgeBasesRetriever") | ||
| @patch("langchain_aws.retrievers.AmazonKnowledgeBasesRetriever") | ||
| def test_retrieve_with_credentials(mock_retriever_class: MagicMock) -> None: |
There was a problem hiding this comment.
🎯 Functional Correctness | 🔴 Critical | ⚡ Quick win
Remove the duplicate @patch decorator.
The @patch("langchain_aws.retrievers.AmazonKnowledgeBasesRetriever") decorator is applied twice, which will cause unittest.mock to pass two mock objects to the test function. Since the function signature only accepts one argument (mock_retriever_class), this will result in a runtime TypeError when the test is executed.
🐛 Proposed fix to remove the duplicate patch
`@patch`("langchain_aws.retrievers.AmazonKnowledgeBasesRetriever")
-@patch("langchain_aws.retrievers.AmazonKnowledgeBasesRetriever")
def test_retrieve_with_credentials(mock_retriever_class: MagicMock) -> None:📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| @patch("langchain_aws.retrievers.AmazonKnowledgeBasesRetriever") | |
| @patch("langchain_aws.retrievers.AmazonKnowledgeBasesRetriever") | |
| def test_retrieve_with_credentials(mock_retriever_class: MagicMock) -> None: | |
| `@patch`("langchain_aws.retrievers.AmazonKnowledgeBasesRetriever") | |
| def test_retrieve_with_credentials(mock_retriever_class: MagicMock) -> None: |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/lfx/tests/unit/components/aws/test_bedrock_knowledge_base_retriever.py`
around lines 95 - 97, Remove the duplicate `@patch` decorator above
test_retrieve_with_credentials, leaving a single patch for
AmazonKnowledgeBasesRetriever so the test receives only the declared
mock_retriever_class argument.
1 participant
feat: add Amazon Bedrock Knowledge Base retriever component
Summary
Adds a new Langflow component that retrieves documents from Amazon Bedrock Knowledge Bases. Wraps
langchain-awsAmazonKnowledgeBasesRetriever as a dedicated Langflow node.Features
Configuration
Required IAM Permissions
bedrock:Retrievebedrock:AgenticRetrieveTesting
boto3 >= 1.43Files Changed
src/lfx/src/lfx/components/aws/bedrock_knowledge_base_retriever.py— componentsrc/lfx/src/lfx/components/aws/test_bedrock_knowledge_base_retriever.py— testssrc/lfx/src/lfx/components/aws/BEDROCK_MANAGED_KB.md— design docSummary by CodeRabbit
New Features
Documentation
Tests