amazon-bedrock-haystack 3.11.0__py3-none-any.whl → 4.1.0__py3-none-any.whl

{amazon_bedrock_haystack-3.11.0.dist-info → amazon_bedrock_haystack-4.1.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: amazon-bedrock-haystack
-Version: 3.11.0
+Version: 4.1.0
 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>=2.16.0
+Requires-Dist: haystack-ai>=2.17.1
 Description-Content-Type: text/markdown
 
 # amazon-bedrock-haystack
@@ -31,7 +31,7 @@ Description-Content-Type: text/markdown
 
 - [Integration page](https://haystack.deepset.ai/integrations/amazon-bedrock)
 - [Changelog](https://github.com/deepset-ai/haystack-core-integrations/blob/main/integrations/amazon_bedrock/CHANGELOG.md)
------
+---
 
 ## Contributing
 

{amazon_bedrock_haystack-3.11.0.dist-info → amazon_bedrock_haystack-4.1.0.dist-info}/RECORD

@@ -12,12 +12,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=_0dpBoZGY9kgK9zQOTskcjElcTifwhyBAixXDliK-vY,24918
-haystack_integrations/components/generators/amazon_bedrock/chat/utils.py,sha256=g2SZV8LdLobaCZpwWCreBJn1BtS1V3-wQkpisStJrcY,29015
+haystack_integrations/components/generators/amazon_bedrock/chat/chat_generator.py,sha256=qArwfXcforWnPzLXrAW-1hkPFpMy3NSdDyJ5GOta25w,26068
+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=Zrc3BSVkEaXYpliEi6hKG9bqW4J7DNk93p50SuoyT1Q,107
 haystack_integrations/components/rankers/amazon_bedrock/ranker.py,sha256=enAjf2QyDwfpidKkFCdLz954cx-Tjh9emrOS3vINJDg,12344
-amazon_bedrock_haystack-3.11.0.dist-info/METADATA,sha256=5nA_v2Ze5xk1p-RQxbshQ0XGa3LYFljVGvNi2VvKU7o,2225
-amazon_bedrock_haystack-3.11.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-amazon_bedrock_haystack-3.11.0.dist-info/licenses/LICENSE.txt,sha256=B05uMshqTA74s-0ltyHKI6yoPfJ3zYgQbvcXfDVGFf8,10280
-amazon_bedrock_haystack-3.11.0.dist-info/RECORD,,
+amazon_bedrock_haystack-4.1.0.dist-info/METADATA,sha256=P6e8VfoRQ0hZrA6gkaTwjCCxZez2NqKjHNAJuFdwK6c,2222
+amazon_bedrock_haystack-4.1.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+amazon_bedrock_haystack-4.1.0.dist-info/licenses/LICENSE.txt,sha256=B05uMshqTA74s-0ltyHKI6yoPfJ3zYgQbvcXfDVGFf8,10280
+amazon_bedrock_haystack-4.1.0.dist-info/RECORD,,

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

@@ -27,6 +27,7 @@ from haystack_integrations.components.generators.amazon_bedrock.chat.utils impor
     _parse_completion_response,
     _parse_streaming_response,
     _parse_streaming_response_async,
+    _validate_guardrail_config,
 )
 
 logger = logging.getLogger(__name__)
@@ -154,10 +155,11 @@ class AmazonBedrockChatGenerator:
         aws_region_name: Optional[Secret] = Secret.from_env_var(["AWS_DEFAULT_REGION"], strict=False),  # noqa: B008
         aws_profile_name: Optional[Secret] = Secret.from_env_var(["AWS_PROFILE"], strict=False),  # noqa: B008
         generation_kwargs: Optional[Dict[str, Any]] = None,
-        stop_words: Optional[List[str]] = None,
         streaming_callback: Optional[StreamingCallbackT] = None,
         boto3_config: Optional[Dict[str, Any]] = None,
         tools: Optional[Union[List[Tool], Toolset]] = None,
+        *,
+        guardrail_config: Optional[Dict[str, str]] = None,
     ) -> None:
         """
         Initializes the `AmazonBedrockChatGenerator` with the provided parameters. The parameters are passed to the
@@ -179,10 +181,6 @@ class AmazonBedrockChatGenerator:
         :param generation_kwargs: Keyword arguments sent to the model. These parameters are specific to a model.
             You can find the model specific arguments in the AWS Bedrock API
             [documentation](https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters.html).
-        :param stop_words: A list of stop words that stop the model from generating more text
-            when encountered. You can provide them using this parameter or using the model's `generation_kwargs`
-            under a model's specific key for stop words.
-            For example, you can provide stop words for Anthropic Claude in the `stop_sequences` key.
         :param streaming_callback: A callback function called when a new token is received from the stream.
             By default, the model is not set up for streaming. To enable streaming, set this parameter to a callback
             function that handles the streaming chunks. The callback function receives a
@@ -190,6 +188,19 @@ class AmazonBedrockChatGenerator:
             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 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).
+            or, in streaming mode (when `streaming_callback` is set),
+            [GuardrailStreamConfiguration](https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_GuardrailStreamConfiguration.html).
+            If `trace` is set to `enabled`, the guardrail trace will be included under the `trace` key in the `meta`
+            attribute of the resulting `ChatMessage`.
+            Note: Enabling guardrails in streaming mode may introduce additional latency.
+            To manage this, you can adjust the `streamProcessingMode` parameter.
+            See the
+            [Guardrails Streaming documentation](https://docs.aws.amazon.com/bedrock/latest/userguide/guardrails-streaming.html)
+            for more information.
+
 
         :raises ValueError: If the model name is empty or None.
         :raises AmazonBedrockConfigurationError: If the AWS environment is not configured correctly or the model is
@@ -204,12 +215,15 @@ class AmazonBedrockChatGenerator:
         self.aws_session_token = aws_session_token
         self.aws_region_name = aws_region_name
         self.aws_profile_name = aws_profile_name
-        self.stop_words = stop_words or []
         self.streaming_callback = streaming_callback
         self.boto3_config = boto3_config
+
         _check_duplicate_tool_names(list(tools or []))  # handles Toolset as well
         self.tools = tools
 
+        _validate_guardrail_config(guardrail_config=guardrail_config, streaming=streaming_callback is not None)
+        self.guardrail_config = guardrail_config
+
         def resolve_secret(secret: Optional[Secret]) -> Optional[str]:
             return secret.resolve_value() if secret else None
 
@@ -237,7 +251,6 @@ class AmazonBedrockChatGenerator:
             raise AmazonBedrockConfigurationError(msg) from exception
 
         self.generation_kwargs = generation_kwargs or {}
-        self.stop_words = stop_words or []
         self.async_session: Optional[aioboto3.Session] = None
 
     def _get_async_session(self) -> aioboto3.Session:
@@ -291,11 +304,11 @@ class AmazonBedrockChatGenerator:
             aws_region_name=self.aws_region_name.to_dict() if self.aws_region_name else None,
             aws_profile_name=self.aws_profile_name.to_dict() if self.aws_profile_name else None,
             model=self.model,
-            stop_words=self.stop_words,
             generation_kwargs=self.generation_kwargs,
             streaming_callback=callback_name,
             boto3_config=self.boto3_config,
             tools=serialize_tools_or_toolset(self.tools),
+            guardrail_config=self.guardrail_config,
         )
 
     @classmethod
@@ -308,6 +321,12 @@ class AmazonBedrockChatGenerator:
             Instance of `AmazonBedrockChatGenerator`.
         """
         init_params = data.get("init_parameters", {})
+
+        stop_words = init_params.pop("stop_words", None)
+        msg = "stop_words parameter will be ignored. Use the `stopSequences` key in `generation_kwargs` instead."
+        if stop_words:
+            logger.warning(msg)
+
         serialized_callback_handler = init_params.get("streaming_callback")
         if serialized_callback_handler:
             data["init_parameters"]["streaming_callback"] = deserialize_callable(serialized_callback_handler)
@@ -387,6 +406,8 @@ class AmazonBedrockChatGenerator:
             params["toolConfig"] = tool_config
         if additional_fields:
             params["additionalModelRequestFields"] = additional_fields
+        if self.guardrail_config:
+            params["guardrailConfig"] = self.guardrail_config
 
         # overloads that exhaust finite Literals(bool) not treated as exhaustive
         # see https://github.com/python/mypy/issues/14764

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

@@ -5,16 +5,20 @@ from typing import Any, Dict, List, Optional, Tuple
 
 from botocore.eventstream import EventStream
 from haystack import logging
+from haystack.components.generators.utils import _convert_streaming_chunks_to_chat_message
 from haystack.dataclasses import (
     AsyncStreamingCallbackT,
     ChatMessage,
     ChatRole,
     ComponentInfo,
+    FinishReason,
     ImageContent,
+    ReasoningContent,
     StreamingChunk,
     SyncStreamingCallbackT,
     TextContent,
     ToolCall,
+    ToolCallDelta,
 )
 from haystack.tools import Tool
 
@@ -24,6 +28,16 @@ logger = logging.getLogger(__name__)
 # see https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ImageBlock.html for supported formats
 IMAGE_SUPPORTED_FORMATS = ["png", "jpeg", "gif", "webp"]
 
+# see https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_MessageStopEvent.html
+FINISH_REASON_MAPPING: Dict[str, FinishReason] = {
+    "end_turn": "stop",
+    "stop_sequence": "stop",
+    "max_tokens": "length",
+    "guardrail_intervened": "content_filter",
+    "content_filtered": "content_filter",
+    "tool_use": "tool_calls",
+}
+
 
 # Haystack to Bedrock util methods
 def _format_tools(tools: Optional[List[Tool]] = None) -> Optional[Dict[str, Any]]:
@@ -57,8 +71,8 @@ def _format_tool_call_message(tool_call_message: ChatMessage) -> Dict[str, Any]:
     content: List[Dict[str, Any]] = []
 
     # tool call messages can contain reasoning content
-    if reasoning_contents := tool_call_message.meta.get("reasoning_contents"):
-        content.extend(_format_reasoning_contents(reasoning_contents=reasoning_contents))
+    if reasoning_content := tool_call_message.reasoning:
+        content.extend(_format_reasoning_content(reasoning_content=reasoning_content))
 
     # Tool call message can contain text
     if tool_call_message.text:
@@ -162,16 +176,16 @@ def _repair_tool_result_messages(bedrock_formatted_messages: List[Dict[str, Any]
     return [msg for _, msg in repaired_bedrock_formatted_messages]
 
 
-def _format_reasoning_contents(reasoning_contents: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+def _format_reasoning_content(reasoning_content: ReasoningContent) -> List[Dict[str, Any]]:
     """
-    Format reasoning contents to match Bedrock's expected structure.
+    Format ReasoningContent to match Bedrock's expected structure.
 
-    :param reasoning_contents: List of reasoning content dictionaries from Haystack ChatMessage metadata.
+    :param reasoning_content: ReasoningContent object containing reasoning contents to format.
     :returns: List of formatted reasoning content dictionaries for Bedrock.
     """
     formatted_contents = []
-    for reasoning_content in reasoning_contents:
-        formatted_content = {"reasoningContent": reasoning_content["reasoning_content"]}
+    for content in reasoning_content.extra.get("reasoning_contents", []):
+        formatted_content = {"reasoningContent": content["reasoning_content"]}
         if reasoning_text := formatted_content["reasoningContent"].pop("reasoning_text", None):
             formatted_content["reasoningContent"]["reasoningText"] = reasoning_text
         if redacted_content := formatted_content["reasoningContent"].pop("redacted_content", None):
@@ -192,8 +206,8 @@ def _format_text_image_message(message: ChatMessage) -> Dict[str, Any]:
 
     bedrock_content_blocks: List[Dict[str, Any]] = []
     # Add reasoning content if available as the first content block
-    if message.meta.get("reasoning_contents"):
-        bedrock_content_blocks.extend(_format_reasoning_contents(reasoning_contents=message.meta["reasoning_contents"]))
+    if message.reasoning:
+        bedrock_content_blocks.extend(_format_reasoning_content(reasoning_content=message.reasoning))
 
     for part in content_parts:
         if isinstance(part, TextContent):
@@ -259,6 +273,7 @@ def _parse_completion_response(response_body: Dict[str, Any], model: str) -> Lis
     :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"]:
         message = response_body["output"]["message"]
@@ -266,10 +281,10 @@ def _parse_completion_response(response_body: Dict[str, Any], model: str) -> Lis
             content_blocks = message["content"]
 
             # Common meta information
-            base_meta = {
+            meta = {
                 "model": model,
                 "index": 0,
-                "finish_reason": response_body.get("stopReason"),
+                "finish_reason": FINISH_REASON_MAPPING.get(response_body.get("stopReason", "")),
                 "usage": {
                     # OpenAI's format for usage for cross ChatGenerator compatibility
                     "prompt_tokens": response_body.get("usage", {}).get("inputTokens", 0),
@@ -277,6 +292,9 @@ def _parse_completion_response(response_body: Dict[str, Any], model: str) -> Lis
                     "total_tokens": response_body.get("usage", {}).get("totalTokens", 0),
                 },
             }
+            # guardrail trace
+            if "trace" in response_body:
+                meta["trace"] = response_body["trace"]
 
             # Process all content blocks and combine them into a single message
             text_content = []
@@ -303,11 +321,26 @@ def _parse_completion_response(response_body: Dict[str, Any], model: str) -> Lis
                         reasoning_content["redacted_content"] = reasoning_content.pop("redactedContent")
                     reasoning_contents.append({"reasoning_content": reasoning_content})
 
-            # If reasoning contents were found, add them to the base meta
-            base_meta.update({"reasoning_contents": reasoning_contents})
+            reasoning_text = ""
+            for content in reasoning_contents:
+                if "reasoning_text" in content["reasoning_content"]:
+                    reasoning_text += content["reasoning_content"]["reasoning_text"]["text"]
+                elif "redacted_content" in content["reasoning_content"]:
+                    reasoning_text += "[REDACTED]"
 
             # Create a single ChatMessage with combined text and tool calls
-            replies.append(ChatMessage.from_assistant(" ".join(text_content), tool_calls=tool_calls, meta=base_meta))
+            replies.append(
+                ChatMessage.from_assistant(
+                    " ".join(text_content),
+                    tool_calls=tool_calls,
+                    meta=meta,
+                    reasoning=ReasoningContent(
+                        reasoning_text=reasoning_text, extra={"reasoning_contents": reasoning_contents}
+                    )
+                    if reasoning_contents
+                    else None,
+                )
+            )
 
     return replies
 
@@ -326,11 +359,11 @@ def _convert_event_to_streaming_chunk(
     :param component_info: ComponentInfo object
     :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()}
-    )
+    base_meta = {"model": model, "received_at": datetime.now(timezone.utc).isoformat()}
+    streaming_chunk = StreamingChunk(content="", meta=base_meta)
 
     if "contentBlockStart" in event:
         # contentBlockStart always has the key "contentBlockIndex"
@@ -340,26 +373,15 @@ def _convert_event_to_streaming_chunk(
             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(),
-                },
+                index=block_idx,
+                tool_calls=[
+                    ToolCallDelta(
+                        index=block_idx,
+                        id=tool_start["toolUseId"],
+                        tool_name=tool_start["name"],
+                    )
+                ],
+                meta=base_meta,
             )
 
     elif "contentBlockDelta" in event:
@@ -370,39 +392,22 @@ def _convert_event_to_streaming_chunk(
         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(),
-                },
+                index=block_idx,
+                meta=base_meta,
             )
         # 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(),
-                },
+                index=block_idx,
+                tool_calls=[
+                    ToolCallDelta(
+                        index=block_idx,
+                        arguments=delta["toolUse"].get("input", ""),
+                    )
+                ],
+                meta=base_meta,
             )
         # This is for accumulating reasoning content deltas
         elif "reasoningContent" in delta:
