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

{amazon_bedrock_haystack-4.0.0.dist-info → amazon_bedrock_haystack-4.2.0.dist-info}/METADATA

@@ -1,7 +1,7 @@
 Metadata-Version: 2.4
 Name: amazon-bedrock-haystack
-Version: 4.0.0
-Summary: An integration of Amazon Bedrock as an AmazonBedrockGenerator component.
+Version: 4.2.0
+Summary: An integration of AWS S3 and Bedrock as a Downloader and Generator components.
 Project-URL: Documentation, https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/amazon_bedrock#readme
 Project-URL: Issues, https://github.com/deepset-ai/haystack-core-integrations/issues
 Project-URL: Source, https://github.com/deepset-ai/haystack-core-integrations/tree/main/integrations/amazon_bedrock

{amazon_bedrock_haystack-4.0.0.dist-info → amazon_bedrock_haystack-4.2.0.dist-info}/RECORD

@@ -1,7 +1,13 @@
 haystack_integrations/common/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 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=ASAwEhInF9F6rhL4CbXFQUFU1pSdscWvG6jcrXkEUhc,2735
+haystack_integrations/common/amazon_bedrock/errors.py,sha256=47w_rg3JTGJ5QfsUELrjYBeuLajxUmQ7-zOFvdz4mT8,856
+haystack_integrations/common/amazon_bedrock/utils.py,sha256=UIJVl1e_hlP9mQr_YvVTFklbc_kqOgI7DnSxcAuSLx0,2846
+haystack_integrations/common/s3/__init__.py,sha256=6GZ8Y3Lw0rLOsOAqi6Tu5mZC977UzQvgDxKpOWr8IQw,110
+haystack_integrations/common/s3/errors.py,sha256=BrTDLdhQvAuQutyg35cFyP5h8PNkDEieLwehi58UqAU,452
+haystack_integrations/common/s3/utils.py,sha256=OJupFj54aQmg6S8VuVq6Lc2qpFZyyJajRVIpwe3_6iA,4744
+haystack_integrations/components/downloaders/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+haystack_integrations/components/downloaders/s3/__init__.py,sha256=2BOd3_N0kGqRJGH-ENrTJqOqzqHryRYaSuNqpLYKMFo,179
+haystack_integrations/components/downloaders/s3/s3_downloader.py,sha256=kptTCSry_uEYtNAca_pU7zQJs_LJwwJKYjuYDVJrZRE,11220
 haystack_integrations/components/embedders/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 haystack_integrations/components/embedders/amazon_bedrock/__init__.py,sha256=7GlhHJ4jFHCxq5QN5losGuGtrGNjvEx2dSQvEYD2yG0,408
 haystack_integrations/components/embedders/amazon_bedrock/document_embedder.py,sha256=DD34-HAGwGwTU7KWGqKXXlFdwIs21JavBRDHrBqC-m4,13060
@@ -12,12 +18,12 @@ haystack_integrations/components/generators/amazon_bedrock/__init__.py,sha256=lv
 haystack_integrations/components/generators/amazon_bedrock/adapters.py,sha256=yBC-3YwV6qAwSXMtdZiLSYh2lUpPQIDy7Efl7w-Cu-k,19640
 haystack_integrations/components/generators/amazon_bedrock/generator.py,sha256=Brzw0XvtPJhz2kR2I3liAqWHRmDR6p5HzJerEAPhoJU,14743
 haystack_integrations/components/generators/amazon_bedrock/chat/__init__.py,sha256=6GZ8Y3Lw0rLOsOAqi6Tu5mZC977UzQvgDxKpOWr8IQw,110
-haystack_integrations/components/generators/amazon_bedrock/chat/chat_generator.py,sha256=_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.2.0.dist-info/METADATA,sha256=HM32juznvZMTZgwNyS34vsLTZeBsLxaru7sUmGz4xQA,2228
+amazon_bedrock_haystack-4.2.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+amazon_bedrock_haystack-4.2.0.dist-info/licenses/LICENSE.txt,sha256=B05uMshqTA74s-0ltyHKI6yoPfJ3zYgQbvcXfDVGFf8,10280
+amazon_bedrock_haystack-4.2.0.dist-info/RECORD,,

