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

{amazon_bedrock_haystack-4.0.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: 4.0.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

{amazon_bedrock_haystack-4.0.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=eF2wldu1IppL64f01N3PIa9_-BZEolQzEz9NjXvEFTQ,25810
+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-4.0.0.dist-info/METADATA,sha256=8iHYVwqO_nLLbLOh59yoqxwnskmfBiDupxxSTWiCvcc,2222
-amazon_bedrock_haystack-4.0.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-amazon_bedrock_haystack-4.0.0.dist-info/licenses/LICENSE.txt,sha256=B05uMshqTA74s-0ltyHKI6yoPfJ3zYgQbvcXfDVGFf8,10280
-amazon_bedrock_haystack-4.0.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

@@ -273,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"]
@@ -280,7 +281,7 @@ 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": FINISH_REASON_MAPPING.get(response_body.get("stopReason", "")),
@@ -291,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 = []
@@ -329,7 +333,7 @@ def _parse_completion_response(response_body: Dict[str, Any], model: str) -> Lis
                 ChatMessage.from_assistant(
                     " ".join(text_content),
                     tool_calls=tool_calls,
-                    meta=base_meta,
+                    meta=meta,
                     reasoning=ReasoningContent(
                         reasoning_text=reasoning_text, extra={"reasoning_contents": reasoning_contents}
                     )
@@ -355,6 +359,7 @@ 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)
     base_meta = {"model": model, "received_at": datetime.now(timezone.utc).isoformat()}
@@ -426,19 +431,23 @@ def _convert_event_to_streaming_chunk(
             meta=base_meta,
         )
 
-    elif "metadata" in event and "usage" in event["metadata"]:
-        metadata = event["metadata"]
-        streaming_chunk = StreamingChunk(
-            content="",
-            meta={
-                **base_meta,
-                "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
 
@@ -547,8 +556,15 @@ def _parse_streaming_response(
             content_block_idxs.add(content_block_idx)
         streaming_callback(streaming_chunk)
         chunks.append(streaming_chunk)
+
     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,
@@ -556,6 +572,7 @@ def _parse_streaming_response(
         tool_calls=reply.tool_calls,
         reasoning=reasoning_content,
     )
+
     return [reply]
 
 
@@ -594,3 +611,24 @@ async def _parse_streaming_response_async(
         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)