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

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

@@ -1,5 +1,11 @@
 # Changelog
 
+## [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.1.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: amazon-bedrock-haystack
-Version: 5.0.0
+Version: 5.1.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

{amazon_bedrock_haystack-5.0.0 → amazon_bedrock_haystack-5.1.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.1.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"