amazon-bedrock-haystack 3.5.2__py3-none-any.whl → 3.6.1__py3-none-any.whl

{amazon_bedrock_haystack-3.5.2.dist-info → amazon_bedrock_haystack-3.6.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: amazon-bedrock-haystack
-Version: 3.5.2
+Version: 3.6.1
 Summary: An integration of Amazon Bedrock as an AmazonBedrockGenerator component.
 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
+Requires-Dist: haystack-ai>=2.13.1
 Description-Content-Type: text/markdown
 
 # amazon-bedrock-haystack

{amazon_bedrock_haystack-3.5.2.dist-info → amazon_bedrock_haystack-3.6.1.dist-info}/RECORD

@@ -1,18 +1,18 @@
 haystack_integrations/common/amazon_bedrock/__init__.py,sha256=6GZ8Y3Lw0rLOsOAqi6Tu5mZC977UzQvgDxKpOWr8IQw,110
 haystack_integrations/common/amazon_bedrock/errors.py,sha256=ReheDbY7L3EJkWcUoih6lWHjbPHg2TlUs9SnXIKK7Gg,744
-haystack_integrations/common/amazon_bedrock/utils.py,sha256=jVXNjzPvYrvWtEJi8FjaW0aBqoTyFp6btp5mXo7F5Go,2687
+haystack_integrations/common/amazon_bedrock/utils.py,sha256=dHUWzHYT0A8_eLDpVkwDhmDpprYbFlWsGg0FOS0uF0I,2720
 haystack_integrations/components/embedders/amazon_bedrock/__init__.py,sha256=CFqYmAVq2aavlMkZHYScKHOTwwETdRzRZITMqGhJ9Kw,298
-haystack_integrations/components/embedders/amazon_bedrock/document_embedder.py,sha256=8Kj3_JiND2IQGdAc8n2GnI244QzbOjNmEHuH_W2T610,13115
-haystack_integrations/components/embedders/amazon_bedrock/text_embedder.py,sha256=KRreCYplvxT6Yvg5TMSb4aDz5pqW1zMeV216cR8C6FY,9005
+haystack_integrations/components/embedders/amazon_bedrock/document_embedder.py,sha256=ZdtM6HpHQwbKfjmfOK6gkIQPPbI0n8_pWRrR6lyXmr8,13321
+haystack_integrations/components/embedders/amazon_bedrock/text_embedder.py,sha256=gpvu6IMoycUXrn4r1OH5yEIheiDxHf2T5fdJJO4DfW0,9202
 haystack_integrations/components/generators/amazon_bedrock/__init__.py,sha256=lv4NouIVm78YavUssWQrHHP_81u-7j21qW8v1kZMJPQ,284
 haystack_integrations/components/generators/amazon_bedrock/adapters.py,sha256=cnlfmie4HfEX4nipSXSDk_3koy7HYZ-ezimGN6BozQ0,19543
-haystack_integrations/components/generators/amazon_bedrock/generator.py,sha256=2L9-ZKJkMI5JxN0xZ8Gn_5MfdjhT4UlFjlrTd9Rdhsg,14514
+haystack_integrations/components/generators/amazon_bedrock/generator.py,sha256=NgywyiKYazEbsLAcGcOPUT4blWhYYOJ9WjO-HWDvu7I,14576
 haystack_integrations/components/generators/amazon_bedrock/chat/__init__.py,sha256=6GZ8Y3Lw0rLOsOAqi6Tu5mZC977UzQvgDxKpOWr8IQw,110
-haystack_integrations/components/generators/amazon_bedrock/chat/chat_generator.py,sha256=f60FRraGvPeMBiosclNN3hLv9uQihQcQFETC-gwGieM,20784
-haystack_integrations/components/generators/amazon_bedrock/chat/utils.py,sha256=Dch2YkrDFu-Dvi9i1bE97taAA1Qf3wMQb4wh2T3Z-lw,11789
-haystack_integrations/components/rankers/amazon_bedrock/__init__.py,sha256=DWsCu-dav2wzr13U2H1jlFjdzdjkV0fnJw7DVDRY8RQ,63
-haystack_integrations/components/rankers/amazon_bedrock/ranker.py,sha256=kq9h_ApGXBDs_GW66QQpo8JwW9grLpqpFT9uXiwT8_M,11643
-amazon_bedrock_haystack-3.5.2.dist-info/METADATA,sha256=gumDQJXyNSXhdWQFhVuVY1OSmhOGJiXP9re3y2uTVE0,2217
-amazon_bedrock_haystack-3.5.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-amazon_bedrock_haystack-3.5.2.dist-info/licenses/LICENSE.txt,sha256=B05uMshqTA74s-0ltyHKI6yoPfJ3zYgQbvcXfDVGFf8,10280
-amazon_bedrock_haystack-3.5.2.dist-info/RECORD,,
+haystack_integrations/components/generators/amazon_bedrock/chat/chat_generator.py,sha256=M9I0sB8LFrXCgoyr5ik2ZPPHyB0b4nJFKX7GARKsk8Y,23384
+haystack_integrations/components/generators/amazon_bedrock/chat/utils.py,sha256=kBSaU_ZqzL-7a7nplezjb4XRBy51pt-4VULoX5lq21A,21148
+haystack_integrations/components/rankers/amazon_bedrock/__init__.py,sha256=Zrc3BSVkEaXYpliEi6hKG9bqW4J7DNk93p50SuoyT1Q,107
+haystack_integrations/components/rankers/amazon_bedrock/ranker.py,sha256=x4QEVkbFM-jMFHx-xmk571wtrohnPLtkIWMhCyg4_II,12278
+amazon_bedrock_haystack-3.6.1.dist-info/METADATA,sha256=oRIb-2Nv642N0wmapQgykocVVjxzwM7rG1_TNOoF6Vs,2225
+amazon_bedrock_haystack-3.6.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+amazon_bedrock_haystack-3.6.1.dist-info/licenses/LICENSE.txt,sha256=B05uMshqTA74s-0ltyHKI6yoPfJ3zYgQbvcXfDVGFf8,10280
+amazon_bedrock_haystack-3.6.1.dist-info/RECORD,,

haystack_integrations/common/amazon_bedrock/utils.py

@@ -1,6 +1,6 @@
-from typing import Optional
+from typing import Optional, Union
 
-import aioboto3  # type: ignore
+import aioboto3
 import boto3
 from botocore.exceptions import BotoCoreError
 
@@ -23,7 +23,7 @@ def get_aws_session(
     aws_profile_name: Optional[str] = None,
     async_mode: bool = False,
     **kwargs,
-):
+) -> Union[boto3.Session, aioboto3.Session]:
     """
     Creates an AWS Session with the given parameters.
     Checks if the provided AWS credentials are valid and can be used to connect to AWS.

haystack_integrations/components/embedders/amazon_bedrock/document_embedder.py

@@ -75,7 +75,7 @@ class AmazonBedrockDocumentEmbedder:
         embedding_separator: str = "\n",
         boto3_config: Optional[Dict[str, Any]] = None,
         **kwargs,
-    ):
+    ) -> None:
         """
         Initializes the AmazonBedrockDocumentEmbedder with the provided parameters. The parameters are passed to the
         Amazon Bedrock client.