@@ -411,55 +416,45 @@ def _convert_event_to_streaming_chunk(
                 reasoning_content["redacted_content"] = reasoning_content.pop("redactedContent")
             streaming_chunk = StreamingChunk(
                 content="",
+                index=block_idx,
                 meta={
-                    "model": model,
-                    "index": 0,
-                    "tool_calls": None,
-                    "finish_reason": None,
-                    "received_at": datetime.now(timezone.utc).isoformat(),
+                    **base_meta,
                     "reasoning_contents": [{"index": block_idx, "reasoning_content": reasoning_content}],
                 },
             )
 
     elif "messageStop" in event:
-        finish_reason = event["messageStop"].get("stopReason")
+        finish_reason = FINISH_REASON_MAPPING.get(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(),
-            },
+            finish_reason=finish_reason,
+            meta=base_meta,
         )
 
-    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),
-                },
-            },
-        )
+    elif "metadata" in event:
+        event_meta = event["metadata"]
+        chunk_meta: Dict[str, Any] = {**base_meta}
+
+        if "usage" in event_meta:
+            usage = event_meta["usage"]
+            chunk_meta["usage"] = {
+                "prompt_tokens": usage.get("inputTokens", 0),
+                "completion_tokens": usage.get("outputTokens", 0),
+                "total_tokens": usage.get("totalTokens", 0),
+            }
+        if "trace" in event_meta:
+            chunk_meta["trace"] = event_meta["trace"]
+
+        # Only create chunk if we added usage or trace data
+        if len(chunk_meta) > len(base_meta):
+            streaming_chunk = StreamingChunk(content="", meta=chunk_meta)
 
     streaming_chunk.component_info = component_info
 
     return streaming_chunk
 
 
