amazon-bedrock-haystack 5.0.0__py3-none-any.whl → 5.1.0__py3-none-any.whl

{amazon_bedrock_haystack-5.0.0.dist-info → amazon_bedrock_haystack-5.1.0.dist-info}/METADATA

@@ -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.dist-info → amazon_bedrock_haystack-5.1.0.dist-info}/RECORD

@@ -7,7 +7,7 @@ haystack_integrations/common/s3/errors.py,sha256=BrTDLdhQvAuQutyg35cFyP5h8PNkDEi
 haystack_integrations/common/s3/utils.py,sha256=OJupFj54aQmg6S8VuVq6Lc2qpFZyyJajRVIpwe3_6iA,4744
 haystack_integrations/components/downloaders/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 haystack_integrations/components/downloaders/s3/__init__.py,sha256=2BOd3_N0kGqRJGH-ENrTJqOqzqHryRYaSuNqpLYKMFo,179
-haystack_integrations/components/downloaders/s3/s3_downloader.py,sha256=kptTCSry_uEYtNAca_pU7zQJs_LJwwJKYjuYDVJrZRE,11220
+haystack_integrations/components/downloaders/s3/s3_downloader.py,sha256=qarIeGxwDCA1BOZ1qdLfE8NcQtMS9bW54a8voEBHTbE,12637
 haystack_integrations/components/embedders/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 haystack_integrations/components/embedders/amazon_bedrock/__init__.py,sha256=7GlhHJ4jFHCxq5QN5losGuGtrGNjvEx2dSQvEYD2yG0,408
 haystack_integrations/components/embedders/amazon_bedrock/document_embedder.py,sha256=DD34-HAGwGwTU7KWGqKXXlFdwIs21JavBRDHrBqC-m4,13060
@@ -23,7 +23,7 @@ haystack_integrations/components/generators/amazon_bedrock/chat/utils.py,sha256=
 haystack_integrations/components/rankers/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 haystack_integrations/components/rankers/amazon_bedrock/__init__.py,sha256=mJQKShAP5AfZvfKQisSh7kfKu6RIXzsYdk4eqMtcaEk,75
 haystack_integrations/components/rankers/amazon_bedrock/ranker.py,sha256=QWtUKfJxMrlfLCWTb8cCP-lKEthnEBwnTd1NSbiFMkg,11812
-amazon_bedrock_haystack-5.0.0.dist-info/METADATA,sha256=9W2R2hivKGIq1Lj6jdls8wmIncjF-k2ekFy-6ePCy0w,2228
-amazon_bedrock_haystack-5.0.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-amazon_bedrock_haystack-5.0.0.dist-info/licenses/LICENSE.txt,sha256=B05uMshqTA74s-0ltyHKI6yoPfJ3zYgQbvcXfDVGFf8,10280
-amazon_bedrock_haystack-5.0.0.dist-info/RECORD,,
+amazon_bedrock_haystack-5.1.0.dist-info/METADATA,sha256=W5IZ1NhzQ11GjMXl5WrDjRJ0MHnV8jpAHAMyZwNe4Yo,2228
+amazon_bedrock_haystack-5.1.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+amazon_bedrock_haystack-5.1.0.dist-info/licenses/LICENSE.txt,sha256=B05uMshqTA74s-0ltyHKI6yoPfJ3zYgQbvcXfDVGFf8,10280
+amazon_bedrock_haystack-5.1.0.dist-info/RECORD,,

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"],