@@ -234,7 +234,7 @@ class AmazonBedrockDocumentEmbedder:
         return documents
 
     @component.output_types(documents=List[Document])
-    def run(self, documents: List[Document]):
+    def run(self, documents: List[Document]) -> Dict[str, List[Document]]:
         """Embed the provided `Document`s using the specified model.
 
         :param documents: The `Document`s to embed.
@@ -253,6 +253,9 @@ class AmazonBedrockDocumentEmbedder:
             documents_with_embeddings = self._embed_cohere(documents=documents)
         elif "titan" in self.model:
             documents_with_embeddings = self._embed_titan(documents=documents)
+        else:
+            msg = f"Model {self.model} is not supported. Supported models are: {', '.join(SUPPORTED_EMBEDDING_MODELS)}."
+            raise ValueError(msg)
 
         return {"documents": documents_with_embeddings}
 

haystack_integrations/components/embedders/amazon_bedrock/text_embedder.py

@@ -64,7 +64,7 @@ class AmazonBedrockTextEmbedder:
         aws_profile_name: Optional[Secret] = Secret.from_env_var("AWS_PROFILE", strict=False),  # noqa: B008
         boto3_config: Optional[Dict[str, Any]] = None,
         **kwargs,
-    ):
+    ) -> None:
         """
         Initializes the AmazonBedrockTextEmbedder with the provided parameters. The parameters are passed to the
         Amazon Bedrock client.
@@ -126,7 +126,7 @@ class AmazonBedrockTextEmbedder:
             raise AmazonBedrockConfigurationError(msg) from exception
 
     @component.output_types(embedding=List[float])
-    def run(self, text: str):
+    def run(self, text: str) -> Dict[str, List[float]]:
         """Embeds the input text using the Amazon Bedrock model.
 
         :param text: The input text to embed.
@@ -173,6 +173,9 @@ class AmazonBedrockTextEmbedder:
             embedding = response_body["embeddings"][0]
         elif "titan" in self.model:
             embedding = response_body["embedding"]
+        else:
+            msg = f"Unsupported model {self.model}. Supported models are: {', '.join(SUPPORTED_EMBEDDING_MODELS)}"
+            raise ValueError(msg)
 
         return {"embedding": embedding}
 

haystack_integrations/components/generators/amazon_bedrock/chat/chat_generator.py

@@ -1,20 +1,21 @@
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple, Union
 
+import aioboto3
 from botocore.config import Config
 from botocore.eventstream import EventStream
 from botocore.exceptions import ClientError
 from haystack import component, default_from_dict, default_to_dict, logging
 from haystack.dataclasses import ChatMessage, StreamingCallbackT, select_streaming_callback
-from haystack.tools import Tool, _check_duplicate_tool_names
+from haystack.tools import (
+    Tool,
+    Toolset,
+    _check_duplicate_tool_names,
+    deserialize_tools_or_toolset_inplace,
+    serialize_tools_or_toolset,
+)
 from haystack.utils.auth import Secret, deserialize_secrets_inplace
 from haystack.utils.callable_serialization import deserialize_callable, serialize_callable
 
-# Compatibility with Haystack 2.12.0 and 2.13.0 - remove after 2.13.0 is released
-try:
-    from haystack.tools import deserialize_tools_or_toolset_inplace
-except ImportError:
-    from haystack.tools import deserialize_tools_inplace as deserialize_tools_or_toolset_inplace
-
 from haystack_integrations.common.amazon_bedrock.errors import (
     AmazonBedrockConfigurationError,
     AmazonBedrockInferenceError,
@@ -140,8 +141,8 @@ class AmazonBedrockChatGenerator:
         stop_words: Optional[List[str]] = None,
         streaming_callback: Optional[StreamingCallbackT] = None,
         boto3_config: Optional[Dict[str, Any]] = None,
-        tools: Optional[List[Tool]] = None,
-    ):
+        tools: Optional[Union[List[Tool], Toolset]] = None,
+    ) -> None:
         """
         Initializes the `AmazonBedrockChatGenerator` with the provided parameters. The parameters are passed to the
         Amazon Bedrock client.
@@ -172,7 +173,7 @@ 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 that the model can use. Each tool should have a unique name.
+        :param tools: A list of Tool objects or a Toolset that the model can use. Each tool should have a unique name.
 
         :raises ValueError: If the model name is empty or None.
         :raises AmazonBedrockConfigurationError: If the AWS environment is not configured correctly or the model is
@@ -190,7 +191,7 @@ class AmazonBedrockChatGenerator:
         self.stop_words = stop_words or []
         self.streaming_callback = streaming_callback
         self.boto3_config = boto3_config
-        _check_duplicate_tool_names(tools)
+        _check_duplicate_tool_names(list(tools or []))  # handles Toolset as well
         self.tools = tools
 
         def resolve_secret(secret: Optional[Secret]) -> Optional[str]:
@@ -228,7 +229,18 @@ class AmazonBedrockChatGenerator:
         self.stop_words = stop_words or []
         self.async_session = None
 