-def _process_reasoning_contents(chunks: List[StreamingChunk]) -> List[Dict[str, Any]]:
+def _process_reasoning_contents(chunks: List[StreamingChunk]) -> Optional[ReasoningContent]:
     """
     Process reasoning contents from a list of StreamingChunk objects into the Bedrock expected format.
 
@@ -491,6 +486,8 @@ def _process_reasoning_contents(chunks: List[StreamingChunk]) -> List[Dict[str,
                     )
                 if redacted_content:
                     formatted_reasoning_contents.append({"reasoning_content": {"redacted_content": redacted_content}})
+
+                # Reset accumulators for new group
                 reasoning_text = ""
                 reasoning_signature = None
                 redacted_content = None
@@ -516,85 +513,22 @@ def _process_reasoning_contents(chunks: List[StreamingChunk]) -> List[Dict[str,
         if redacted_content:
             formatted_reasoning_contents.append({"reasoning_content": {"redacted_content": redacted_content}})
 
-    return formatted_reasoning_contents
-
-
-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.
-    """
-    # Join all text content from the chunks
-    text = "".join([chunk.content for chunk in chunks])
-
-    # If reasoning content is present in any chunk, accumulate it
-    reasoning_contents = _process_reasoning_contents(chunks=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.get("arguments", "{}")) if call_data.get("arguments") else {}
-            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,
-        "reasoning_contents": reasoning_contents,
-    }
-
-    return ChatMessage.from_assistant(text=text or None, tool_calls=tool_calls, meta=meta)
+    # Combine all reasoning texts into a single string for the main reasoning_text field
+    final_reasoning_text = ""
+    for content in formatted_reasoning_contents:
+        if "reasoning_text" in content["reasoning_content"]:
+            # mypy somehow thinks that content["reasoning_content"]["reasoning_text"]["text"] can be of type None
+            final_reasoning_text += content["reasoning_content"]["reasoning_text"]["text"]  # type: ignore[operator]
+        elif "redacted_content" in content["reasoning_content"]:
+            final_reasoning_text += "[REDACTED]"
+
+    return (
+        ReasoningContent(
+            reasoning_text=final_reasoning_text, extra={"reasoning_contents": formatted_reasoning_contents}
+        )
+        if formatted_reasoning_contents
+        else None
+    )
 
 
 def _parse_streaming_response(
@@ -612,13 +546,34 @@ def _parse_streaming_response(
     :param component_info: ComponentInfo object
     :return: List of ChatMessage objects
     """
+    content_block_idxs = set()
     chunks: List[StreamingChunk] = []
     for event in response_stream:
         streaming_chunk = _convert_event_to_streaming_chunk(event=event, model=model, component_info=component_info)
+        content_block_idx = streaming_chunk.index
+        if content_block_idx is not None and content_block_idx not in content_block_idxs:
+            streaming_chunk.start = True
+            content_block_idxs.add(content_block_idx)
         streaming_callback(streaming_chunk)
         chunks.append(streaming_chunk)
-    replies = [_convert_streaming_chunks_to_chat_message(chunks=chunks)]
-    return replies
+
+    reply = _convert_streaming_chunks_to_chat_message(chunks=chunks)
+
+    # both the reasoning content and the trace are ignored in _convert_streaming_chunks_to_chat_message
+    # so we need to process them separately
+    reasoning_content = _process_reasoning_contents(chunks=chunks)
+    if chunks[-1].meta and "trace" in chunks[-1].meta:
+        reply.meta["trace"] = chunks[-1].meta["trace"]
+
+    reply = ChatMessage.from_assistant(
+        text=reply.text,
+        meta=reply.meta,
+        name=reply.name,
+        tool_calls=reply.tool_calls,
+        reasoning=reasoning_content,
+    )
+
+    return [reply]
 
 
 async def _parse_streaming_response_async(
@@ -636,10 +591,44 @@ async def _parse_streaming_response_async(
     :param component_info: ComponentInfo object
     :return: List of ChatMessage objects
     """
+    content_block_idxs = set()
     chunks: List[StreamingChunk] = []
     async for event in response_stream:
         streaming_chunk = _convert_event_to_streaming_chunk(event=event, model=model, component_info=component_info)
+        content_block_idx = streaming_chunk.index
+        if content_block_idx is not None and content_block_idx not in content_block_idxs:
+            streaming_chunk.start = True
+            content_block_idxs.add(content_block_idx)
         await streaming_callback(streaming_chunk)
         chunks.append(streaming_chunk)
-    replies = [_convert_streaming_chunks_to_chat_message(chunks=chunks)]
-    return replies
+    reply = _convert_streaming_chunks_to_chat_message(chunks=chunks)
+    reasoning_content = _process_reasoning_contents(chunks=chunks)
+    reply = ChatMessage.from_assistant(
+        text=reply.text,
+        meta=reply.meta,
+        name=reply.name,
+        tool_calls=reply.tool_calls,
+        reasoning=reasoning_content,
+    )
+    return [reply]
+
+
+def _validate_guardrail_config(guardrail_config: Optional[Dict[str, str]] = None, streaming: bool = False) -> None:
+    """
+    Validate the guardrail configuration.
+
+    :param guardrail_config: The guardrail configuration.
+    :param streaming: Whether the streaming is enabled.
+
+    :raises ValueError: If the guardrail configuration is invalid.
+    """
+    if guardrail_config is None:
+        return
+
+    required_fields = {"guardrailIdentifier", "guardrailVersion"}
+    if not required_fields.issubset(guardrail_config):
+        msg = "`guardrailIdentifier` and `guardrailVersion` fields are required in guardrail configuration."
+        raise ValueError(msg)
+    if not streaming and "streamProcessingMode" in guardrail_config:
+        msg = "`streamProcessingMode` field is only supported for streaming (when `streaming_callback` is not None)."
+        raise ValueError(msg)