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

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

@@ -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.0.0.dist-info → amazon_bedrock_haystack-5.2.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
@@ -18,12 +18,12 @@ haystack_integrations/components/generators/amazon_bedrock/__init__.py,sha256=lv
 haystack_integrations/components/generators/amazon_bedrock/adapters.py,sha256=yBC-3YwV6qAwSXMtdZiLSYh2lUpPQIDy7Efl7w-Cu-k,19640
 haystack_integrations/components/generators/amazon_bedrock/generator.py,sha256=Brzw0XvtPJhz2kR2I3liAqWHRmDR6p5HzJerEAPhoJU,14743
 haystack_integrations/components/generators/amazon_bedrock/chat/__init__.py,sha256=6GZ8Y3Lw0rLOsOAqi6Tu5mZC977UzQvgDxKpOWr8IQw,110
-haystack_integrations/components/generators/amazon_bedrock/chat/chat_generator.py,sha256=qArwfXcforWnPzLXrAW-1hkPFpMy3NSdDyJ5GOta25w,26068
+haystack_integrations/components/generators/amazon_bedrock/chat/chat_generator.py,sha256=fwjNUrQM4UziGBVZKvsVEZVcKPF00ZY7CY0g2Advwxs,26246
 haystack_integrations/components/generators/amazon_bedrock/chat/utils.py,sha256=1M_k8CG2WH23Yz-sB7a1kiIqVh2QB8Pqi0zbWXyMUL8,27255
 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.2.0.dist-info/METADATA,sha256=DO5DjQps1pXcLOiyLmzvoRnf4xgMKE0kYSSFPiFfQeo,2228
+amazon_bedrock_haystack-5.2.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+amazon_bedrock_haystack-5.2.0.dist-info/licenses/LICENSE.txt,sha256=B05uMshqTA74s-0ltyHKI6yoPfJ3zYgQbvcXfDVGFf8,10280
+amazon_bedrock_haystack-5.2.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"],

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.