diff --git a/src/lfx/src/lfx/_assets/component_index.json b/src/lfx/src/lfx/_assets/component_index.json index 708c4493e4b3..5d8cd614b6a7 100644 --- a/src/lfx/src/lfx/_assets/component_index.json +++ b/src/lfx/src/lfx/_assets/component_index.json @@ -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", { @@ -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" } diff --git a/src/lfx/src/lfx/components/aws/BEDROCK_MANAGED_KB.md b/src/lfx/src/lfx/components/aws/BEDROCK_MANAGED_KB.md new file mode 100644 index 000000000000..085041244cdf --- /dev/null +++ b/src/lfx/src/lfx/components/aws/BEDROCK_MANAGED_KB.md @@ -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:::knowledge-base/" +} +``` diff --git a/src/lfx/src/lfx/components/aws/__init__.py b/src/lfx/src/lfx/components/aws/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/src/lfx/src/lfx/components/aws/bedrock_knowledge_base_retriever.py b/src/lfx/src/lfx/components/aws/bedrock_knowledge_base_retriever.py new file mode 100644 index 000000000000..abb3a2f05a39 --- /dev/null +++ b/src/lfx/src/lfx/components/aws/bedrock_knowledge_base_retriever.py @@ -0,0 +1,191 @@ +"""Amazon Bedrock Knowledge Base Retriever component for Langflow. + +Exposes the langchain-aws AmazonKnowledgeBasesRetriever as a dedicated Langflow node. +""" + +from langflow.custom import Component +from langflow.io import BoolInput, IntInput, MessageTextInput, Output, SecretStrInput +from langflow.schema import Data + + +def _get_source_uri(result: dict) -> str: + """Extract source URI from a retrieval result, handling all location types.""" + location = result.get("location", {}) + loc_type = location.get("type", "") + if loc_type == "S3" or "s3Location" in location: + return location.get("s3Location", {}).get("uri", "") + if loc_type == "WEB" or "webLocation" in location: + return location.get("webLocation", {}).get("url", "") + if "confluenceLocation" in location: + return location.get("confluenceLocation", {}).get("url", "") + if "salesforceLocation" in location: + return location.get("salesforceLocation", {}).get("url", "") + if "sharePointLocation" in location: + return location.get("sharePointLocation", {}).get("url", "") + if "customDocumentLocation" in location: + return location.get("customDocumentLocation", {}).get("id", "") + # Fallback to metadata._source_uri (for agentic results) + return result.get("metadata", {}).get("_source_uri", "") + + +class BedrockKnowledgeBaseRetrieverComponent(Component): + """Retrieves documents from an Amazon Bedrock Knowledge Base. + + Wraps langchain-aws's AmazonKnowledgeBasesRetriever as a dedicated Langflow node. + Supports both Managed Knowledge Bases (recommended) and Vector search. + """ + + display_name = "Amazon Bedrock Knowledge Base" + description = "Retrieve documents using langchain-aws AmazonKnowledgeBasesRetriever with managed or vector search." + icon = "Amazon" + name = "BedrockKnowledgeBaseRetriever" + + inputs = [ + MessageTextInput( + name="knowledge_base_id", + display_name="Knowledge Base ID", + info="The ID of the Amazon Bedrock Knowledge Base (10 alphanumeric characters).", + required=True, + ), + MessageTextInput( + name="query", + display_name="Query", + info="The search query to retrieve relevant documents.", + required=True, + ), + BoolInput( + name="use_agentic_retrieval", + display_name="Use Agentic Retrieval", + value=True, + info=( + "If enabled, tries AgenticRetrieveStream first (query decomposition + managed reranking)" + " with fallback to standard Retrieve." + ), + ), + MessageTextInput( + name="region_name", + display_name="AWS Region", + value="us-east-1", + info="AWS region where the Knowledge Base is located.", + ), + IntInput( + name="number_of_results", + display_name="Number of Results", + value=5, + info="Maximum number of results to return.", + ), + SecretStrInput( + name="aws_access_key_id", + display_name="AWS Access Key ID", + info="AWS access key. Optional if using IAM role or environment credentials.", + required=False, + ), + SecretStrInput( + name="aws_secret_access_key", + display_name="AWS Secret Access Key", + info="AWS secret key. Optional if using IAM role or environment credentials.", + required=False, + ), + ] + + outputs = [ + Output(display_name="Retrieved Documents", name="documents", method="retrieve"), + ] + + def retrieve(self) -> list[Data]: + """Retrieve documents using langchain-aws AmazonKnowledgeBasesRetriever.""" + use_agentic_retrieval = self.use_agentic_retrieval + + # Try agentic retrieval first + if use_agentic_retrieval: + try: + import boto3 + from botocore.config import Config + + # Build credentials kwargs for boto3 + boto_kwargs: dict = {"region_name": self.region_name} + if self.aws_access_key_id and self.aws_secret_access_key: + boto_kwargs["aws_access_key_id"] = self.aws_access_key_id + boto_kwargs["aws_secret_access_key"] = self.aws_secret_access_key + + client = boto3.client( + "bedrock-agent-runtime", + config=Config(user_agent_extra="langflow/bedrock-kb"), + **boto_kwargs, + ) + response = client.agentic_retrieve_stream( + messages=[{"content": {"text": self.query}, "role": "user"}], + retrievers=[ + { + "configuration": { + "knowledgeBase": { + "knowledgeBaseId": self.knowledge_base_id, + "retrievalOverrides": {"maxNumberOfResults": self.number_of_results}, + } + } + } + ], + agenticRetrieveConfiguration={ + "foundationModelType": "MANAGED", + "rerankingModelType": "MANAGED", + }, + ) + results = [] + for event in response.get("stream", []): + if "result" in event and "results" in event["result"]: + results.extend( + Data( + text=result.get("content", {}).get("text", ""), + data={ + "source": _get_source_uri(result), + "score": result.get("score", 0.0), + "knowledge_base_id": self.knowledge_base_id, + }, + ) + for result in event["result"]["results"] + ) + if results: + return results + except Exception as e: # noqa: BLE001 + import logging + + logging.getLogger(__name__).debug("Agentic retrieval unavailable, falling back: %s", e) + + try: + from langchain_aws.retrievers import AmazonKnowledgeBasesRetriever + except ImportError as err: + msg = "langchain-aws is required. Install with: pip install langchain-aws>=0.2.0" + raise ImportError(msg) from err + + # Build retrieval config + retrieval_config = {"managedSearchConfiguration": {"numberOfResults": self.number_of_results}} + # Build credentials kwargs + credentials_kwargs = {} + if self.aws_access_key_id and self.aws_secret_access_key: + credentials_kwargs["credentials_profile_name"] = None + credentials_kwargs["aws_access_key_id"] = self.aws_access_key_id + credentials_kwargs["aws_secret_access_key"] = self.aws_secret_access_key + + retriever = AmazonKnowledgeBasesRetriever( + knowledge_base_id=self.knowledge_base_id, + region_name=self.region_name, + retrieval_config=retrieval_config, + **credentials_kwargs, + ) + + docs = retriever.invoke(self.query) + + results = [] + for doc in docs: + results.append( + Data( + text=doc.page_content, + data={ + "source": doc.metadata.get("source", ""), + "score": doc.metadata.get("score", 0.0), + "knowledge_base_id": self.knowledge_base_id, + }, + ) + ) + + return results diff --git a/src/lfx/tests/unit/components/aws/__init__.py b/src/lfx/tests/unit/components/aws/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/src/lfx/tests/unit/components/aws/test_bedrock_knowledge_base_retriever.py b/src/lfx/tests/unit/components/aws/test_bedrock_knowledge_base_retriever.py new file mode 100644 index 000000000000..d0feca7cc701 --- /dev/null +++ b/src/lfx/tests/unit/components/aws/test_bedrock_knowledge_base_retriever.py @@ -0,0 +1,114 @@ +"""Tests for the Amazon Bedrock Knowledge Base Retriever Langflow component.""" + +from __future__ import annotations + +import importlib.util +import sys +import types +from pathlib import Path +from unittest.mock import MagicMock, patch + +# Mock langflow and langchain_aws modules before importing the component +mock_langflow = types.ModuleType("langflow") +mock_langflow.__path__ = [] +mock_langflow_custom = types.ModuleType("langflow.custom") +mock_langflow_io = types.ModuleType("langflow.io") +mock_langflow_schema = types.ModuleType("langflow.schema") + + +class MockComponent: + """Mock Component base class.""" + + +mock_langflow_custom.Component = MockComponent + + +class MockInput: + """Mock IO input class.""" + + def __init__(self, **kwargs: object) -> None: + for k, v in kwargs.items(): + setattr(self, k, v) + + +mock_langflow_io.DropdownInput = MockInput +mock_langflow_io.IntInput = MockInput +mock_langflow_io.MessageTextInput = MockInput +mock_langflow_io.Output = MockInput +mock_langflow_io.SecretStrInput = MockInput + + +class MockData: + """Mock Data class.""" + + def __init__(self, text: str = "", data: dict | None = None) -> None: + self.text = text + self.data = data or {} + + +mock_langflow_schema.Data = MockData + +# Register all mock modules +sys.modules["langflow"] = mock_langflow +sys.modules["langflow.custom"] = mock_langflow_custom +sys.modules["langflow.io"] = mock_langflow_io +sys.modules["langflow.schema"] = mock_langflow_schema + +# Import component via relative path +_component_path = ( + Path(__file__).resolve().parents[4] / "src" / "lfx" / "components" / "aws" / "bedrock_knowledge_base_retriever.py" +) +spec = importlib.util.spec_from_file_location("bedrock_knowledge_base_retriever", _component_path) +mod = importlib.util.module_from_spec(spec) +spec.loader.exec_module(mod) +BedrockKnowledgeBaseRetrieverComponent = mod.BedrockKnowledgeBaseRetrieverComponent + + +@patch("langchain_aws.retrievers.AmazonKnowledgeBasesRetriever") +def test_retrieve_with_managed_config(mock_retriever_class: MagicMock) -> None: + """Test retrieval uses langchain-aws retriever with managed config.""" + mock_retriever = MagicMock() + mock_retriever.invoke.return_value = [ + MagicMock(page_content="Document 1", metadata={"source": "s3://b/1", "score": 0.9}), + MagicMock(page_content="Document 2", metadata={"source": "s3://b/2", "score": 0.8}), + ] + mock_retriever_class.return_value = mock_retriever + + component = BedrockKnowledgeBaseRetrieverComponent() + component.knowledge_base_id = "TEST123456" + component.query = "What is managed KB?" + component.region_name = "us-west-2" + component.number_of_results = 5 + component.aws_access_key_id = "" + component.aws_secret_access_key = "" + + results = component.retrieve() + + mock_retriever_class.assert_called_once() + call_kwargs = mock_retriever_class.call_args.kwargs + assert call_kwargs["knowledge_base_id"] == "TEST123456" + assert "managedSearchConfiguration" in call_kwargs["retrieval_config"] + assert len(results) == 2 + assert results[0].text == "Document 1" + + +@patch("langchain_aws.retrievers.AmazonKnowledgeBasesRetriever") +def test_retrieve_with_credentials(mock_retriever_class: MagicMock) -> None: + """Test credentials are passed when provided.""" + mock_retriever = MagicMock() + mock_retriever.invoke.return_value = [] + mock_retriever_class.return_value = mock_retriever + + component = BedrockKnowledgeBaseRetrieverComponent() + component.knowledge_base_id = "TEST123456" + component.query = "test" + component.region_name = "us-east-1" + component.number_of_results = 5 + component.aws_access_key_id = "AKID123" + component.aws_secret_access_key = "SECRET456" # noqa: S105 + + component.retrieve() + + call_kwargs = mock_retriever_class.call_args.kwargs + assert call_kwargs["aws_access_key_id"] == "AKID123" + assert call_kwargs["aws_secret_access_key"] == "SECRET456" # noqa: S105