-    def _get_async_session(self):
+    def _get_async_session(self) -> aioboto3.Session:
+        """
+        Initializes and returns an asynchronous AWS session for accessing Amazon Bedrock.
+
+        If the session is already created, it is reused. Otherwise, a new session is created using the provided AWS
+        credentials and configuration.
+
+        :returns:
+            An async-compatible boto3 session configured for use with Amazon Bedrock.
+        :raises AmazonBedrockConfigurationError:
+            If unable to establish an async session due to misconfiguration.
+        """
         if self.async_session:
             return self.async_session
 
@@ -260,7 +272,6 @@ class AmazonBedrockChatGenerator:
             Dictionary with serialized data.
         """
         callback_name = serialize_callable(self.streaming_callback) if self.streaming_callback else None
-        serialized_tools = [tool.to_dict() for tool in self.tools] if self.tools 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,
@@ -273,7 +284,7 @@ class AmazonBedrockChatGenerator:
             generation_kwargs=self.generation_kwargs,
             streaming_callback=callback_name,
             boto3_config=self.boto3_config,
-            tools=serialized_tools,
+            tools=serialize_tools_or_toolset(self.tools),
         )
 
     @classmethod
@@ -301,19 +312,27 @@ class AmazonBedrockChatGenerator:
         messages: List[ChatMessage],
         streaming_callback: Optional[StreamingCallbackT] = None,
         generation_kwargs: Optional[Dict[str, Any]] = None,
-        tools: Optional[List[Tool]] = None,
+        tools: Optional[Union[List[Tool], Toolset]] = None,
         requires_async: bool = False,
     ) -> Tuple[Dict[str, Any], Optional[StreamingCallbackT]]:
         """
-        Prepares the request parameters for both sync and async run methods.
-
-        :param messages: List of ChatMessage objects representing the conversation history.
-        :param streaming_callback: Optional callback function for handling streaming responses.
-        :param generation_kwargs: Optional dictionary of generation parameters.
-        :param tools: Optional list of Tool objects that the model can use.
-        :param requires_async: Boolean indicating whether the request is for async execution.
-            This affects how the streaming callback is selected.
-        :return: Tuple of (request parameters dict, callback function)
+        Prepares and formats parameters required to call the Amazon Bedrock Converse API.
+
+        This includes merging default and runtime generation parameters, formatting messages and tools, and
+        selecting the appropriate streaming callback.
+
+        :param messages: List of `ChatMessage` objects representing the conversation history.
+        :param streaming_callback: Optional streaming callback provided at runtime.
+        :param generation_kwargs: Optional dictionary of generation parameters. Some common parameters are:
+            - `maxTokens`: Maximum number of tokens to generate.
+            - `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 requires_async: Boolean flag to indicate if an async-compatible streaming callback function is needed.
+
+        :returns:
+            A tuple of (API-ready parameter dictionary, streaming callback function).
         """
         generation_kwargs = generation_kwargs or {}
 
@@ -331,9 +350,12 @@ class AmazonBedrockChatGenerator:
 
         # Handle tools - either toolConfig or Haystack Tool objects but not both
         tools = tools or self.tools
-        _check_duplicate_tool_names(tools)
+        _check_duplicate_tool_names(list(tools or []))
         tool_config = merged_kwargs.pop("toolConfig", None)
         if tools:
+            # Convert Toolset to list if needed
+            if isinstance(tools, Toolset):
+                tools = list(tools)
             # Format Haystack tools to Bedrock format
             tool_config = _format_tools(tools)
 
@@ -369,8 +391,27 @@ class AmazonBedrockChatGenerator:
         messages: List[ChatMessage],
         streaming_callback: Optional[StreamingCallbackT] = None,
         generation_kwargs: Optional[Dict[str, Any]] = None,
-        tools: Optional[List[Tool]] = None,
-    ):
+        tools: Optional[Union[List[Tool], Toolset]] = None,
+    ) -> Dict[str, List[ChatMessage]]:
+        """
+        Executes a synchronous inference call to the Amazon Bedrock model using the Converse API.
+
+        Supports both standard and streaming responses depending on whether a streaming callback is provided.
+
+        :param messages: A list of `ChatMessage` objects forming the chat history.
+        :param streaming_callback: Optional callback for handling streaming outputs.
+        :param generation_kwargs: Optional dictionary of generation parameters. Some common parameters are:
+            - `maxTokens`: Maximum number of tokens to generate.
+            - `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.
+
+        :returns:
+            A dictionary containing the model-generated replies under the `"replies"` key.
+        :raises AmazonBedrockInferenceError:
+            If the Bedrock inference API call fails.
+        """
         params, callback = self._prepare_request_params(
             messages=messages,
             streaming_callback=streaming_callback,
@@ -402,16 +443,26 @@ class AmazonBedrockChatGenerator:
         messages: List[ChatMessage],
         streaming_callback: Optional[StreamingCallbackT] = None,
         generation_kwargs: Optional[Dict[str, Any]] = None,
-        tools: Optional[List[Tool]] = None,
-    ):
+        tools: Optional[Union[List[Tool], Toolset]] = None,
+    ) -> Dict[str, List[ChatMessage]]:
         """
-        Async version of the run method. Completes chats using LLMs hosted on Amazon Bedrock.
+        Executes an asynchronous inference call to the Amazon Bedrock model using the Converse API.
+
+        Designed for use cases where non-blocking or concurrent execution is desired.
 
