amazon-bedrock-haystack 5.0.0__tar.gz → 5.2.0__tar.gz

{amazon_bedrock_haystack-5.0.0 → amazon_bedrock_haystack-5.2.0}/CHANGELOG.md

@@ -1,5 +1,17 @@
 # Changelog
 
+## [integrations/amazon_bedrock-v5.1.0] - 2025-09-29
+
+### 🚀 Features
+
+- S3Downloader - add `s3_key_generation_function` param to customize S3 key generation (#2343)
+
+## [integrations/amazon_bedrock-v5.0.0] - 2025-09-22
+
+### 🧹 Chores
+
+- [**breaking**] Remove deprecated `BedrockRanker` (use `AmazonBedrockRanker` instead) (#2287)
+
 ## [integrations/amazon_bedrock-v4.2.0] - 2025-09-19
 
 ### 🚀 Features

{amazon_bedrock_haystack-5.0.0 → amazon_bedrock_haystack-5.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: amazon-bedrock-haystack
-Version: 5.0.0
+Version: 5.2.0
 Summary: An integration of AWS S3 and Bedrock as a Downloader and Generator components.
 Project-URL: Documentation, https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/amazon_bedrock#readme
 Project-URL: Issues, https://github.com/deepset-ai/haystack-core-integrations/issues
@@ -21,7 +21,7 @@ Classifier: Programming Language :: Python :: Implementation :: PyPy
 Requires-Python: >=3.9
 Requires-Dist: aioboto3>=14.0.0
 Requires-Dist: boto3>=1.28.57
-Requires-Dist: haystack-ai>=2.17.1
+Requires-Dist: haystack-ai>=2.19.0
 Description-Content-Type: text/markdown
 
 # amazon-bedrock-haystack

amazon_bedrock_haystack-5.2.0/pydoc/config_docusaurus.yml

@@ -0,0 +1,41 @@
+loaders:
+- ignore_when_discovered:
+  - __init__
+  modules:
+  - haystack_integrations.common.amazon_bedrock.errors
+  - haystack_integrations.components.embedders.amazon_bedrock.document_embedder
+  - haystack_integrations.components.embedders.amazon_bedrock.text_embedder
+  - haystack_integrations.components.embedders.amazon_bedrock.document_image_embedder
+  - haystack_integrations.components.generators.amazon_bedrock.generator
+  - haystack_integrations.components.generators.amazon_bedrock.adapters
+  - haystack_integrations.common.amazon_bedrock.errors
+  - haystack_integrations.components.generators.amazon_bedrock.chat.chat_generator
+  - haystack_integrations.components.rankers.amazon_bedrock.ranker
+  - haystack_integrations.components.downloaders.s3.s3_downloader
+  - haystack_integrations.common.s3.utils
+  - haystack_integrations.common.s3.errors
+  search_path:
+  - ../src
+  type: haystack_pydoc_tools.loaders.CustomPythonLoader
+processors:
+- do_not_filter_modules: false
+  documented_only: true
+  expression: null
+  skip_empty_modules: true
+  type: filter
+- expression: name not in ['BedrockRanker']
+  type: filter
+- type: smart
+- type: crossref
+renderer:
+  description: Amazon Bedrock integration for Haystack
+  id: integrations-amazon-bedrock
+  markdown:
+    add_member_class_prefix: false
+    add_method_class_prefix: true
+    classdef_code_block: false
+    descriptive_class_title: false
+    descriptive_module_title: true
+    filename: amazon_bedrock.md
+  title: Amazon Bedrock
+  type: haystack_pydoc_tools.renderers.DocusaurusRenderer

{amazon_bedrock_haystack-5.0.0 → amazon_bedrock_haystack-5.2.0}/pyproject.toml

@@ -23,7 +23,7 @@ classifiers = [
   "Programming Language :: Python :: Implementation :: CPython",
   "Programming Language :: Python :: Implementation :: PyPy",
 ]
-dependencies = ["haystack-ai>=2.17.1", "boto3>=1.28.57", "aioboto3>=14.0.0"]
+dependencies = ["haystack-ai>=2.19.0", "boto3>=1.28.57", "aioboto3>=14.0.0"]
 
 [project.urls]
 Documentation = "https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/amazon_bedrock#readme"

{amazon_bedrock_haystack-5.0.0 → amazon_bedrock_haystack-5.2.0}/src/haystack_integrations/components/downloaders/s3/s3_downloader.py

@@ -5,12 +5,13 @@
 import os
 from concurrent.futures import ThreadPoolExecutor
 from pathlib import Path
-from typing import Any, Dict, List, Optional
+from typing import Any, Callable, Dict, List, Optional
 
 from botocore.config import Config
 from haystack import component, default_from_dict, default_to_dict, logging
 from haystack.dataclasses import Document
 from haystack.utils.auth import Secret, deserialize_secrets_inplace
+from haystack.utils.callable_serialization import deserialize_callable, serialize_callable
 
 from haystack_integrations.common.amazon_bedrock.utils import get_aws_session
 from haystack_integrations.common.s3.utils import S3Storage
@@ -41,6 +42,7 @@ class S3Downloader:
         file_name_meta_key: str = "file_name",
         max_workers: int = 32,
         max_cache_size: int = 100,
+        s3_key_generation_function: Optional[Callable[[Document], str]] = None,
     ) -> None:
         """
         Initializes the `S3Downloader` with the provided parameters.
@@ -64,9 +66,15 @@ class S3Downloader:
             By default, all file extensions are allowed.
         :param max_workers: The maximum number of workers to use for concurrent downloads.
         :param max_cache_size: The maximum number of files to cache.
-        :param file_name_meta_key: The name of the meta key that contains the file name to download.
+        :param file_name_meta_key: The name of the meta key that contains the file name to download. The file name
+            will also be used to create local file path for download.
             By default, the `Document.meta["file_name"]` is used. If you want to use a
             different key in `Document.meta`, you can set it here.
+        :param s3_key_generation_function: An optional function that generates the S3 key for the file to download.
+            If not provided, the default behavior is to use `Document.meta[file_name_meta_key]`.
+            The function must accept a `Document` object and return a string.
+            If the environment variable `S3_DOWNLOADER_PREFIX` is set, its value will be automatically
+            prefixed to the generated S3 key.
         :raises ValueError: If the `file_root_path` is not set through
             the constructor or the `FILE_ROOT_PATH` environment variable.
 
@@ -94,6 +102,7 @@ class S3Downloader:
         self.max_workers = max_workers
         self.max_cache_size = max_cache_size
         self.file_name_meta_key = file_name_meta_key
+        self.s3_key_generation_function = s3_key_generation_function
 
         self._storage: Optional[S3Storage] = None
 
@@ -186,8 +195,9 @@ class S3Downloader:
             file_path.touch()
 
         else:
+            s3_key = self.s3_key_generation_function(document) if self.s3_key_generation_function else file_name
             # we know that _storage is not None after warm_up() is called, but mypy does not know that
-            self._storage.download(key=file_name, local_file_path=file_path)  # type: ignore[union-attr]
+            self._storage.download(key=s3_key, local_file_path=file_path)  # type: ignore[union-attr]
 
         document.meta["file_path"] = str(file_path)
         return document
@@ -216,6 +226,11 @@ class S3Downloader:
 
     def to_dict(self) -> Dict[str, Any]:
         """Serialize the component to a dictionary."""
+
+        s3_key_generation_function_name = (
+            serialize_callable(self.s3_key_generation_function) if self.s3_key_generation_function else None
+        )
+
         return default_to_dict(
             self,
             aws_access_key_id=self.aws_access_key_id.to_dict() if self.aws_access_key_id else None,
@@ -228,6 +243,7 @@ class S3Downloader:
             max_cache_size=self.max_cache_size,
             file_extensions=self.file_extensions,
             file_name_meta_key=self.file_name_meta_key,
+            s3_key_generation_function=s3_key_generation_function_name,
         )
 
     @classmethod
@@ -239,6 +255,11 @@ class S3Downloader:
         :returns:
             Deserialized component.
         """
+        s3_key_generation_function_name = data["init_parameters"].get("s3_key_generation_function")
+        if s3_key_generation_function_name:
+            data["init_parameters"]["s3_key_generation_function"] = deserialize_callable(
+                s3_key_generation_function_name
+            )
         deserialize_secrets_inplace(
             data["init_parameters"],
             ["aws_access_key_id", "aws_secret_access_key", "aws_session_token", "aws_region_name", "aws_profile_name"],

{amazon_bedrock_haystack-5.0.0 → amazon_bedrock_haystack-5.2.0}/src/haystack_integrations/components/generators/amazon_bedrock/chat/chat_generator.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Tuple
 
 import aioboto3
 from botocore.config import Config
@@ -7,10 +7,10 @@ from botocore.exceptions import ClientError
 from haystack import component, default_from_dict, default_to_dict, logging
 from haystack.dataclasses import ChatMessage, ComponentInfo, StreamingCallbackT, select_streaming_callback
 from haystack.tools import (
-    Tool,
-    Toolset,
+    ToolsType,
     _check_duplicate_tool_names,
     deserialize_tools_or_toolset_inplace,
+    flatten_tools_or_toolsets,
     serialize_tools_or_toolset,
 )
 from haystack.utils.auth import Secret, deserialize_secrets_inplace
@@ -157,7 +157,7 @@ class AmazonBedrockChatGenerator:
         generation_kwargs: Optional[Dict[str, Any]] = None,
         streaming_callback: Optional[StreamingCallbackT] = None,
         boto3_config: Optional[Dict[str, Any]] = None,
-        tools: Optional[Union[List[Tool], Toolset]] = None,
+        tools: Optional[ToolsType] = None,
         *,
         guardrail_config: Optional[Dict[str, str]] = None,
     ) -> None:
@@ -187,7 +187,8 @@ class AmazonBedrockChatGenerator:
             [StreamingChunk](https://docs.haystack.deepset.ai/docs/data-classes#streamingchunk) object and switches
             the streaming mode on.
         :param boto3_config: The configuration for the boto3 client.
-        :param tools: A list of Tool objects or a Toolset that the model can use. Each tool should have a unique name.
+        :param tools: A list of Tool and/or Toolset objects, or a single Toolset for which the model can prepare calls.
+            Each tool should have a unique name.
         :param guardrail_config: Optional configuration for a guardrail that has been created in Amazon Bedrock.
             This must be provided as a dictionary matching either
             [GuardrailConfiguration](https://docs.aws.amazon.com/bedrock/latest/APIReference/API_GuardrailConfiguration.html).
@@ -218,7 +219,7 @@ class AmazonBedrockChatGenerator:
         self.streaming_callback = streaming_callback
         self.boto3_config = boto3_config
 
-        _check_duplicate_tool_names(list(tools or []))  # handles Toolset as well
+        _check_duplicate_tool_names(flatten_tools_or_toolsets(tools))
         self.tools = tools
 
         _validate_guardrail_config(guardrail_config=guardrail_config, streaming=streaming_callback is not None)
@@ -342,7 +343,7 @@ class AmazonBedrockChatGenerator:
         messages: List[ChatMessage],
         streaming_callback: Optional[StreamingCallbackT] = None,
         generation_kwargs: Optional[Dict[str, Any]] = None,
-        tools: Optional[Union[List[Tool], Toolset]] = None,
+        tools: Optional[ToolsType] = None,
         requires_async: bool = False,
     ) -> Tuple[Dict[str, Any], Optional[StreamingCallbackT]]:
         """
@@ -358,7 +359,8 @@ class AmazonBedrockChatGenerator:
             - `stopSequences`: List of stop sequences to stop generation.
             - `temperature`: Sampling temperature.
             - `topP`: Nucleus sampling parameter.
-        :param tools: Optional list of Tool objects or a Toolset that the model can use.
+        :param tools: A list of Tool and/or Toolset objects, or a single Toolset for which the model can prepare calls.
+            Each tool should have a unique name.
         :param requires_async: Boolean flag to indicate if an async-compatible streaming callback function is needed.
 
         :returns:
@@ -380,14 +382,12 @@ class AmazonBedrockChatGenerator:
 
         # Handle tools - either toolConfig or Haystack Tool objects but not both
         tools = tools or self.tools
-        _check_duplicate_tool_names(list(tools or []))
+        flattened_tools = flatten_tools_or_toolsets(tools)
+        _check_duplicate_tool_names(flattened_tools)
         tool_config = merged_kwargs.pop("toolConfig", None)
-        if tools:
-            # Convert Toolset to list if needed
-            if isinstance(tools, Toolset):
-                tools = list(tools)
+        if flattened_tools:
             # Format Haystack tools to Bedrock format
-            tool_config = _format_tools(tools)
+            tool_config = _format_tools(flattened_tools)
 
         # Any remaining kwargs go to additionalModelRequestFields
         additional_fields = merged_kwargs if merged_kwargs else None
@@ -425,7 +425,7 @@ class AmazonBedrockChatGenerator:
         messages: List[ChatMessage],
         streaming_callback: Optional[StreamingCallbackT] = None,
         generation_kwargs: Optional[Dict[str, Any]] = None,
-        tools: Optional[Union[List[Tool], Toolset]] = None,
+        tools: Optional[ToolsType] = None,
     ) -> Dict[str, List[ChatMessage]]:
         """
         Executes a synchronous inference call to the Amazon Bedrock model using the Converse API.
@@ -439,7 +439,8 @@ class AmazonBedrockChatGenerator:
             - `stopSequences`: List of stop sequences to stop generation.
             - `temperature`: Sampling temperature.
             - `topP`: Nucleus sampling parameter.
-        :param tools: Optional list of Tools that the model may call during execution.
+        :param tools: A list of Tool and/or Toolset objects, or a single Toolset for which the model can prepare calls.
+            Each tool should have a unique name.
 
         :returns:
             A dictionary containing the model-generated replies under the `"replies"` key.
@@ -485,7 +486,7 @@ class AmazonBedrockChatGenerator:
         messages: List[ChatMessage],
         streaming_callback: Optional[StreamingCallbackT] = None,
         generation_kwargs: Optional[Dict[str, Any]] = None,
-        tools: Optional[Union[List[Tool], Toolset]] = None,
+        tools: Optional[ToolsType] = None,
     ) -> Dict[str, List[ChatMessage]]:
         """
         Executes an asynchronous inference call to the Amazon Bedrock model using the Converse API.
@@ -499,7 +500,8 @@ class AmazonBedrockChatGenerator:
             - `stopSequences`: List of stop sequences to stop generation.
             - `temperature`: Sampling temperature.
             - `topP`: Nucleus sampling parameter.
-        :param tools: Optional list of Tool objects or a Toolset that the model can use.
+        :param tools: A list of Tool and/or Toolset objects, or a single Toolset for which the model can prepare calls.
+            Each tool should have a unique name.
 
         :returns:
             A dictionary containing the model-generated replies under the `"replies"` key.

{amazon_bedrock_haystack-5.0.0 → amazon_bedrock_haystack-5.2.0}/tests/test_chat_generator.py

@@ -6,7 +6,7 @@ from haystack import Pipeline
 from haystack.components.generators.utils import print_streaming_chunk
 from haystack.components.tools import ToolInvoker
 from haystack.dataclasses import ChatMessage, ChatRole, ImageContent, StreamingChunk, ToolCall
-from haystack.tools import Tool
+from haystack.tools import Tool, Toolset
 
 from haystack_integrations.components.generators.amazon_bedrock import AmazonBedrockChatGenerator
 
@@ -59,6 +59,11 @@ def weather(city: str):
     return f"The weather in {city} is sunny and 32°C"
 
 
+def population(city: str):
+    """Get population for a given city."""
+    return f"The population of {city} is 2.2 million"
+
+
 @pytest.fixture
 def tools():
     tool_parameters = {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}
@@ -71,6 +76,25 @@ def tools():
     return [tool]
 
 
+@pytest.fixture
+def mixed_tools():
+    """Fixture that returns a mixed list of Tool and Toolset."""
+    weather_tool = Tool(
+        name="weather",
+        description="useful to determine the weather in a given location",
+        parameters={"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]},
+        function=weather,
+    )
+    population_tool = Tool(
+        name="population",
+        description="useful to determine the population of a given location",
+        parameters={"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]},
+        function=population,
+    )
+    toolset = Toolset([population_tool])
+    return [weather_tool, toolset]
+
+
 @pytest.fixture
 def chat_messages():
     messages = [
@@ -305,6 +329,86 @@ class TestAmazonBedrockChatGenerator:
         assert request_params["messages"] == [{"content": [{"text": "What's the capital of France?"}], "role": "user"}]
         assert request_params["guardrailConfig"] == {"guardrailIdentifier": "test", "guardrailVersion": "test"}
 
+    def test_init_with_mixed_tools(self, mock_boto3_session, set_env_variables):
+        def tool_fn(city: str) -> str:
+            return city
+
+        weather_tool = Tool(
+            name="weather",
+            description="Weather lookup",
+            parameters={"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]},
+            function=tool_fn,
+        )
+        population_tool = Tool(
+            name="population",
+            description="Population lookup",
+            parameters={"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]},
+            function=tool_fn,
+        )
+        toolset = Toolset([population_tool])
+
+        generator = AmazonBedrockChatGenerator(
+            model="anthropic.claude-3-5-sonnet-20240620-v1:0",
+            tools=[weather_tool, toolset],
+        )
+
+        assert generator.tools == [weather_tool, toolset]
+
+    def test_prepare_request_params_with_mixed_tools(self, mock_boto3_session, set_env_variables):
+        def tool_fn(city: str) -> str:
+            return city
+
+        weather_tool = Tool(
+            name="weather",
+            description="Weather lookup",
+            parameters={"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]},
+            function=tool_fn,
+        )
+        population_tool = Tool(
+            name="population",
+            description="Population lookup",
+            parameters={"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]},
+            function=tool_fn,
+        )
+        toolset = Toolset([population_tool])
+
+        generator = AmazonBedrockChatGenerator(model="anthropic.claude-3-5-sonnet-20240620-v1:0")
+        request_params, _ = generator._prepare_request_params(
+            messages=[ChatMessage.from_user("What's the capital of France?")],
+            tools=[weather_tool, toolset],
+        )
+
+        assert request_params["toolConfig"] == {
+            "tools": [
+                {
+                    "toolSpec": {
+                        "name": "weather",
+                        "description": "Weather lookup",
+                        "inputSchema": {
+                            "json": {
+                                "type": "object",
+                                "properties": {"city": {"type": "string"}},
+                                "required": ["city"],
+                            }
+                        },
+                    }
+                },
+                {
+                    "toolSpec": {
+                        "name": "population",
+                        "description": "Population lookup",
+                        "inputSchema": {
+                            "json": {
+                                "type": "object",
+                                "properties": {"city": {"type": "string"}},
+                                "required": ["city"],
+                            }
+                        },
+                    }
+                },
+            ]
+        }
+
 
 # In the CI, those tests are skipped if AWS Authentication fails
 @pytest.mark.integration
@@ -423,6 +527,66 @@ class TestAmazonBedrockChatGeneratorInference:
         assert "paris" in final_message.text.lower()
         assert "berlin" in final_message.text.lower()
 
+    @pytest.mark.parametrize("model_name", MODELS_TO_TEST_WITH_TOOLS)
+    def test_live_run_with_mixed_tools(self, model_name, mixed_tools):
+        """
+        Integration test that verifies AmazonBedrockChatGenerator works with mixed Tool and Toolset.
+        This tests that the LLM can correctly invoke tools from both a standalone Tool and a Toolset.
+        """
+        initial_messages = [
+            ChatMessage.from_user("What's the weather like in Paris and what is the population of Berlin?")
+        ]
+        component = AmazonBedrockChatGenerator(model=model_name, tools=mixed_tools)
+        results = component.run(messages=initial_messages)
+
+        assert len(results["replies"]) > 0, "No replies received"
+
+        # Find the message with tool calls
+        tool_call_message = None
+        for message in results["replies"]:
+            if message.tool_calls:
+                tool_call_message = message
+                break
+
+        assert tool_call_message is not None, "No message with tool call found"
+        assert isinstance(tool_call_message, ChatMessage), "Tool message is not a ChatMessage instance"
+        assert ChatMessage.is_from(tool_call_message, ChatRole.ASSISTANT), "Tool message is not from the assistant"
+
+        tool_calls = tool_call_message.tool_calls
+        assert len(tool_calls) == 2, f"Expected 2 tool calls, got {len(tool_calls)}"
+
+        # Verify we got calls to both weather and population tools
+        tool_names = {tc.tool_name for tc in tool_calls}
+        assert "weather" in tool_names, "Expected 'weather' tool call"
+        assert "population" in tool_names, "Expected 'population' tool call"
+
+        # Verify tool call details
+        for tool_call in tool_calls:
+            assert tool_call.id, "Tool call does not contain value for 'id' key"
+            assert tool_call.tool_name in ["weather", "population"]
+            assert "city" in tool_call.arguments
+            assert tool_call.arguments["city"] in ["Paris", "Berlin"]
+            assert tool_call_message.meta["finish_reason"] == "tool_calls"
+
+        # Mock the response we'd get from ToolInvoker
+        tool_result_messages = []
+        for tool_call in tool_calls:
+            if tool_call.tool_name == "weather":
+                result = "The weather in Paris is sunny and 32°C"
+            else:  # population
+                result = "The population of Berlin is 2.2 million"
+            tool_result_messages.append(ChatMessage.from_tool(tool_result=result, origin=tool_call))
+
+        new_messages = [*initial_messages, tool_call_message, *tool_result_messages]
+        results = component.run(new_messages)
+
+        assert len(results["replies"]) == 1
+        final_message = results["replies"][0]
+        assert not final_message.tool_call
+        assert len(final_message.text) > 0
+        assert "paris" in final_message.text.lower()
+        assert "berlin" in final_message.text.lower()
+
     @pytest.mark.parametrize("model_name", MODELS_TO_TEST_WITH_THINKING)
     def test_live_run_with_tool_call_and_thinking(self, model_name, tools):
         initial_messages = [ChatMessage.from_user("What's the weather like in Paris?")]

{amazon_bedrock_haystack-5.0.0 → amazon_bedrock_haystack-5.2.0}/tests/test_s3_downloader.py

@@ -30,6 +30,10 @@ def mock_s3_storage():
         yield mock
 
 
+def s3_key_generation_function(document: Document) -> str:
+    return document.meta["file_name"] + "_suffix"
+
+
 class TestS3Downloader:
     def test_init(self, mock_boto3_session, set_env_variables, tmp_path):
         S3Downloader(file_root_path=str(tmp_path))
@@ -74,6 +78,7 @@ class TestS3Downloader:
                 "max_cache_size": 100,
                 "max_workers": 32,
                 "file_name_meta_key": "file_name",
+                "s3_key_generation_function": None,
             },
         }
         assert d.to_dict() == expected
@@ -89,6 +94,7 @@ class TestS3Downloader:
                 "aws_session_token": {"type": "env_var", "env_vars": ["AWS_SESSION_TOKEN"], "strict": False},
                 "aws_profile_name": {"type": "env_var", "env_vars": ["AWS_PROFILE"], "strict": False},
                 "file_root_path": str(tmp_path),
+                "s3_key_generation_function": None,
             },
         }
         d = S3Downloader.from_dict(data)
@@ -101,6 +107,7 @@ class TestS3Downloader:
             max_cache_size=400,
             max_workers=40,
             file_name_meta_key="new_file_key",
+            s3_key_generation_function=s3_key_generation_function,
         )
         expected = {
             "type": TYPE,
@@ -115,6 +122,7 @@ class TestS3Downloader:
                 "max_cache_size": 400,
                 "max_workers": 40,
                 "file_name_meta_key": "new_file_key",
+                "s3_key_generation_function": "tests.test_s3_downloader.s3_key_generation_function",
             },
         }
         assert d.to_dict() == expected
@@ -158,14 +166,50 @@ class TestS3Downloader:
         assert len(out["documents"]) == 1
         assert out["documents"][0].meta["custom_file_key"] == "a.txt"
 
+    def test_run_with_s3_key_generation_function(self, tmp_path, mock_s3_storage, mock_boto3_session):
+        d = S3Downloader(file_root_path=str(tmp_path), s3_key_generation_function=s3_key_generation_function)
+        S3Downloader.warm_up(d)
+        d._storage = mock_s3_storage
+
+        docs = [
+            Document(meta={"file_id": str(uuid4()), "file_name": "a.txt"}),
+        ]
+        out = d.run(documents=docs)
+        assert len(out["documents"]) == 1
+        assert out["documents"][0].meta["file_name"] == "a.txt"
+
+        mock_s3_storage.download.assert_called_once()
+        assert mock_s3_storage.download.call_args.kwargs["key"] == "a.txt_suffix"
+
+    def test_run_with_s3_key_generation_function_and_file_extensions(
+        self, tmp_path, mock_s3_storage, mock_boto3_session
+    ):
+        d = S3Downloader(
+            file_root_path=str(tmp_path),
+            s3_key_generation_function=s3_key_generation_function,
+            file_extensions=[".txt"],
+        )
+        S3Downloader.warm_up(d)
+        d._storage = mock_s3_storage
+
+        docs = [
+            Document(meta={"file_id": str(uuid4()), "file_name": "a.txt"}),
+            Document(meta={"file_id": str(uuid4()), "file_name": "b.pdf"}),
+        ]
+        out = d.run(documents=docs)
+        assert len(out["documents"]) == 1
+        assert out["documents"][0].meta["file_name"] == "a.txt"
+        mock_s3_storage.download.assert_called_once()
+        assert mock_s3_storage.download.call_args.kwargs["key"] == "a.txt_suffix"
+
     @pytest.mark.integration
     @pytest.mark.skipif(
         not os.environ.get("S3_DOWNLOADER_BUCKET", None),
         reason="Export an env var called `S3_DOWNLOADER_BUCKET` containing the S3 bucket to run this test.",
     )
-    def test_live_run(self, tmp_path):
+    def test_live_run(self, tmp_path, monkeypatch):
         d = S3Downloader(file_root_path=str(tmp_path))
-        os.environ["S3_DOWNLOADER_PREFIX"] = ""
+        monkeypatch.setenv("S3_DOWNLOADER_PREFIX", "")
         S3Downloader.warm_up(d)
 
         docs = [
@@ -200,9 +244,9 @@ class TestS3Downloader:
         not os.environ.get("S3_DOWNLOADER_BUCKET", None),
         reason="Export an env var called `S3_DOWNLOADER_BUCKET` containing the S3 bucket to run this test.",
     )
-    def test_live_run_with_custom_meta_key(self, tmp_path):
+    def test_live_run_with_custom_meta_key(self, tmp_path, monkeypatch):
         d = S3Downloader(file_root_path=str(tmp_path), file_name_meta_key="custom_name")
-        os.environ["S3_DOWNLOADER_PREFIX"] = ""
+        monkeypatch.setenv("S3_DOWNLOADER_PREFIX", "")
         S3Downloader.warm_up(d)
         docs = [
             Document(meta={"custom_name": "text-sample.txt"}),
@@ -216,9 +260,9 @@ class TestS3Downloader:
         not os.environ.get("S3_DOWNLOADER_BUCKET", None),
         reason="Export an env var called `S3_DOWNLOADER_BUCKET` containing the S3 bucket to run this test.",
     )
-    def test_live_run_with_prefix(self, tmp_path):
+    def test_live_run_with_prefix(self, tmp_path, monkeypatch):
         d = S3Downloader(file_root_path=str(tmp_path))
-        os.environ["S3_DOWNLOADER_PREFIX"] = "subfolder/"
+        monkeypatch.setenv("S3_DOWNLOADER_PREFIX", "subfolder/")
 
         S3Downloader.warm_up(d)
         docs = [
@@ -227,3 +271,25 @@ class TestS3Downloader:
         out = d.run(documents=docs)
         assert len(out["documents"]) == 1
         assert out["documents"][0].meta["file_name"] == "employees.json"
+
+    @pytest.mark.integration
+    @pytest.mark.skipif(
+        not os.environ.get("S3_DOWNLOADER_BUCKET", None),
+        reason="Export an env var called `S3_DOWNLOADER_BUCKET` containing the S3 bucket to run this test.",
+    )
+    def test_live_run_with_s3_key_generation_function_and_file_extensions(self, tmp_path):
+        # the file in the s3 bucket has this key: "dog.jpg_suffix"
+
+        d = S3Downloader(
+            file_root_path=str(tmp_path),
+            file_extensions=[".jpg"],
+            file_name_meta_key="file_name",
+            s3_key_generation_function=s3_key_generation_function,
+        )
+        S3Downloader.warm_up(d)
+        docs = [
+            Document(meta={"file_name": "dog.jpg"}),
+        ]
+        out = d.run(documents=docs)
+        assert len(out["documents"]) == 1
+        assert out["documents"][0].meta["file_name"] == "dog.jpg"