haystack_integrations/common/amazon_bedrock/errors.py

@@ -1,3 +1,8 @@
+# SPDX-FileCopyrightText: 2023-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+
 class AmazonBedrockError(Exception):
     """
     Any error generated by the Amazon Bedrock integration.

haystack_integrations/common/amazon_bedrock/utils.py

@@ -1,3 +1,7 @@
+# SPDX-FileCopyrightText: 2023-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
 from typing import Any, Optional, Union
 
 import aioboto3

haystack_integrations/common/s3/__init__.py

@@ -0,0 +1,3 @@
+# SPDX-FileCopyrightText: 2023-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0

haystack_integrations/common/s3/errors.py

@@ -0,0 +1,15 @@
+# SPDX-FileCopyrightText: 2023-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+
+class S3Error(Exception):
+    """Exception for issues that occur in the S3 based components"""
+
+
+class S3ConfigurationError(S3Error):
+    """Exception raised when AmazonS3 node is not configured correctly"""
+
+
+class S3StorageError(S3Error):
+    """This exception is raised when an error occurs while interacting with a S3Storage object."""

haystack_integrations/common/s3/utils.py

@@ -0,0 +1,120 @@
+# SPDX-FileCopyrightText: 2023-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+from dataclasses import dataclass
+from http import HTTPStatus
+from pathlib import Path
+from typing import Optional
+
+from boto3.session import Session
+from botocore.config import Config
+from botocore.exceptions import ClientError, NoCredentialsError, PartialCredentialsError
+
+from haystack_integrations.common.s3.errors import S3ConfigurationError, S3StorageError
+
+
+@dataclass
+class S3Storage:
+    """This class provides a storage class for downloading files from an AWS S3 bucket."""
+
+    def __init__(
+        self,
+        s3_bucket: str,
+        session: Session,
+        s3_prefix: Optional[str] = None,
+        endpoint_url: Optional[str] = None,
+        config: Optional[Config] = None,
+    ) -> None:
+        """
+        Initializes the S3Storage object with the provided parameters.
+
+        :param s3_bucket: The name of the S3 bucket to download files from.
+        :param session: The session to use for the S3 client.
+        :param s3_prefix: The optional prefix of the files in the S3 bucket.
+        Can be used to specify folder or naming structure.
+            For example, if the file is in the folder "folder/subfolder/file.txt",
+            the s3_prefix should be "folder/subfolder/". If the file is in the root of the S3 bucket,
+            the s3_prefix should be None.
+        :param endpoint_url: The endpoint URL of the S3 bucket to download files from.
+        :param config: The configuration to use for the S3 client.
+        """
+
+        self.s3_bucket = s3_bucket
+        self.s3_prefix = s3_prefix
+        self.endpoint_url = endpoint_url
+        self.session = session
+        self.config = config
+
+        try:
+            self._client = self.session.client("s3", endpoint_url=self.endpoint_url, config=self.config)
+        except Exception as e:
+            msg = f"Failed to create S3 session client: {e}"
+            raise S3ConfigurationError(msg) from e
+
+    def download(self, key: str, local_file_path: Path) -> None:
+        """Download a file from S3.
+
+        :param key: The key of the file to download.
+        :param local_file_path: The folder path to download the file to.
+        It will be created if it does not exist. The file will be downloaded to
+        the folder with the same name as the key.
+        :raises S3ConfigurationError: If the S3 session client cannot be created.
+        :raises S3StorageError: If the file does not exist in the S3 bucket
+        or the file cannot be downloaded.
+        """
+
+        if self.s3_prefix:
+            s3_key = f"{self.s3_prefix}{key}"
+        else:
+            s3_key = key
+
+        try:
+            self._client.download_file(self.s3_bucket, s3_key, str(local_file_path))
+
+        except (NoCredentialsError, PartialCredentialsError) as e:
+            msg = (
+                f"Missing AWS credentials. Please check your AWS credentials (access key, secret key, region)."
+                f"Error: {e}"
+            )
+            raise S3ConfigurationError(msg) from e
+
+        except ClientError as e:
+            error_code = int(e.response["Error"]["Code"])
+
+            if error_code == HTTPStatus.FORBIDDEN:
+                msg = (
+                    f"Failed to access S3 bucket {self.s3_bucket!r}. "
+                    f"Please check your AWS credentials (access key, secret key, region) and ensure "
+                    f"they have the necessary S3 permissions. "
+                    f"Error: {e}"
+                )
+                raise S3ConfigurationError(msg) from e
+
+            elif error_code == HTTPStatus.NOT_FOUND:
+                msg = f"The object {s3_key!r} does not exist in the S3 bucket {self.s3_bucket!r}. \n Error: {e}"
+                raise S3StorageError(msg) from e
+            else:
+                msg = f"Failed to download file {s3_key!r} from S3. Error: {e}"
+                raise S3StorageError(msg) from e
+
+    @classmethod
+    def from_env(cls, *, session: Session, config: Config) -> "S3Storage":
+        """Create a S3Storage object from environment variables."""
+        s3_bucket = os.getenv("S3_DOWNLOADER_BUCKET")
+        if not s3_bucket:
+            msg = (
+                "Missing environment variable S3_DOWNLOADER_BUCKET."
+                "Please set it to the name of the S3 bucket to download files from."
+            )
+            raise ValueError(msg)
+        s3_prefix = os.getenv("S3_DOWNLOADER_PREFIX") or None
+        endpoint_url = os.getenv("AWS_ENDPOINT_URL") or None
+        return cls(
+            s3_bucket=s3_bucket,
+            s3_prefix=s3_prefix,
+            endpoint_url=endpoint_url,
+            session=session,
+            config=config,
+        )

haystack_integrations/components/downloaders/s3/__init__.py

@@ -0,0 +1,7 @@
+# SPDX-FileCopyrightText: 2023-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from .s3_downloader import S3Downloader
+
+__all__ = ["S3Downloader"]

haystack_integrations/components/downloaders/s3/s3_downloader.py

@@ -0,0 +1,246 @@
+# SPDX-FileCopyrightText: 2023-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+from botocore.config import Config
+from haystack import component, default_from_dict, default_to_dict, logging
+from haystack.dataclasses import Document
+from haystack.utils.auth import Secret, deserialize_secrets_inplace
+
+from haystack_integrations.common.amazon_bedrock.utils import get_aws_session
+from haystack_integrations.common.s3.utils import S3Storage
+
+logger = logging.getLogger(__name__)
+
+
+@component
+class S3Downloader:
+    """
+    A component for downloading files from AWS S3 Buckets to local filesystem.
+    Supports filtering by file extensions.
+    """
+
+    def __init__(
+        self,
+        *,
+        aws_access_key_id: Optional[Secret] = Secret.from_env_var("AWS_ACCESS_KEY_ID", strict=False),  # noqa: B008
+        aws_secret_access_key: Optional[Secret] = Secret.from_env_var(  # noqa: B008
+            "AWS_SECRET_ACCESS_KEY", strict=False
+        ),
+        aws_session_token: Optional[Secret] = Secret.from_env_var("AWS_SESSION_TOKEN", strict=False),  # noqa: B008
+        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
+        boto3_config: Optional[Dict[str, Any]] = None,
+        file_root_path: Optional[str] = None,
+        file_extensions: Optional[List[str]] = None,
+        file_name_meta_key: str = "file_name",
+        max_workers: int = 32,
+        max_cache_size: int = 100,
+    ) -> None:
+        """
+        Initializes the `S3Downloader` with the provided parameters.
+
+        Note that the AWS credentials are not required if the AWS environment is configured correctly. These are loaded
+        automatically from the environment or the AWS configuration file and do not need to be provided explicitly via
+        the constructor. If the AWS environment is not configured users need to provide the AWS credentials via the
+        constructor. Three required parameters are `aws_access_key_id`, `aws_secret_access_key`,
+        and `aws_region_name`.
+
+        :param aws_access_key_id: AWS access key ID.
+        :param aws_secret_access_key: AWS secret access key.
+        :param aws_session_token: AWS session token.
+        :param aws_region_name: AWS region name.
+        :param aws_profile_name: AWS profile name.
+        :param boto3_config: The configuration for the boto3 client.
+        :param file_root_path: The path where the file will be downloaded.
+            Can be set through this parameter or the `FILE_ROOT_PATH` environment variable.
+            If none of them is set, a `ValueError` is raised.
+        :param file_extensions: The file extensions that are permitted to be downloaded.
+            By default, all file extensions are allowed.
+        :param max_workers: The maximum number of workers to use for concurrent downloads.
+        :param max_cache_size: The maximum number of files to cache.
+        :param file_name_meta_key: The name of the meta key that contains the file name to download.
+            By default, the `Document.meta["file_name"]` is used. If you want to use a
+            different key in `Document.meta`, you can set it here.
+        :raises ValueError: If the `file_root_path` is not set through
+            the constructor or the `FILE_ROOT_PATH` environment variable.
+
+        """
+
+        # Set up download directory
+        file_root_path = file_root_path or os.getenv("FILE_ROOT_PATH")
+
+        if file_root_path is None:
+            msg = (
+                "The path where files will be downloaded is not set. Please set the "
+                "`file_root_path` init parameter or the `FILE_ROOT_PATH` environment variable."
+            )
+            raise ValueError(msg)
+
+        self.file_root_path = Path(file_root_path)
+
+        self.aws_access_key_id = aws_access_key_id
+        self.aws_secret_access_key = aws_secret_access_key
+        self.aws_region_name = aws_region_name
+        self.aws_session_token = aws_session_token
+        self.aws_profile_name = aws_profile_name
+        self.boto3_config = boto3_config
+        self.file_extensions = [e.lower() for e in file_extensions] if file_extensions else None
+        self.max_workers = max_workers
+        self.max_cache_size = max_cache_size
+        self.file_name_meta_key = file_name_meta_key
+
+        self._storage: Optional[S3Storage] = None
+
+        def resolve_secret(secret: Optional[Secret]) -> Optional[str]:
+            return secret.resolve_value() if secret else None
+
+        self._session = get_aws_session(
+            aws_access_key_id=resolve_secret(aws_access_key_id),
+            aws_secret_access_key=resolve_secret(aws_secret_access_key),
+            aws_session_token=resolve_secret(aws_session_token),
+            aws_region_name=resolve_secret(aws_region_name),
+            aws_profile_name=resolve_secret(aws_profile_name),
+        )
+        self._config = Config(
+            user_agent_extra="x-client-framework:haystack", **(self.boto3_config if self.boto3_config else {})
+        )
+
+    def warm_up(self) -> None:
+        """Warm up the component by initializing the settings and storage."""
+        if self._storage is None:
+            self.file_root_path.mkdir(parents=True, exist_ok=True)
+            self._storage = S3Storage.from_env(session=self._session, config=self._config)
+
+    @component.output_types(documents=List[Document])
+    def run(
+        self,
+        documents: List[Document],
+    ) -> Dict[str, List[Document]]:
+        """Download files from AWS S3 Buckets to local filesystem.
+
+        Return enriched `Document`s with the path of the downloaded file.
+        :param documents: Document containing the name of the file to download in the meta field.
+        :returns: A dictionary with:
+            - `documents`: The downloaded `Document`s; each has `meta['file_path']`.
+        :raises S3Error: If a download attempt fails or the file does not exist in the S3 bucket.
+        :raises ValueError: If the path where files will be downloaded is not set.
+        """
+
+        if self._storage is None:
+            msg = f"The component {self.__class__.__name__} was not warmed up. Call 'warm_up()' before calling run()."
+            raise RuntimeError(msg)
+
+        filtered_documents = self._filter_documents_by_extensions(documents) if self.file_extensions else documents
+
+        if not filtered_documents:
+            return {"documents": []}
+
+        try:
+            max_workers = min(self.max_workers, len(filtered_documents) if filtered_documents else self.max_workers)
+            with ThreadPoolExecutor(max_workers=max_workers) as executor:
+                iterable = executor.map(self._download_file, filtered_documents)
+        finally:
+            self._cleanup_cache(filtered_documents)
+
+        downloaded_documents = [d for d in iterable if d is not None]
+        return {"documents": downloaded_documents}
+
+    def _filter_documents_by_extensions(self, documents: List[Document]) -> List[Document]:
+        """Filter documents by file extensions."""
+        if not self.file_extensions:
+            return documents
+        return [
+            doc
+            for doc in documents
+            if Path(doc.meta.get(self.file_name_meta_key, "")).suffix.lower() in self.file_extensions
+        ]
+
+    def _download_file(self, document: Document) -> Optional[Document]:
+        """
+        Download a single file from AWS S3 Bucket to local filesystem.
+
+        :param document: `Document` with the name of the file to download in the meta field.
+        :returns:
+            The same `Document` with `meta` containing the `file_path` of the
+            downloaded file.
+        :raises S3Error: If the download or head request fails or the file does not exist in the S3 bucket.
+        """
+
+        file_name = document.meta.get(self.file_name_meta_key)
+        if not file_name:
+            logger.warning(
+                f"Document missing required file name metadata key '{self.file_name_meta_key}'. Skipping download."
+            )
+            return None
+
+        file_path = self.file_root_path / Path(file_name)
+
+        if file_path.is_file():
+            # set access and modification time to now without redownloading the file
+            file_path.touch()
+
+        else:
+            # we know that _storage is not None after warm_up() is called, but mypy does not know that
+            self._storage.download(key=file_name, local_file_path=file_path)  # type: ignore[union-attr]
+
+        document.meta["file_path"] = str(file_path)
+        return document
+
+    def _cleanup_cache(self, documents: List[Document]) -> None:
+        """
+        Remove least-recently-accessed cache files when cache exceeds `max_cache_size`.
+
+        :param documents: List of Document objects being used containing `cache_id` metadata.
+        """
+        requested_ids = {
+            str(abs(hash(str(doc.meta.get("cache_id", ""))))) for doc in documents if doc.meta.get("cache_id")
+        }
+
+        all_files = [p for p in self.file_root_path.iterdir() if p.is_file()]
+        misses = [p for p in all_files if p.stem not in requested_ids]
+
+        overflow = len(misses) + len(requested_ids) - self.max_cache_size
+        if overflow > 0:
+            misses.sort(key=lambda p: p.stat().st_atime)
+            for p in misses[:overflow]:
+                try:
+                    p.unlink()
+                except Exception as error:
+                    logger.warning("Failed to remove cache file at {path} with error: {e}", path=p, e=error)
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Serialize the component to a dictionary."""
+        return default_to_dict(
+            self,
+            aws_access_key_id=self.aws_access_key_id.to_dict() if self.aws_access_key_id else None,
+            aws_secret_access_key=self.aws_secret_access_key.to_dict() if self.aws_secret_access_key else None,
+            aws_session_token=self.aws_session_token.to_dict() if self.aws_session_token else None,
+            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,
+            file_root_path=str(self.file_root_path),
+            max_workers=self.max_workers,
+            max_cache_size=self.max_cache_size,
+            file_extensions=self.file_extensions,
+            file_name_meta_key=self.file_name_meta_key,
+        )
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "S3Downloader":
+        """
+        Deserializes the component from a dictionary.
+        :param data:
+            Dictionary to deserialize from.
+        :returns:
+            Deserialized component.
+        """
+        deserialize_secrets_inplace(
+            data["init_parameters"],
+            ["aws_access_key_id", "aws_secret_access_key", "aws_session_token", "aws_region_name", "aws_profile_name"],
+        )
+        return default_from_dict(cls, data)

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)