-        :param messages: List of ChatMessage objects representing the conversation history.
-        :param streaming_callback: Optional callback function for handling streaming responses.
-        :param generation_kwargs: Optional dictionary of generation parameters.
-        :param tools: Optional list of Tool objects that the model can use.
-        :return: Dictionary containing the model's replies as a list of ChatMessage objects.
+        :param messages: A list of `ChatMessage` objects forming the chat history.
+        :param streaming_callback: Optional async-compatible callback for handling streaming outputs.
+        :param generation_kwargs: Optional dictionary of generation parameters. Some common parameters are:
+            - `maxTokens`: Maximum number of tokens to generate.
+            - `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.
+
+        :returns:
+            A dictionary containing the model-generated replies under the `"replies"` key.
+        :raises AmazonBedrockInferenceError:
+            If the Bedrock inference API call fails.
         """
         params, callback = self._prepare_request_params(
             messages=messages,

haystack_integrations/components/generators/amazon_bedrock/chat/utils.py

@@ -1,17 +1,30 @@
 import json
-from typing import Any, Callable, Dict, List, Optional, Tuple
+from datetime import datetime, timezone
+from typing import Any, Dict, List, Optional, Tuple
 
 from botocore.eventstream import EventStream
-from haystack.dataclasses import ChatMessage, ChatRole, StreamingCallbackT, StreamingChunk, ToolCall
+from haystack import logging
+from haystack.dataclasses import (
+    AsyncStreamingCallbackT,
+    ChatMessage,
+    ChatRole,
+    StreamingChunk,
+    SyncStreamingCallbackT,
+    ToolCall,
+)
 from haystack.tools import Tool
 
+logger = logging.getLogger(__name__)
 
+
+# Haystack to Bedrock util methods
 def _format_tools(tools: Optional[List[Tool]] = None) -> Optional[Dict[str, Any]]:
     """
     Format Haystack Tool(s) to Amazon Bedrock toolConfig format.
 
     :param tools: List of Tool objects to format
-    :return: Dictionary in Bedrock toolConfig format or None if no tools
+    :returns:
+        Dictionary in Bedrock toolConfig format or None if no tools are provided
     """
     if not tools:
         return None
@@ -25,69 +38,160 @@ def _format_tools(tools: Optional[List[Tool]] = None) -> Optional[Dict[str, Any]
     return {"tools": tool_specs} if tool_specs else None
 
 
+def _format_tool_call_message(tool_call_message: ChatMessage) -> Dict[str, Any]:
+    """
+    Format a Haystack ChatMessage containing tool calls into Bedrock format.
+
+    :param tool_call_message: ChatMessage object containing tool calls to be formatted.
+    :returns:
+        Dictionary representing the tool call message in Bedrock's expected format
+    """
+    content = []
+    # Tool call message can contain text
+    if tool_call_message.text:
+        content.append({"text": tool_call_message.text})
+
+    for tool_call in tool_call_message.tool_calls:
+        content.append(
+            {"toolUse": {"toolUseId": tool_call.id, "name": tool_call.tool_name, "input": tool_call.arguments}}
+        )
+    return {"role": tool_call_message.role.value, "content": content}
+
+
+def _format_tool_result_message(tool_call_result_message: ChatMessage) -> Dict[str, Any]:
+    """
+    Format a Haystack ChatMessage containing tool call results into Bedrock format.
+
+    :param tool_call_result_message: ChatMessage object containing tool call results to be formatted.
+    :returns: Dictionary representing the tool result message in Bedrock's expected format
+    """
+    # Assuming tool call result messages will only contain tool results
+    tool_results = []
+    for result in tool_call_result_message.tool_call_results:
+        try:
+            json_result = json.loads(result.result)
+            content = [{"json": json_result}]
+        except json.JSONDecodeError:
+            content = [{"text": result.result}]
+
+        tool_results.append(
+            {
+                "toolResult": {
+                    "toolUseId": result.origin.id,
+                    "content": content,
+                    **({"status": "error"} if result.error else {}),
+                }
+            }
+        )
+    # role must be user
+    return {"role": "user", "content": tool_results}
+
+
+def _repair_tool_result_messages(bedrock_formatted_messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    """
+    Repair and reorganize tool result messages to maintain proper ordering and grouping.
+
+    Ensures tool result messages are properly grouped in the same way as their corresponding tool call messages
+    and maintains the original message ordering.
+
+    :param bedrock_formatted_messages: List of Bedrock-formatted messages that may need repair.
+    :returns: List of properly organized Bedrock-formatted messages with correctly grouped tool results.
+    """
+    tool_call_messages = []
+    tool_result_messages = []
+    for idx, msg in enumerate(bedrock_formatted_messages):
+        content = msg.get("content", [])
+        if content:
+            if any("toolUse" in c for c in content):
+                tool_call_messages.append((idx, msg))
+            elif any("toolResult" in c for c in content):
+                tool_result_messages.append((idx, msg))
+
+    # Determine the tool call IDs for each tool call message
+    group_to_tool_call_ids: Dict[int, Any] = {idx: [] for idx, _ in tool_call_messages}
+    for idx, tool_call in tool_call_messages:
+        tool_use_contents = [c for c in tool_call["content"] if "toolUse" in c]
+        for content in tool_use_contents:
+            group_to_tool_call_ids[idx].append(content["toolUse"]["toolUseId"])
+
+    # Regroups the tool_result_prompts based on the tool_call_prompts
+    # Makes sure to:
+    # - Within the new group the tool call IDs of the tool result messages are in the same order as the tool call
+    #   messages
+    # - The tool result messages are in the same order as the original message list
+    repaired_tool_result_prompts = []
+    for tool_call_ids in group_to_tool_call_ids.values():
+        regrouped_tool_result = []
+        original_idx = None
+        for tool_call_id in tool_call_ids:
+            for idx, tool_result in tool_result_messages:
+                tool_result_contents = [c for c in tool_result["content"] if "toolResult" in c]
+                for content in tool_result_contents:
+                    if content["toolResult"]["toolUseId"] == tool_call_id:
+                        regrouped_tool_result.append(content)
+                        # Keep track of the original index of the last tool result message
+                        original_idx = idx
+        if regrouped_tool_result and original_idx is not None:
+            repaired_tool_result_prompts.append((original_idx, {"role": "user", "content": regrouped_tool_result}))
+
+    # Remove the tool result messages from bedrock_formatted_messages
+    bedrock_formatted_messages_minus_tool_results: List[Tuple[int, Any]] = []
+    for idx, msg in enumerate(bedrock_formatted_messages):
+        # Assumes the content of tool result messages only contains 'toolResult': {...} objects (e.g. no 'text')
+        if msg.get("content") and "toolResult" not in msg["content"][0]:
+            bedrock_formatted_messages_minus_tool_results.append((idx, msg))
+
+    # Add the repaired tool result messages and sort to maintain the correct order
+    repaired_bedrock_formatted_messages = bedrock_formatted_messages_minus_tool_results + repaired_tool_result_prompts
+    repaired_bedrock_formatted_messages.sort(key=lambda x: x[0])
+
+    # Drop the index and return only the messages
+    return [msg for _, msg in repaired_bedrock_formatted_messages]
+
+
 def _format_messages(messages: List[ChatMessage]) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]:
     """
-    Format a list of ChatMessages to the format expected by Bedrock API.
-    Separates system messages and handles tool results and tool calls.
+    Format a list of Haystack ChatMessages to the format expected by Bedrock API.
 
-    :param messages: List of ChatMessages to format
-    :return: Tuple of (system_prompts, non_system_messages) in Bedrock format
+    Processes and separates system messages from other message types and handles special formatting for tool calls
+    and tool results.
+
+    :param messages: List of ChatMessage objects to format for Bedrock API.
+    :returns: Tuple containing (system_prompts, non_system_messages) in Bedrock format,
+              where system_prompts is a list of system message dictionaries and
+              non_system_messages is a list of properly formatted message dictionaries.
     """
+    # Separate system messages, tool calls, and tool results
     system_prompts = []
-    non_system_messages = []
-
+    bedrock_formatted_messages = []
     for msg in messages:
         if msg.is_from(ChatRole.SYSTEM):
+            # Assuming system messages can only contain text
+            # Don't need to track idx since system_messages are handled separately
             system_prompts.append({"text": msg.text})
-            continue
-
-        # Handle tool results - must role these as user messages
-        if msg.tool_call_results:
-            tool_results = []
-            for result in msg.tool_call_results:
-                try:
-                    json_result = json.loads(result.result)
-                    content = [{"json": json_result}]
-                except json.JSONDecodeError:
-                    content = [{"text": result.result}]
-
-                tool_results.append(
-                    {
-                        "toolResult": {
-                            "toolUseId": result.origin.id,
-                            "content": content,
-                            **({"status": "error"} if result.error else {}),
-                        }
-                    }
-                )
-            non_system_messages.append({"role": "user", "content": tool_results})
-            continue
-
-        content = []
-        # Handle text content
-        if msg.text:
-            content.append({"text": msg.text})
-
-        # Handle tool calls
-        if msg.tool_calls:
-            for tool_call in msg.tool_calls:
-                content.append(
-                    {"toolUse": {"toolUseId": tool_call.id, "name": tool_call.tool_name, "input": tool_call.arguments}}
-                )
-
-        if content:  # Only add message if it has content
-            non_system_messages.append({"role": msg.role.value, "content": content})
+        elif msg.tool_calls:
+            bedrock_formatted_messages.append(_format_tool_call_message(msg))
+        elif msg.tool_call_results:
+            bedrock_formatted_messages.append(_format_tool_result_message(msg))
+        else:
+            # regular user or assistant messages with only text content
+            bedrock_formatted_messages.append({"role": msg.role.value, "content": [{"text": msg.text}]})
 
-    return system_prompts, non_system_messages
+    repaired_bedrock_formatted_messages = _repair_tool_result_messages(bedrock_formatted_messages)
+    return system_prompts, repaired_bedrock_formatted_messages
 
 
+# Bedrock to Haystack util method
 def _parse_completion_response(response_body: Dict[str, Any], model: str) -> List[ChatMessage]:
     """
-    Parse a Bedrock response to a list of ChatMessage objects.
+    Parse a Bedrock API response into Haystack ChatMessage objects.
 
-    :param response_body: Raw response from Bedrock API
-    :param model: The model ID used for generation
-    :return: List of ChatMessage objects
+    Extracts text content, tool calls, and metadata from the Bedrock response and converts them into the appropriate
+    Haystack format.
+
+    :param response_body: Raw JSON response from Bedrock API.
+    :param model: The model ID used for generation, included in message metadata.
+    :returns: List of ChatMessage objects containing the assistant's response(s) with appropriate metadata.
     """
     replies = []
     if "output" in response_body and "message" in response_body["output"]:
@@ -130,9 +234,209 @@ def _parse_completion_response(response_body: Dict[str, Any], model: str) -> Lis
     return replies
 
 
+# Bedrock streaming to Haystack util methods
+def _convert_event_to_streaming_chunk(event: Dict[str, Any], model: str) -> StreamingChunk:
+    """
+    Convert a Bedrock streaming event to a Haystack StreamingChunk.
+
+    Handles different event types (contentBlockStart, contentBlockDelta, messageStop, metadata) and extracts relevant
+    information to create StreamingChunk objects in the same format used by Haystack's OpenAIChatGenerator.
+
+    :param event: Dictionary containing a Bedrock streaming event.
+    :param model: The model ID used for generation, included in chunk metadata.
+    :returns: StreamingChunk object containing the content and metadata extracted from the event.
+    """
+    # Initialize an empty StreamingChunk to return if no relevant event is found
+    # (e.g. for messageStart and contentBlockStop)
+    streaming_chunk = StreamingChunk(
+        content="", meta={"model": model, "received_at": datetime.now(timezone.utc).isoformat()}
+    )
+
+    if "contentBlockStart" in event:
+        # contentBlockStart always has the key "contentBlockIndex"
+        block_start = event["contentBlockStart"]
+        block_idx = block_start["contentBlockIndex"]
+        if "start" in block_start and "toolUse" in block_start["start"]:
+            tool_start = block_start["start"]["toolUse"]
+            streaming_chunk = StreamingChunk(
+                content="",
+                meta={
+                    "model": model,
+                    # This is always 0 b/c it represents the choice index
+                    "index": 0,
+                    # We follow the same format used in the OpenAIChatGenerator
+                    "tool_calls": [  # Optional[List[ChoiceDeltaToolCall]]
+                        {
+                            "index": block_idx,  # int
+                            "id": tool_start["toolUseId"],  # Optional[str]
+                            "function": {  # Optional[ChoiceDeltaToolCallFunction]
+                                # Will accumulate deltas as string
+                                "arguments": "",  # Optional[str]
+                                "name": tool_start["name"],  # Optional[str]
+                            },
+                            "type": "function",  # Optional[Literal["function"]]
+                        }
+                    ],
+                    "finish_reason": None,
+                    "received_at": datetime.now(timezone.utc).isoformat(),
+                },
+            )
+
+    elif "contentBlockDelta" in event:
+        # contentBlockDelta always has the key "contentBlockIndex" and "delta"
+        block_idx = event["contentBlockDelta"]["contentBlockIndex"]
+        delta = event["contentBlockDelta"]["delta"]
+        # This is for accumulating text deltas
+        if "text" in delta:
+            streaming_chunk = StreamingChunk(
+                content=delta["text"],
+                meta={
+                    "model": model,
+                    # This is always 0 b/c it represents the choice index
+                    "index": 0,
+                    "tool_calls": None,
+                    "finish_reason": None,
+                    "received_at": datetime.now(timezone.utc).isoformat(),
+                },
+            )
+        # This only occurs when accumulating the arguments for a toolUse
+        # The content_block for this tool should already exist at this point
+        elif "toolUse" in delta:
+            streaming_chunk = StreamingChunk(
+                content="",
+                meta={
+                    "model": model,
+                    # This is always 0 b/c it represents the choice index
+                    "index": 0,
+                    "tool_calls": [  # Optional[List[ChoiceDeltaToolCall]]
+                        {
+                            "index": block_idx,  # int
+                            "id": None,  # Optional[str]
+                            "function": {  # Optional[ChoiceDeltaToolCallFunction]
+                                # Will accumulate deltas as string
+                                "arguments": delta["toolUse"].get("input", ""),  # Optional[str]
+                                "name": None,  # Optional[str]
+                            },
+                            "type": "function",  # Optional[Literal["function"]]
+                        }
+                    ],
+                    "finish_reason": None,
+                    "received_at": datetime.now(timezone.utc).isoformat(),
+                },
+            )
+
+    elif "messageStop" in event:
+        finish_reason = event["messageStop"].get("stopReason")
+        streaming_chunk = StreamingChunk(
+            content="",
+            meta={
+                "model": model,
+                # This is always 0 b/c it represents the choice index
+                "index": 0,
+                "tool_calls": None,
+                "finish_reason": finish_reason,
+                "received_at": datetime.now(timezone.utc).isoformat(),
+            },
+        )
+
+    elif "metadata" in event and "usage" in event["metadata"]:
+        metadata = event["metadata"]
+        streaming_chunk = StreamingChunk(
+            content="",
+            meta={
+                "model": model,
+                # This is always 0 b/c it represents the choice index
+                "index": 0,
+                "tool_calls": None,
+                "finish_reason": None,
+                "received_at": datetime.now(timezone.utc).isoformat(),
+                "usage": {
+                    "prompt_tokens": metadata["usage"].get("inputTokens", 0),
+                    "completion_tokens": metadata["usage"].get("outputTokens", 0),
+                    "total_tokens": metadata["usage"].get("totalTokens", 0),
+                },
+            },
+        )
+
+    return streaming_chunk
+
+
+def _convert_streaming_chunks_to_chat_message(chunks: List[StreamingChunk]) -> ChatMessage:
+    """
+    Converts a list of streaming chunks into a ChatMessage object.
+
+    The function processes streaming chunks to build a ChatMessage object, including extracting and constructing
+    tool calls, managing metadata such as model type, finish reason, and usage information.
+    The tool call processing handles accumulating data across the chunks and attempts to parse JSON-formatted
+    arguments for tool calls.
+
+    :param chunks: A list of StreamingChunk objects representing parts of the assistant's response.
+
+    :returns:
+        A ChatMessage object constructed from the streaming chunks, containing the aggregated text, processed tool
+        calls, and metadata.
+    """
+    text = "".join([chunk.content for chunk in chunks])
+
+    # Process tool calls if present in any chunk
+    tool_calls = []
+    tool_call_data: Dict[int, Dict[str, str]] = {}  # Track tool calls by index
+    for chunk_payload in chunks:
+        tool_calls_meta = chunk_payload.meta.get("tool_calls")
+        if tool_calls_meta is not None:
+            for delta in tool_calls_meta:
+                # We use the index of the tool call to track it across chunks since the ID is not always provided
+                if delta["index"] not in tool_call_data:
+                    tool_call_data[delta["index"]] = {"id": "", "name": "", "arguments": ""}
+
+                # Save the ID if present
+                if delta.get("id"):
+                    tool_call_data[delta["index"]]["id"] = delta["id"]
+
+                if delta.get("function"):
+                    if delta["function"].get("name"):
+                        tool_call_data[delta["index"]]["name"] += delta["function"]["name"]
+                    if delta["function"].get("arguments"):
+                        tool_call_data[delta["index"]]["arguments"] += delta["function"]["arguments"]
+
+    # Convert accumulated tool call data into ToolCall objects
+    for call_data in tool_call_data.values():
+        try:
+            arguments = json.loads(call_data["arguments"])
+            tool_calls.append(ToolCall(id=call_data["id"], tool_name=call_data["name"], arguments=arguments))
+        except json.JSONDecodeError:
+            logger.warning(
+                "Amazon Bedrock returned a malformed JSON string for tool call arguments. This tool call will be "
+                "skipped. Tool call ID: {tool_id}, Tool name: {tool_name}, Arguments: {tool_arguments}",
+                tool_id=call_data["id"],
+                tool_name=call_data["name"],
+                tool_arguments=call_data["arguments"],
+            )
+
+    # finish_reason can appear in different places so we look for the last one
+    finish_reasons = [
+        chunk.meta.get("finish_reason") for chunk in chunks if chunk.meta.get("finish_reason") is not None
+    ]
+    finish_reason = finish_reasons[-1] if finish_reasons else None
+
+    # usage is usually last but we look for it as well
+    usages = [chunk.meta.get("usage") for chunk in chunks if chunk.meta.get("usage") is not None]
+    usage = usages[-1] if usages else None
+
+    meta = {
+        "model": chunks[-1].meta["model"],
+        "index": 0,
+        "finish_reason": finish_reason,
+        "completion_start_time": chunks[0].meta.get("received_at"),  # first chunk received
+        "usage": usage,
+    }
+
+    return ChatMessage.from_assistant(text=text or None, tool_calls=tool_calls, meta=meta)
+
+
 def _parse_streaming_response(
     response_stream: EventStream,
-    streaming_callback: Callable[[StreamingChunk], None],
+    streaming_callback: SyncStreamingCallbackT,
     model: str,
 ) -> List[ChatMessage]:
     """
@@ -143,78 +447,18 @@ def _parse_streaming_response(
     :param model: The model ID used for generation
     :return: List of ChatMessage objects
     """
-    replies = []
-    current_content = ""
-    current_tool_call: Optional[Dict[str, Any]] = None
-    base_meta = {"model": model, "index": 0}
-
+    chunks: List[StreamingChunk] = []
     for event in response_stream:
-        if "contentBlockStart" in event:
-            # Reset accumulators for new message
-            current_content = ""
-            current_tool_call = None
-            block_start = event["contentBlockStart"]
-            if "start" in block_start and "toolUse" in block_start["start"]:
-                tool_start = block_start["start"]["toolUse"]
-                current_tool_call = {
-                    "id": tool_start["toolUseId"],
-                    "name": tool_start["name"],
-                    "arguments": "",  # Will accumulate deltas as string
-                }
-
-        elif "contentBlockDelta" in event:
-            delta = event["contentBlockDelta"]["delta"]
-            if "text" in delta:
-                delta_text = delta["text"]
-                current_content += delta_text
-                streaming_chunk = StreamingChunk(content=delta_text, meta={})
-                streaming_callback(streaming_chunk)
-            elif "toolUse" in delta and current_tool_call:
-                # Accumulate tool use input deltas
-                current_tool_call["arguments"] += delta["toolUse"].get("input", "")
-
-        elif "contentBlockStop" in event:
-            if current_tool_call:
-                # Parse accumulated input if it's a JSON string
-                try:
-                    input_json = json.loads(current_tool_call["arguments"])
-                    current_tool_call["arguments"] = input_json
-                except json.JSONDecodeError:
-                    # Keep as string if not valid JSON
-                    pass
-
-                tool_call = ToolCall(
-                    id=current_tool_call["id"],
-                    tool_name=current_tool_call["name"],
-                    arguments=current_tool_call["arguments"],
-                )
-                replies.append(ChatMessage.from_assistant("", tool_calls=[tool_call], meta=base_meta.copy()))
-            elif current_content:
-                replies.append(ChatMessage.from_assistant(current_content, meta=base_meta.copy()))
-
-        elif "messageStop" in event:
-            # Update finish reason for all replies
-            for reply in replies:
-                reply.meta["finish_reason"] = event["messageStop"].get("stopReason")
-
-        elif "metadata" in event:
-            metadata = event["metadata"]
-            # Update usage stats for all replies
-            for reply in replies:
-                if "usage" in metadata:
-                    usage = metadata["usage"]
-                    reply.meta["usage"] = {
-                        "prompt_tokens": usage.get("inputTokens", 0),
-                        "completion_tokens": usage.get("outputTokens", 0),
-                        "total_tokens": usage.get("totalTokens", 0),
-                    }
-
+        streaming_chunk = _convert_event_to_streaming_chunk(event=event, model=model)
+        streaming_callback(streaming_chunk)
+        chunks.append(streaming_chunk)
+    replies = [_convert_streaming_chunks_to_chat_message(chunks=chunks)]
     return replies
 
 
 async def _parse_streaming_response_async(
     response_stream: EventStream,
-    streaming_callback: StreamingCallbackT,
+    streaming_callback: AsyncStreamingCallbackT,
     model: str,
 ) -> List[ChatMessage]:
     """
@@ -225,70 +469,10 @@ async def _parse_streaming_response_async(
     :param model: The model ID used for generation
     :return: List of ChatMessage objects
     """
-    replies = []
-    current_content = ""
-    current_tool_call: Optional[Dict[str, Any]] = None
-    base_meta = {"model": model, "index": 0}
-
+    chunks: List[StreamingChunk] = []
     async for event in response_stream:
-        if "contentBlockStart" in event:
-            # Reset accumulators for new message
-            current_content = ""
-            current_tool_call = None
-            block_start = event["contentBlockStart"]
-            if "start" in block_start and "toolUse" in block_start["start"]:
-                tool_start = block_start["start"]["toolUse"]
-                current_tool_call = {
-                    "id": tool_start["toolUseId"],
-                    "name": tool_start["name"],
-                    "arguments": "",  # Will accumulate deltas as string
-                }
-
-        elif "contentBlockDelta" in event:
-            delta = event["contentBlockDelta"]["delta"]
-            if "text" in delta:
-                delta_text = delta["text"]
-                current_content += delta_text
-                streaming_chunk = StreamingChunk(content=delta_text, meta={})
-                await streaming_callback(streaming_chunk)
-            elif "toolUse" in delta and current_tool_call:
-                # Accumulate tool use input deltas
-                current_tool_call["arguments"] += delta["toolUse"].get("input", "")
-
-        elif "contentBlockStop" in event:
-            if current_tool_call:
-                # Parse accumulated input if it's a JSON string
-                try:
-                    input_json = json.loads(current_tool_call["arguments"])
-                    current_tool_call["arguments"] = input_json
-                except json.JSONDecodeError:
-                    # Keep as string if not valid JSON
-                    pass
-
-                tool_call = ToolCall(
-                    id=current_tool_call["id"],
-                    tool_name=current_tool_call["name"],
-                    arguments=current_tool_call["arguments"],
-                )
-                replies.append(ChatMessage.from_assistant("", tool_calls=[tool_call], meta=base_meta.copy()))
-            elif current_content:
-                replies.append(ChatMessage.from_assistant(current_content, meta=base_meta.copy()))
-
-        elif "messageStop" in event:
-            # Update finish reason for all replies
-            for reply in replies:
-                reply.meta["finish_reason"] = event["messageStop"].get("stopReason")
-
-        elif "metadata" in event:
-            metadata = event["metadata"]
-            # Update usage stats for all replies
-            for reply in replies:
-                if "usage" in metadata:
-                    usage = metadata["usage"]
-                    reply.meta["usage"] = {
-                        "prompt_tokens": usage.get("inputTokens", 0),
-                        "completion_tokens": usage.get("outputTokens", 0),
-                        "total_tokens": usage.get("totalTokens", 0),
-                    }
-
+        streaming_chunk = _convert_event_to_streaming_chunk(event=event, model=model)
+        await streaming_callback(streaming_chunk)
+        chunks.append(streaming_chunk)
+    replies = [_convert_streaming_chunks_to_chat_message(chunks=chunks)]
     return replies

haystack_integrations/components/generators/amazon_bedrock/generator.py

@@ -1,7 +1,7 @@
 import json
 import re
 import warnings
-from typing import Any, Callable, ClassVar, Dict, List, Literal, Optional, Type, get_args
+from typing import Any, Callable, ClassVar, Dict, List, Literal, Optional, Type, Union, get_args
 
 from botocore.config import Config
 from botocore.exceptions import ClientError
@@ -108,7 +108,7 @@ class AmazonBedrockGenerator:
         boto3_config: Optional[Dict[str, Any]] = None,
         model_family: Optional[MODEL_FAMILIES] = None,
         **kwargs,
-    ):
+    ) -> None:
         """
         Create a new `AmazonBedrockGenerator` instance.
 
@@ -189,7 +189,7 @@ class AmazonBedrockGenerator:
         prompt: str,
         streaming_callback: Optional[Callable[[StreamingChunk], None]] = None,
         generation_kwargs: Optional[Dict[str, Any]] = None,
-    ):
+    ) -> Dict[str, Union[List[str], Dict[str, Any]]]:
         """
         Generates a list of string response to the given prompt.
 

haystack_integrations/components/rankers/amazon_bedrock/__init__.py

@@ -1,3 +1,3 @@
-from .ranker import BedrockRanker
+from .ranker import AmazonBedrockRanker, BedrockRanker
 
-__all__ = ["BedrockRanker"]
+__all__ = ["AmazonBedrockRanker", "BedrockRanker"]

haystack_integrations/components/rankers/amazon_bedrock/ranker.py

@@ -1,3 +1,4 @@
+import warnings
 from typing import Any, Dict, List, Optional
 
 from botocore.exceptions import ClientError
@@ -16,7 +17,7 @@ MAX_NUM_DOCS_FOR_BEDROCK_RANKER = 1000
 
 
 @component
-class BedrockRanker:
+class AmazonBedrockRanker:
     """
     Ranks Documents based on their similarity to the query using Amazon Bedrock's Cohere Rerank model.
 
@@ -30,9 +31,13 @@ class BedrockRanker:
     ```python
     from haystack import Document
     from haystack.utils import Secret
-    from haystack_integrations.components.rankers.amazon_bedrock import BedrockRanker
+    from haystack_integrations.components.rankers.amazon_bedrock import AmazonBedrockRanker
 
-    ranker = BedrockRanker(model="cohere.rerank-v3-5:0", top_k=2, aws_region_name=Secret.from_token("eu-central-1"))
+    ranker = AmazonBedrockRanker(
+        model="cohere.rerank-v3-5:0",
+        top_k=2,
+        aws_region_name=Secret.from_token("eu-central-1")
+    )
 
     docs = [Document(content="Paris"), Document(content="Berlin")]
     query = "What is the capital of germany?"
@@ -40,7 +45,7 @@ class BedrockRanker:
     docs = output["documents"]
     ```
 
-    BedrockRanker uses AWS for authentication. You can use the AWS CLI to authenticate through your IAM.
+    AmazonBedrockRanker uses AWS for authentication. You can use the AWS CLI to authenticate through your IAM.
     For more information on setting up an IAM identity-based policy, see [Amazon Bedrock documentation]
     (https://docs.aws.amazon.com/bedrock/latest/userguide/security_iam_id-based-policy-examples.html).
 
@@ -66,12 +71,12 @@ class BedrockRanker:
         max_chunks_per_doc: Optional[int] = None,
         meta_fields_to_embed: Optional[List[str]] = None,
         meta_data_separator: str = "\n",
-    ):
+    ) -> None:
         if not model:
             msg = "'model' cannot be None or empty string"
             raise ValueError(msg)
         """
-        Creates an instance of the 'BedrockRanker'.
+        Creates an instance of the 'AmazonBedrockRanker'.
 
         :param model: Amazon Bedrock model name for Cohere Rerank. Default is "cohere.rerank-v3-5:0".
         :param top_k: The maximum number of documents to return.
@@ -140,7 +145,7 @@ class BedrockRanker:
         )
 
     @classmethod
-    def from_dict(cls, data: Dict[str, Any]) -> "BedrockRanker":
+    def from_dict(cls, data: Dict[str, Any]) -> "AmazonBedrockRanker":
         """
         Deserializes the component from a dictionary.
 
@@ -173,7 +178,7 @@ class BedrockRanker:
         return concatenated_input_list
 
     @component.output_types(documents=List[Document])
-    def run(self, query: str, documents: List[Document], top_k: Optional[int] = None):
+    def run(self, query: str, documents: List[Document], top_k: Optional[int] = None) -> Dict[str, List[Document]]:
         """
         Use the Amazon Bedrock Reranker to re-rank the list of documents based on the query.
 
@@ -260,3 +265,20 @@ class BedrockRanker:
         except Exception as e:
             msg = f"Error during Amazon Bedrock API call: {e!s}"
             raise AmazonBedrockInferenceError(msg) from e
+
+
+class BedrockRanker(AmazonBedrockRanker):
+    """
+    Deprecated alias for AmazonBedrockRanker.
+    This class will be removed in a future version.
+    Please use AmazonBedrockRanker instead.
+    """
+
+    def __init__(self, *args, **kwargs):
+        warnings.warn(
+            "BedrockRanker is deprecated and will be removed in a future version. "
+            "Please use AmazonBedrockRanker instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        super().__init__(*args, **kwargs)