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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: amazon-bedrock-haystack
-Version: 5.2.0
+Version: 5.3.1
 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

{amazon_bedrock_haystack-5.2.0.dist-info → amazon_bedrock_haystack-5.3.1.dist-info}/RECORD

@@ -7,23 +7,23 @@ haystack_integrations/common/s3/errors.py,sha256=BrTDLdhQvAuQutyg35cFyP5h8PNkDEi
 haystack_integrations/common/s3/utils.py,sha256=OJupFj54aQmg6S8VuVq6Lc2qpFZyyJajRVIpwe3_6iA,4744
 haystack_integrations/components/downloaders/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 haystack_integrations/components/downloaders/s3/__init__.py,sha256=2BOd3_N0kGqRJGH-ENrTJqOqzqHryRYaSuNqpLYKMFo,179
-haystack_integrations/components/downloaders/s3/s3_downloader.py,sha256=qarIeGxwDCA1BOZ1qdLfE8NcQtMS9bW54a8voEBHTbE,12637
+haystack_integrations/components/downloaders/s3/s3_downloader.py,sha256=jj9VMQ5QciTZvJ8Qki_zha0tcK0-APm2rMhMPObtr1M,12625
 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
-haystack_integrations/components/embedders/amazon_bedrock/document_image_embedder.py,sha256=CHNH0Dt7JQqYNbZi1lKsGvarnEhJn3UNGdghF0IhqWw,16163
-haystack_integrations/components/embedders/amazon_bedrock/text_embedder.py,sha256=3eSqt3XpH2thblTeOPf-ej1V2UbdG2z50d3jInq1bYc,9144
+haystack_integrations/components/embedders/amazon_bedrock/document_embedder.py,sha256=EKLiKzN0al7OLvfIAnXXNgUTI4TbSc-qCkfjjDJeCd4,13216
+haystack_integrations/components/embedders/amazon_bedrock/document_image_embedder.py,sha256=qKsWRGkubFMHDeqTbPgA5qvEag_CDNTnc5kEmDbc6IQ,16202
+haystack_integrations/components/embedders/amazon_bedrock/text_embedder.py,sha256=rYmU39q_-9uKANBebcQsPtGE0EjrAkOQSFHL3e6WYag,9349
 haystack_integrations/components/generators/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 haystack_integrations/components/generators/amazon_bedrock/__init__.py,sha256=lv4NouIVm78YavUssWQrHHP_81u-7j21qW8v1kZMJPQ,284
-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/adapters.py,sha256=SDCi_EA-1OUmTKrnbcql9oQQuLnoAGzF_XwrKnQ1_jk,19628
+haystack_integrations/components/generators/amazon_bedrock/generator.py,sha256=1vuA0x6k8UOvPjeKjYvcYORpsTFXs1rKMUXlima51IY,14725
 haystack_integrations/components/generators/amazon_bedrock/chat/__init__.py,sha256=6GZ8Y3Lw0rLOsOAqi6Tu5mZC977UzQvgDxKpOWr8IQw,110
-haystack_integrations/components/generators/amazon_bedrock/chat/chat_generator.py,sha256=fwjNUrQM4UziGBVZKvsVEZVcKPF00ZY7CY0g2Advwxs,26246
-haystack_integrations/components/generators/amazon_bedrock/chat/utils.py,sha256=1M_k8CG2WH23Yz-sB7a1kiIqVh2QB8Pqi0zbWXyMUL8,27255
+haystack_integrations/components/generators/amazon_bedrock/chat/chat_generator.py,sha256=batM98HDpUyZEwG-pGYhgTQ6Z9uoGE2TUb3EaZjmox4,26227
+haystack_integrations/components/generators/amazon_bedrock/chat/utils.py,sha256=sWgh-58YeDIatDFNLX2YG1KPJghsyXK5p5uVM4pe8U4,27236
 haystack_integrations/components/rankers/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 haystack_integrations/components/rankers/amazon_bedrock/__init__.py,sha256=mJQKShAP5AfZvfKQisSh7kfKu6RIXzsYdk4eqMtcaEk,75
-haystack_integrations/components/rankers/amazon_bedrock/ranker.py,sha256=QWtUKfJxMrlfLCWTb8cCP-lKEthnEBwnTd1NSbiFMkg,11812
-amazon_bedrock_haystack-5.2.0.dist-info/METADATA,sha256=DO5DjQps1pXcLOiyLmzvoRnf4xgMKE0kYSSFPiFfQeo,2228
-amazon_bedrock_haystack-5.2.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-amazon_bedrock_haystack-5.2.0.dist-info/licenses/LICENSE.txt,sha256=B05uMshqTA74s-0ltyHKI6yoPfJ3zYgQbvcXfDVGFf8,10280
-amazon_bedrock_haystack-5.2.0.dist-info/RECORD,,
+haystack_integrations/components/rankers/amazon_bedrock/ranker.py,sha256=oFNTF2BnKNA1VERiF42ttTFDzxOiaBIA1H5ozLNFFFU,11800
+amazon_bedrock_haystack-5.3.1.dist-info/METADATA,sha256=-icmdgZ9-poMEUHeBb470jU0VnpZwAzlHOcaUS57bLA,2228
+amazon_bedrock_haystack-5.3.1.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+amazon_bedrock_haystack-5.3.1.dist-info/licenses/LICENSE.txt,sha256=B05uMshqTA74s-0ltyHKI6yoPfJ3zYgQbvcXfDVGFf8,10280
+amazon_bedrock_haystack-5.3.1.dist-info/RECORD,,

{amazon_bedrock_haystack-5.2.0.dist-info → amazon_bedrock_haystack-5.3.1.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: hatchling 1.27.0
+Generator: hatchling 1.28.0
 Root-Is-Purelib: true
 Tag: py3-none-any

haystack_integrations/components/downloaders/s3/s3_downloader.py

@@ -5,7 +5,7 @@
 import os
 from concurrent.futures import ThreadPoolExecutor
 from pathlib import Path
-from typing import Any, Callable, Dict, List, Optional
+from typing import Any, Callable, Optional
 
 from botocore.config import Config
 from haystack import component, default_from_dict, default_to_dict, logging
@@ -36,9 +36,9 @@ class S3Downloader:
         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,
+        boto3_config: Optional[dict[str, Any]] = None,
         file_root_path: Optional[str] = None,
-        file_extensions: Optional[List[str]] = None,
+        file_extensions: Optional[list[str]] = None,
         file_name_meta_key: str = "file_name",
         max_workers: int = 32,
         max_cache_size: int = 100,
@@ -126,11 +126,11 @@ class S3Downloader:
             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])
+    @component.output_types(documents=list[Document])
     def run(
         self,
-        documents: List[Document],
-    ) -> Dict[str, List[Document]]:
+        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.
@@ -160,7 +160,7 @@ class S3Downloader:
         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]:
+    def _filter_documents_by_extensions(self, documents: list[Document]) -> list[Document]:
         """Filter documents by file extensions."""
         if not self.file_extensions:
             return documents
@@ -202,7 +202,7 @@ class S3Downloader:
         document.meta["file_path"] = str(file_path)
         return document
 
-    def _cleanup_cache(self, documents: List[Document]) -> None:
+    def _cleanup_cache(self, documents: list[Document]) -> None:
         """
         Remove least-recently-accessed cache files when cache exceeds `max_cache_size`.
 
@@ -224,7 +224,7 @@ class S3Downloader:
                 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]:
+    def to_dict(self) -> dict[str, Any]:
         """Serialize the component to a dictionary."""
 
         s3_key_generation_function_name = (
@@ -247,7 +247,7 @@ class S3Downloader:
         )
 
     @classmethod
-    def from_dict(cls, data: Dict[str, Any]) -> "S3Downloader":
+    def from_dict(cls, data: dict[str, Any]) -> "S3Downloader":
         """
         Deserializes the component from a dictionary.
         :param data:

haystack_integrations/components/embedders/amazon_bedrock/document_embedder.py

@@ -1,5 +1,5 @@
 import json
-from typing import Any, Dict, List, Literal, Optional
+from typing import Any, Optional
 
 from botocore.config import Config
 from botocore.exceptions import ClientError
@@ -16,14 +16,6 @@ from haystack_integrations.common.amazon_bedrock.utils import get_aws_session
 
 logger = logging.getLogger(__name__)
 
-SUPPORTED_EMBEDDING_MODELS = [
-    "amazon.titan-embed-text-v1",
-    "cohere.embed-english-v3",
-    "cohere.embed-multilingual-v3",
-    "amazon.titan-embed-text-v2:0",
-    "amazon.titan-embed-image-v1",
-]
-
 
 @component
 class AmazonBedrockDocumentEmbedder:
@@ -57,13 +49,7 @@ class AmazonBedrockDocumentEmbedder:
 
     def __init__(
         self,
-        model: Literal[
-            "amazon.titan-embed-text-v1",
-            "cohere.embed-english-v3",
-            "cohere.embed-multilingual-v3",
-            "amazon.titan-embed-text-v2:0",
-            "amazon.titan-embed-image-v1",
-        ],
+        model: str,
         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
@@ -73,9 +59,9 @@ class AmazonBedrockDocumentEmbedder:
         aws_profile_name: Optional[Secret] = Secret.from_env_var("AWS_PROFILE", strict=False),  # noqa: B008
         batch_size: int = 32,
         progress_bar: bool = True,
-        meta_fields_to_embed: Optional[List[str]] = None,
+        meta_fields_to_embed: Optional[list[str]] = None,
         embedding_separator: str = "\n",
-        boto3_config: Optional[Dict[str, Any]] = None,
+        boto3_config: Optional[dict[str, Any]] = None,
         **kwargs: Any,
     ) -> None:
         """
@@ -88,8 +74,13 @@ class AmazonBedrockDocumentEmbedder:
         constructor. Aside from model, three required parameters are `aws_access_key_id`, `aws_secret_access_key`,
          and `aws_region_name`.
 
-        :param model: The embedding model to use. The model has to be specified in the format outlined in the Amazon
-            Bedrock [documentation](https://docs.aws.amazon.com/bedrock/latest/userguide/model-ids.html).
+        :param model: The embedding model to use.
+            Amazon Titan and Cohere embedding models are supported, for example:
+            "amazon.titan-embed-text-v1", "amazon.titan-embed-text-v2:0", "amazon.titan-embed-image-v1",
+            "cohere.embed-english-v3", "cohere.embed-multilingual-v3", "cohere.embed-v4:0".
+            To find all supported models, refer to the Amazon Bedrock
+            [documentation](https://docs.aws.amazon.com/bedrock/latest/userguide/models-supported.html) and
+            filter for "embedding", then select models from the Amazon Titan and Cohere series.
         :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.
@@ -107,11 +98,8 @@ class AmazonBedrockDocumentEmbedder:
         :raises ValueError: If the model is not supported.
         :raises AmazonBedrockConfigurationError: If the AWS environment is not configured correctly.
         """
-
-        if not model or model not in SUPPORTED_EMBEDDING_MODELS:
-            msg = "Please provide a valid model from the list of supported models: " + ", ".join(
-                SUPPORTED_EMBEDDING_MODELS
-            )
+        if "titan" not in model and "cohere" not in model:
+            msg = f"Model {model} is not supported. Only Amazon Titan and Cohere embedding models are supported."
             raise ValueError(msg)
 
         self.model = model
@@ -149,7 +137,7 @@ class AmazonBedrockDocumentEmbedder:
             )
             raise AmazonBedrockConfigurationError(msg) from exception
 
-    def _prepare_texts_to_embed(self, documents: List[Document]) -> List[str]:
+    def _prepare_texts_to_embed(self, documents: list[Document]) -> list[str]:
         """
         Prepare the texts to embed by concatenating the Document text with the metadata fields to embed.
         """
@@ -162,7 +150,7 @@ class AmazonBedrockDocumentEmbedder:
             texts_to_embed.append(text_to_embed)
         return texts_to_embed
 
-    def _embed_cohere(self, documents: List[Document]) -> List[Document]:
+    def _embed_cohere(self, documents: list[Document]) -> list[Document]:
         """
         Internal method to embed Documents using Cohere models.
         Batch inference is supported.
@@ -191,15 +179,19 @@ class AmazonBedrockDocumentEmbedder:
                 msg = f"Could not perform inference for Amazon Bedrock model {self.model} due to:\n{exception}"
                 raise AmazonBedrockInferenceError(msg) from exception
 
-            response_body = json.loads(response.get("body").read())
-            all_embeddings.extend(response_body["embeddings"])
+            cohere_embeddings = json.loads(response.get("body").read())["embeddings"]
+            # depending on the model, Cohere returns a dict with the embedding types as keys or a list of lists
+            embeddings_list = (
+                next(iter(cohere_embeddings.values())) if isinstance(cohere_embeddings, dict) else cohere_embeddings
+            )
+            all_embeddings.extend(embeddings_list)
 
         for doc, emb in zip(documents, all_embeddings):
             doc.embedding = emb
 
         return documents
 
-    def _embed_titan(self, documents: List[Document]) -> List[Document]:
+    def _embed_titan(self, documents: list[Document]) -> list[Document]:
         """
         Internal method to embed Documents using Amazon Titan models.
         NOTE: Batch inference is not supported, so embeddings are created one by one.
@@ -227,8 +219,8 @@ class AmazonBedrockDocumentEmbedder:
 
         return documents
 
-    @component.output_types(documents=List[Document])
-    def run(self, documents: List[Document]) -> Dict[str, List[Document]]:
+    @component.output_types(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.
@@ -248,12 +240,12 @@ class AmazonBedrockDocumentEmbedder:
         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)}."
+            msg = f"Model {self.model} is not supported. Only Amazon Titan and Cohere embedding models are supported."
             raise ValueError(msg)
 
         return {"documents": documents_with_embeddings}
 
-    def to_dict(self) -> Dict[str, Any]:
+    def to_dict(self) -> dict[str, Any]:
         """
         Serializes the component to a dictionary.
 
@@ -277,7 +269,7 @@ class AmazonBedrockDocumentEmbedder:
         )
 
     @classmethod
-    def from_dict(cls, data: Dict[str, Any]) -> "AmazonBedrockDocumentEmbedder":
+    def from_dict(cls, data: dict[str, Any]) -> "AmazonBedrockDocumentEmbedder":
         """
         Deserializes the component from a dictionary.
 

haystack_integrations/components/embedders/amazon_bedrock/document_image_embedder.py

@@ -4,7 +4,7 @@
 
 import json
 from dataclasses import replace
-from typing import Any, Dict, List, Literal, Optional, Tuple
+from typing import Any, Optional
 
 from botocore.config import Config
 from botocore.exceptions import ClientError
@@ -27,8 +27,6 @@ from haystack_integrations.common.amazon_bedrock.utils import get_aws_session
 
 logger = logging.getLogger(__name__)
 
-SUPPORTED_EMBEDDING_MODELS = ["amazon.titan-embed-image-v1", "cohere.embed-english-v3", "cohere.embed-multilingual-v3"]
-
 
 @component
 class AmazonBedrockDocumentImageEmbedder:
@@ -69,7 +67,7 @@ class AmazonBedrockDocumentImageEmbedder:
     def __init__(
         self,
         *,
-        model: Literal["amazon.titan-embed-image-v1", "cohere.embed-english-v3", "cohere.embed-multilingual-v3"],
+        model: str,
         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
@@ -79,20 +77,21 @@ class AmazonBedrockDocumentImageEmbedder:
         aws_profile_name: Optional[Secret] = Secret.from_env_var("AWS_PROFILE", strict=False),  # noqa: B008
         file_path_meta_field: str = "file_path",
         root_path: Optional[str] = None,
-        image_size: Optional[Tuple[int, int]] = None,
+        image_size: Optional[tuple[int, int]] = None,
         progress_bar: bool = True,
-        boto3_config: Optional[Dict[str, Any]] = None,
+        boto3_config: Optional[dict[str, Any]] = None,
         **kwargs: Any,
     ) -> None:
         """
         Creates a AmazonBedrockDocumentImageEmbedder component.
 
-        :param model:
-            The Bedrock model to use for calculating embeddings. Pass a valid model ID.
-            Supported models:
-            - "amazon.titan-embed-image-v1"
-            - "cohere.embed-english-v3"
-            - "cohere.embed-multilingual-v3"
+        :param model: The embedding model to use.
+            Amazon Titan and Cohere multimodal embedding models are supported, for example:
+            "amazon.titan-embed-image-v1", "cohere.embed-english-v3", "cohere.embed-multilingual-v3",
+            "cohere.embed-v4:0".
+            To find all supported models, refer to the Amazon Bedrock
+            [documentation](https://docs.aws.amazon.com/bedrock/latest/userguide/models-supported.html) and
+            filter for "embedding", then select multimodal models from the Amazon Titan and Cohere series.
         :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.
@@ -114,9 +113,10 @@ class AmazonBedrockDocumentImageEmbedder:
         :raises ValueError: If the model is not supported.
         :raises AmazonBedrockConfigurationError: If the AWS environment is not configured correctly.
         """
-        if not model or model not in SUPPORTED_EMBEDDING_MODELS:
-            msg = "Please provide a valid model from the list of supported models: " + ", ".join(
-                SUPPORTED_EMBEDDING_MODELS
+        if "titan" not in model and "cohere" not in model:
+            msg = (
+                f"Model {model} is not supported. "
+                "Only Amazon Titan and Cohere multimodal embedding models are supported."
             )
             raise ValueError(msg)
 
@@ -135,14 +135,14 @@ class AmazonBedrockDocumentImageEmbedder:
         self.kwargs = kwargs
         self.embedding_types = None
 
-        if emmbedding_types := self.kwargs.get("embedding_types"):
-            if len(emmbedding_types) > 1:
+        if embedding_types := self.kwargs.get("embedding_types"):
+            if len(embedding_types) > 1:
                 msg = (
                     "You have provided multiple embedding_types for Cohere model. "
                     "AmazonBedrockDocumentImageEmbedder only supports one embedding_type at a time."
                 )
                 raise ValueError(msg)
-            self.embedding_types = emmbedding_types
+            self.embedding_types = embedding_types
 
         def resolve_secret(secret: Optional[Secret]) -> Optional[str]:
             return secret.resolve_value() if secret else None
@@ -280,7 +280,10 @@ class AmazonBedrockDocumentImageEmbedder:
         elif "titan" in self.model:
             embeddings = self._embed_titan(images=images_to_embed)
         else:
-            msg = f"Model {self.model} is not supported. Supported models are: {', '.join(SUPPORTED_EMBEDDING_MODELS)}."
+            msg = (
+                f"Model {self.model} is not supported. "
+                "Only Amazon Titan and Cohere multimodal embedding models are supported."
+            )
             raise ValueError(msg)
 
         docs_with_embeddings = []
@@ -296,7 +299,7 @@ class AmazonBedrockDocumentImageEmbedder:
 
         return {"documents": docs_with_embeddings}
 
-    def _embed_titan(self, images: List[str]) -> List[List[float]]:
+    def _embed_titan(self, images: list[str]) -> list[list[float]]:
         """
         Internal method to embed base64 images using Amazon Titan models.
 
@@ -326,7 +329,7 @@ class AmazonBedrockDocumentImageEmbedder:
 
         return all_embeddings
 
-    def _embed_cohere(self, image_uris: List[str]) -> List[List[float]]:
+    def _embed_cohere(self, image_uris: list[str]) -> list[list[float]]:
         """
         Internal method to embed base64 images using Cohere models.
 
@@ -351,15 +354,13 @@ class AmazonBedrockDocumentImageEmbedder:
                 raise AmazonBedrockInferenceError(msg) from exception
 
             response_body = json.loads(response.get("body").read())
-            embeddings = response_body["embeddings"]
+            cohere_embeddings = response_body["embeddings"]
 
-            # if embedding_types is specified, cohere returns a dict with the embedding types as keys
-            if isinstance(embeddings, dict):
-                for embedding in embeddings.values():
-                    all_embeddings.append(embedding[0])
-            else:
-                # if embedding_types is not specified, cohere returns
-                # a nested list of float embeddings
-                all_embeddings.append(embeddings[0])
+            # depending on the model and embedding_types, Cohere returns a dict with the embedding types as keys
+            # or a list of lists
+            embeddings_list = (
+                next(iter(cohere_embeddings.values())) if isinstance(cohere_embeddings, dict) else cohere_embeddings
+            )
+            all_embeddings.extend(embeddings_list)
 
         return all_embeddings

haystack_integrations/components/embedders/amazon_bedrock/text_embedder.py

@@ -1,5 +1,5 @@
 import json
-from typing import Any, Dict, List, Literal, Optional
+from typing import Any, Optional
 
 from botocore.config import Config
 from botocore.exceptions import ClientError
@@ -14,14 +14,6 @@ from haystack_integrations.common.amazon_bedrock.utils import get_aws_session
 
 logger = logging.getLogger(__name__)
 
-SUPPORTED_EMBEDDING_MODELS = [
-    "amazon.titan-embed-text-v1",
-    "cohere.embed-english-v3",
-    "cohere.embed-multilingual-v3",
-    "amazon.titan-embed-text-v2:0",
-    "amazon.titan-embed-image-v1",
-]
-
 
 @component
 class AmazonBedrockTextEmbedder:
@@ -50,13 +42,7 @@ class AmazonBedrockTextEmbedder:
 
     def __init__(
         self,
-        model: Literal[
-            "amazon.titan-embed-text-v1",
-            "cohere.embed-english-v3",
-            "cohere.embed-multilingual-v3",
-            "amazon.titan-embed-text-v2:0",
-            "amazon.titan-embed-image-v1",
-        ],
+        model: str,
         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
@@ -64,7 +50,7 @@ class AmazonBedrockTextEmbedder:
         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,
+        boto3_config: Optional[dict[str, Any]] = None,
         **kwargs: Any,
     ) -> None:
         """
@@ -77,8 +63,13 @@ class AmazonBedrockTextEmbedder:
         constructor. Aside from model, three required parameters are `aws_access_key_id`, `aws_secret_access_key`,
          and `aws_region_name`.
 
-        :param model: The embedding model to use. The model has to be specified in the format outlined in the Amazon
-            Bedrock [documentation](https://docs.aws.amazon.com/bedrock/latest/userguide/model-ids.html).
+        :param model: The embedding model to use.
+            Amazon Titan and Cohere embedding models are supported, for example:
+            "amazon.titan-embed-text-v1", "amazon.titan-embed-text-v2:0", "amazon.titan-embed-image-v1",
+            "cohere.embed-english-v3", "cohere.embed-multilingual-v3", "cohere.embed-v4:0".
+            To find all supported models, refer to the Amazon Bedrock
+            [documentation](https://docs.aws.amazon.com/bedrock/latest/userguide/models-supported.html) and
+            filter for "embedding", then select models from the Amazon Titan and Cohere series.
         :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.
@@ -90,10 +81,8 @@ class AmazonBedrockTextEmbedder:
         :raises ValueError: If the model is not supported.
         :raises AmazonBedrockConfigurationError: If the AWS environment is not configured correctly.
         """
-        if not model or model not in SUPPORTED_EMBEDDING_MODELS:
-            msg = "Please provide a valid model from the list of supported models: " + ", ".join(
-                SUPPORTED_EMBEDDING_MODELS
-            )
+        if "titan" not in model and "cohere" not in model:
+            msg = f"Model {model} is not supported. Only Amazon Titan and Cohere embedding models are supported."
             raise ValueError(msg)
 
         self.model = model
@@ -127,8 +116,8 @@ class AmazonBedrockTextEmbedder:
             )
             raise AmazonBedrockConfigurationError(msg) from exception
 
-    @component.output_types(embedding=List[float])
-    def run(self, text: str) -> Dict[str, List[float]]:
+    @component.output_types(embedding=list[float])
+    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.
@@ -168,16 +157,21 @@ class AmazonBedrockTextEmbedder:
         response_body = json.loads(response.get("body").read())
 
         if "cohere" in self.model:
-            embedding = response_body["embeddings"][0]
+            cohere_embeddings = response_body["embeddings"]
+            # depending on the model, Cohere returns a dict with the embedding types as keys or a list of lists
+            embeddings_list = (
+                next(iter(cohere_embeddings.values())) if isinstance(cohere_embeddings, dict) else cohere_embeddings
+            )
+            embedding = embeddings_list[0]
         elif "titan" in self.model:
             embedding = response_body["embedding"]
         else:
-            msg = f"Unsupported model {self.model}. Supported models are: {', '.join(SUPPORTED_EMBEDDING_MODELS)}"
+            msg = f"Model {self.model} is not supported. Only Amazon Titan and Cohere embedding models are supported."
             raise ValueError(msg)
 
         return {"embedding": embedding}
 
-    def to_dict(self) -> Dict[str, Any]:
+    def to_dict(self) -> dict[str, Any]:
         """
         Serializes the component to a dictionary.
 
@@ -197,7 +191,7 @@ class AmazonBedrockTextEmbedder:
         )
 
     @classmethod
-    def from_dict(cls, data: Dict[str, Any]) -> "AmazonBedrockTextEmbedder":
+    def from_dict(cls, data: dict[str, Any]) -> "AmazonBedrockTextEmbedder":
         """
         Deserializes the component from a dictionary.
 

haystack_integrations/components/generators/amazon_bedrock/adapters.py

@@ -1,6 +1,6 @@
 import json
 from abc import ABC, abstractmethod
-from typing import Any, Dict, List, Optional
+from typing import Any, Optional
 
 from botocore.eventstream import EventStream
 from haystack.dataclasses import StreamingChunk, SyncStreamingCallbackT
@@ -19,12 +19,12 @@ class BedrockModelAdapter(ABC):
         It will be overridden by the corresponding parameter in the `model_kwargs` if it is present.
     """
 
-    def __init__(self, model_kwargs: Dict[str, Any], max_length: Optional[int]) -> None:
+    def __init__(self, model_kwargs: dict[str, Any], max_length: Optional[int]) -> None:
         self.model_kwargs = model_kwargs
         self.max_length = max_length
 
     @abstractmethod
-    def prepare_body(self, prompt: str, **inference_kwargs: Any) -> Dict[str, Any]:
+    def prepare_body(self, prompt: str, **inference_kwargs: Any) -> dict[str, Any]:
         """
         Prepares the body for the Amazon Bedrock request.
         Each subclass should implement this method to prepare the request body for the specific model.
@@ -34,7 +34,7 @@ class BedrockModelAdapter(ABC):
         :returns: A dictionary containing the body for the request.
         """
 
-    def get_responses(self, response_body: Dict[str, Any]) -> List[str]:
+    def get_responses(self, response_body: dict[str, Any]) -> list[str]:
         """
         Extracts the responses from the Amazon Bedrock response.
 
@@ -45,7 +45,7 @@ class BedrockModelAdapter(ABC):
         responses = [completion.lstrip() for completion in completions]
         return responses
 
-    def get_stream_responses(self, stream: EventStream, streaming_callback: SyncStreamingCallbackT) -> List[str]:
+    def get_stream_responses(self, stream: EventStream, streaming_callback: SyncStreamingCallbackT) -> list[str]:
         """
         Extracts the responses from the Amazon Bedrock streaming response.
 
@@ -53,7 +53,7 @@ class BedrockModelAdapter(ABC):
         :param streaming_callback: The handler for the streaming response.
         :returns: A list of string responses.
         """
-        streaming_chunks: List[StreamingChunk] = []
+        streaming_chunks: list[StreamingChunk] = []
         for event in stream:
             chunk = event.get("chunk")
             if chunk:
@@ -65,7 +65,7 @@ class BedrockModelAdapter(ABC):
         responses = ["".join(streaming_chunk.content for streaming_chunk in streaming_chunks).lstrip()]
         return responses
 
-    def _get_params(self, inference_kwargs: Dict[str, Any], default_params: Dict[str, Any]) -> Dict[str, Any]:
+    def _get_params(self, inference_kwargs: dict[str, Any], default_params: dict[str, Any]) -> dict[str, Any]:
         """
         Merges the default params with the inference kwargs and model kwargs.
 
@@ -83,7 +83,7 @@ class BedrockModelAdapter(ABC):
         }
 
     @abstractmethod
-    def _extract_completions_from_response(self, response_body: Dict[str, Any]) -> List[str]:
+    def _extract_completions_from_response(self, response_body: dict[str, Any]) -> list[str]:
         """
         Extracts the responses from the Amazon Bedrock response.
 
@@ -92,7 +92,7 @@ class BedrockModelAdapter(ABC):
         """
 
     @abstractmethod
-    def _build_streaming_chunk(self, chunk: Dict[str, Any]) -> StreamingChunk:
+    def _build_streaming_chunk(self, chunk: dict[str, Any]) -> StreamingChunk:
         """
         Extracts the content and meta from a streaming chunk.
 
@@ -115,7 +115,7 @@ class AnthropicClaudeAdapter(BedrockModelAdapter):
     :param max_length: Maximum length of generated text
     """
 
-    def __init__(self, model_kwargs: Dict[str, Any], max_length: Optional[int]) -> None:
+    def __init__(self, model_kwargs: dict[str, Any], max_length: Optional[int]) -> None:
         self.use_messages_api = model_kwargs.get("use_messages_api", True)
         self.include_thinking = model_kwargs.get("include_thinking", True)
         self.thinking_tag = model_kwargs.get("thinking_tag", "thinking")
@@ -123,7 +123,7 @@ class AnthropicClaudeAdapter(BedrockModelAdapter):
         self.thinking_tag_end = f"</{self.thinking_tag}>\n\n" if self.thinking_tag else "\n\n"
         super().__init__(model_kwargs, max_length)
 
-    def prepare_body(self, prompt: str, **inference_kwargs: Any) -> Dict[str, Any]:
+    def prepare_body(self, prompt: str, **inference_kwargs: Any) -> dict[str, Any]:
         """
         Prepares the body for the Claude model
 
@@ -134,7 +134,7 @@ class AnthropicClaudeAdapter(BedrockModelAdapter):
             - specified inference parameters.
         """
         if self.use_messages_api:
-            default_params: Dict[str, Any] = {
+            default_params: dict[str, Any] = {
                 "anthropic_version": "bedrock-2023-05-31",
                 "max_tokens": self.max_length,
                 "system": None,
@@ -160,7 +160,7 @@ class AnthropicClaudeAdapter(BedrockModelAdapter):
             body = {"prompt": f"\n\nHuman: {prompt}\n\nAssistant:", **params}
         return body
 
-    def _extract_completions_from_response(self, response_body: Dict[str, Any]) -> List[str]:
+    def _extract_completions_from_response(self, response_body: dict[str, Any]) -> list[str]:
         """
         Extracts the responses from the Amazon Bedrock response.
 
@@ -181,7 +181,7 @@ class AnthropicClaudeAdapter(BedrockModelAdapter):
 
         return [response_body["completion"]]
 
-    def _build_streaming_chunk(self, chunk: Dict[str, Any]) -> StreamingChunk:
+    def _build_streaming_chunk(self, chunk: dict[str, Any]) -> StreamingChunk:
         """
         Extracts the content and meta from a streaming chunk.
 
@@ -211,7 +211,7 @@ class MistralAdapter(BedrockModelAdapter):
     Adapter for the Mistral models.
     """
 
-    def prepare_body(self, prompt: str, **inference_kwargs: Any) -> Dict[str, Any]:
+    def prepare_body(self, prompt: str, **inference_kwargs: Any) -> dict[str, Any]:
         """
         Prepares the body for the Mistral model
 
@@ -221,7 +221,7 @@ class MistralAdapter(BedrockModelAdapter):
             - `prompt`: The prompt to be sent to the model.
             - specified inference parameters.
         """
-        default_params: Dict[str, Any] = {
+        default_params: dict[str, Any] = {
             "max_tokens": self.max_length,
             "stop": [],
             "temperature": None,
@@ -233,7 +233,7 @@ class MistralAdapter(BedrockModelAdapter):
         formatted_prompt = f"<s>[INST] {prompt} [/INST]" if "INST" not in prompt else prompt
         return {"prompt": formatted_prompt, **params}
 
-    def _extract_completions_from_response(self, response_body: Dict[str, Any]) -> List[str]:
+    def _extract_completions_from_response(self, response_body: dict[str, Any]) -> list[str]:
         """
         Extracts the responses from the Amazon Bedrock response.
 
@@ -242,7 +242,7 @@ class MistralAdapter(BedrockModelAdapter):
         """
         return [output.get("text", "") for output in response_body.get("outputs", [])]
 
-    def _build_streaming_chunk(self, chunk: Dict[str, Any]) -> StreamingChunk:
+    def _build_streaming_chunk(self, chunk: dict[str, Any]) -> StreamingChunk:
         """
         Extracts the content and meta from a streaming chunk.
 
@@ -261,7 +261,7 @@ class CohereCommandAdapter(BedrockModelAdapter):
     Adapter for the Cohere Command model.
     """
 
-    def prepare_body(self, prompt: str, **inference_kwargs: Any) -> Dict[str, Any]:
+    def prepare_body(self, prompt: str, **inference_kwargs: Any) -> dict[str, Any]:
         """
         Prepares the body for the Command model
 
@@ -288,7 +288,7 @@ class CohereCommandAdapter(BedrockModelAdapter):
         body = {"prompt": prompt, **params}
         return body
 
-    def _extract_completions_from_response(self, response_body: Dict[str, Any]) -> List[str]:
+    def _extract_completions_from_response(self, response_body: dict[str, Any]) -> list[str]:
         """
         Extracts the responses from the Cohere Command model response.
 
@@ -298,7 +298,7 @@ class CohereCommandAdapter(BedrockModelAdapter):
         responses = [generation["text"] for generation in response_body["generations"]]
         return responses
 
-    def _build_streaming_chunk(self, chunk: Dict[str, Any]) -> StreamingChunk:
+    def _build_streaming_chunk(self, chunk: dict[str, Any]) -> StreamingChunk:
         """
         Extracts the content and meta from a streaming chunk.
 
@@ -313,7 +313,7 @@ class CohereCommandRAdapter(BedrockModelAdapter):
     Adapter for the Cohere Command R models.
     """
 
-    def prepare_body(self, prompt: str, **inference_kwargs: Any) -> Dict[str, Any]:
+    def prepare_body(self, prompt: str, **inference_kwargs: Any) -> dict[str, Any]:
         """
         Prepares the body for the Command model
 
@@ -347,7 +347,7 @@ class CohereCommandRAdapter(BedrockModelAdapter):
         body = {"message": prompt, **params}
         return body
 
-    def _extract_completions_from_response(self, response_body: Dict[str, Any]) -> List[str]:
+    def _extract_completions_from_response(self, response_body: dict[str, Any]) -> list[str]:
         """
         Extracts the responses from the Cohere Command model response.
 
@@ -357,7 +357,7 @@ class CohereCommandRAdapter(BedrockModelAdapter):
         responses = [response_body["text"]]
         return responses
 
-    def _build_streaming_chunk(self, chunk: Dict[str, Any]) -> StreamingChunk:
+    def _build_streaming_chunk(self, chunk: dict[str, Any]) -> StreamingChunk:
         """
         Extracts the content and meta from a streaming chunk.
 
@@ -373,7 +373,7 @@ class AI21LabsJurassic2Adapter(BedrockModelAdapter):
     Model adapter for AI21 Labs' Jurassic 2 models.
     """
 
-    def prepare_body(self, prompt: str, **inference_kwargs: Any) -> Dict[str, Any]:
+    def prepare_body(self, prompt: str, **inference_kwargs: Any) -> dict[str, Any]:
         """Prepares the body for the Jurassic 2 model.
 
         :param prompt: The prompt to be sent to the model.
@@ -397,11 +397,11 @@ class AI21LabsJurassic2Adapter(BedrockModelAdapter):
         body = {"prompt": prompt, **params}
         return body
 
-    def _extract_completions_from_response(self, response_body: Dict[str, Any]) -> List[str]:
+    def _extract_completions_from_response(self, response_body: dict[str, Any]) -> list[str]:
         responses = [completion["data"]["text"] for completion in response_body["completions"]]
         return responses
 
-    def _build_streaming_chunk(self, chunk: Dict[str, Any]) -> StreamingChunk:
+    def _build_streaming_chunk(self, chunk: dict[str, Any]) -> StreamingChunk:
         msg = "Streaming is not supported for AI21 Jurassic 2 models."
         raise NotImplementedError(msg)
 
@@ -411,7 +411,7 @@ class AmazonTitanAdapter(BedrockModelAdapter):
     Adapter for Amazon's Titan models.
     """
 
-    def prepare_body(self, prompt: str, **inference_kwargs: Any) -> Dict[str, Any]:
+    def prepare_body(self, prompt: str, **inference_kwargs: Any) -> dict[str, Any]:
         """
         Prepares the body for the Titan model
 
@@ -432,7 +432,7 @@ class AmazonTitanAdapter(BedrockModelAdapter):
         body = {"inputText": prompt, "textGenerationConfig": params}
         return body
 
-    def _extract_completions_from_response(self, response_body: Dict[str, Any]) -> List[str]:
+    def _extract_completions_from_response(self, response_body: dict[str, Any]) -> list[str]:
         """
         Extracts the responses from the Titan model response.
 
@@ -442,7 +442,7 @@ class AmazonTitanAdapter(BedrockModelAdapter):
         responses = [result["outputText"] for result in response_body["results"]]
         return responses
 
-    def _build_streaming_chunk(self, chunk: Dict[str, Any]) -> StreamingChunk:
+    def _build_streaming_chunk(self, chunk: dict[str, Any]) -> StreamingChunk:
         """
         Extracts the content and meta from a streaming chunk.
 
@@ -457,7 +457,7 @@ class MetaLlamaAdapter(BedrockModelAdapter):
     Adapter for Meta's Llama2 models.
     """
 
-    def prepare_body(self, prompt: str, **inference_kwargs: Any) -> Dict[str, Any]:
+    def prepare_body(self, prompt: str, **inference_kwargs: Any) -> dict[str, Any]:
         """
         Prepares the body for the Llama2 model
 
@@ -477,7 +477,7 @@ class MetaLlamaAdapter(BedrockModelAdapter):
         body = {"prompt": prompt, **params}
         return body
 
-    def _extract_completions_from_response(self, response_body: Dict[str, Any]) -> List[str]:
+    def _extract_completions_from_response(self, response_body: dict[str, Any]) -> list[str]:
         """
         Extracts the responses from the Llama2 model response.
 
@@ -486,7 +486,7 @@ class MetaLlamaAdapter(BedrockModelAdapter):
         """
         return [response_body["generation"]]
 
-    def _build_streaming_chunk(self, chunk: Dict[str, Any]) -> StreamingChunk:
+    def _build_streaming_chunk(self, chunk: dict[str, Any]) -> StreamingChunk:
         """
         Extracts the content and meta from a streaming chunk.
 

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

@@ -1,4 +1,4 @@
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Optional
 
 import aioboto3
 from botocore.config import Config
@@ -154,12 +154,12 @@ class AmazonBedrockChatGenerator:
         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
-        generation_kwargs: Optional[Dict[str, Any]] = None,
+        generation_kwargs: Optional[dict[str, Any]] = None,
         streaming_callback: Optional[StreamingCallbackT] = None,
-        boto3_config: Optional[Dict[str, Any]] = None,
+        boto3_config: Optional[dict[str, Any]] = None,
         tools: Optional[ToolsType] = None,
         *,
-        guardrail_config: Optional[Dict[str, str]] = None,
+        guardrail_config: Optional[dict[str, str]] = None,
     ) -> None:
         """
         Initializes the `AmazonBedrockChatGenerator` with the provided parameters. The parameters are passed to the
@@ -289,7 +289,7 @@ class AmazonBedrockChatGenerator:
             )
             raise AmazonBedrockConfigurationError(msg) from exception
 
-    def to_dict(self) -> Dict[str, Any]:
+    def to_dict(self) -> dict[str, Any]:
         """
         Serializes the component to a dictionary.
 
@@ -313,7 +313,7 @@ class AmazonBedrockChatGenerator:
         )
 
     @classmethod
-    def from_dict(cls, data: Dict[str, Any]) -> "AmazonBedrockChatGenerator":
+    def from_dict(cls, data: dict[str, Any]) -> "AmazonBedrockChatGenerator":
         """
         Deserializes the component from a dictionary.
 
@@ -340,12 +340,12 @@ class AmazonBedrockChatGenerator:
 
     def _prepare_request_params(
         self,
-        messages: List[ChatMessage],
+        messages: list[ChatMessage],
         streaming_callback: Optional[StreamingCallbackT] = None,
-        generation_kwargs: Optional[Dict[str, Any]] = None,
+        generation_kwargs: Optional[dict[str, Any]] = None,
         tools: Optional[ToolsType] = None,
         requires_async: bool = False,
-    ) -> Tuple[Dict[str, Any], Optional[StreamingCallbackT]]:
+    ) -> tuple[dict[str, Any], Optional[StreamingCallbackT]]:
         """
         Prepares and formats parameters required to call the Amazon Bedrock Converse API.
 
@@ -419,14 +419,14 @@ class AmazonBedrockChatGenerator:
 
         return params, callback
 
-    @component.output_types(replies=List[ChatMessage])
+    @component.output_types(replies=list[ChatMessage])
     def run(
         self,
-        messages: List[ChatMessage],
+        messages: list[ChatMessage],
         streaming_callback: Optional[StreamingCallbackT] = None,
-        generation_kwargs: Optional[Dict[str, Any]] = None,
+        generation_kwargs: Optional[dict[str, Any]] = None,
         tools: Optional[ToolsType] = None,
-    ) -> Dict[str, List[ChatMessage]]:
+    ) -> dict[str, list[ChatMessage]]:
         """
         Executes a synchronous inference call to the Amazon Bedrock model using the Converse API.
 
@@ -480,14 +480,14 @@ class AmazonBedrockChatGenerator:
 
         return {"replies": replies}
 
-    @component.output_types(replies=List[ChatMessage])
+    @component.output_types(replies=list[ChatMessage])
     async def run_async(
         self,
-        messages: List[ChatMessage],
+        messages: list[ChatMessage],
         streaming_callback: Optional[StreamingCallbackT] = None,
-        generation_kwargs: Optional[Dict[str, Any]] = None,
+        generation_kwargs: Optional[dict[str, Any]] = None,
         tools: Optional[ToolsType] = None,
-    ) -> Dict[str, List[ChatMessage]]:
+    ) -> dict[str, list[ChatMessage]]:
         """
         Executes an asynchronous inference call to the Amazon Bedrock model using the Converse API.
 

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

@@ -1,7 +1,7 @@
 import base64
 import json
 from datetime import datetime, timezone
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Optional
 
 from botocore.eventstream import EventStream
 from haystack import logging
@@ -29,7 +29,7 @@ logger = logging.getLogger(__name__)
 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] = {
+FINISH_REASON_MAPPING: dict[str, FinishReason] = {
     "end_turn": "stop",
     "stop_sequence": "stop",
     "max_tokens": "length",
@@ -40,7 +40,7 @@ FINISH_REASON_MAPPING: Dict[str, FinishReason] = {
 
 
 # Haystack to Bedrock util methods
-def _format_tools(tools: Optional[List[Tool]] = None) -> Optional[Dict[str, Any]]:
+def _format_tools(tools: Optional[list[Tool]] = None) -> Optional[dict[str, Any]]:
     """
     Format Haystack Tool(s) to Amazon Bedrock toolConfig format.
 
@@ -60,7 +60,7 @@ 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]:
+def _format_tool_call_message(tool_call_message: ChatMessage) -> dict[str, Any]:
     """
     Format a Haystack ChatMessage containing tool calls into Bedrock format.
 
@@ -68,7 +68,7 @@ def _format_tool_call_message(tool_call_message: ChatMessage) -> Dict[str, Any]:
     :returns:
         Dictionary representing the tool call message in Bedrock's expected format
     """
-    content: List[Dict[str, Any]] = []
+    content: list[dict[str, Any]] = []
 
     # tool call messages can contain reasoning content
     if reasoning_content := tool_call_message.reasoning:
@@ -85,7 +85,7 @@ def _format_tool_call_message(tool_call_message: ChatMessage) -> Dict[str, Any]:
     return {"role": tool_call_message.role.value, "content": content}
 
 
-def _format_tool_result_message(tool_call_result_message: ChatMessage) -> Dict[str, Any]:
+def _format_tool_result_message(tool_call_result_message: ChatMessage) -> dict[str, Any]:
     """
     Format a Haystack ChatMessage containing tool call results into Bedrock format.
 
@@ -114,7 +114,7 @@ def _format_tool_result_message(tool_call_result_message: ChatMessage) -> Dict[s
     return {"role": "user", "content": tool_results}
 
 
-def _repair_tool_result_messages(bedrock_formatted_messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+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.
 
@@ -135,7 +135,7 @@ def _repair_tool_result_messages(bedrock_formatted_messages: List[Dict[str, Any]
                 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}
+    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:
@@ -162,7 +162,7 @@ def _repair_tool_result_messages(bedrock_formatted_messages: List[Dict[str, Any]
             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]] = []
+    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]:
@@ -176,7 +176,7 @@ def _repair_tool_result_messages(bedrock_formatted_messages: List[Dict[str, Any]
     return [msg for _, msg in repaired_bedrock_formatted_messages]
 
 
-def _format_reasoning_content(reasoning_content: ReasoningContent) -> List[Dict[str, Any]]:
+def _format_reasoning_content(reasoning_content: ReasoningContent) -> list[dict[str, Any]]:
     """
     Format ReasoningContent to match Bedrock's expected structure.
 
@@ -194,7 +194,7 @@ def _format_reasoning_content(reasoning_content: ReasoningContent) -> List[Dict[
     return formatted_contents
 
 
-def _format_text_image_message(message: ChatMessage) -> Dict[str, Any]:
+def _format_text_image_message(message: ChatMessage) -> dict[str, Any]:
     """
     Format a Haystack ChatMessage containing text and optional image content into Bedrock format.
 
@@ -204,7 +204,7 @@ def _format_text_image_message(message: ChatMessage) -> Dict[str, Any]:
     """
     content_parts = message._content
 
-    bedrock_content_blocks: List[Dict[str, Any]] = []
+    bedrock_content_blocks: list[dict[str, Any]] = []
     # Add reasoning content if available as the first content block
     if message.reasoning:
         bedrock_content_blocks.extend(_format_reasoning_content(reasoning_content=message.reasoning))
@@ -231,7 +231,7 @@ def _format_text_image_message(message: ChatMessage) -> Dict[str, Any]:
     return {"role": message.role.value, "content": bedrock_content_blocks}
 
 
-def _format_messages(messages: List[ChatMessage]) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]:
+def _format_messages(messages: list[ChatMessage]) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
     """
     Format a list of Haystack ChatMessages to the format expected by Bedrock API.
 
@@ -262,7 +262,7 @@ def _format_messages(messages: List[ChatMessage]) -> Tuple[List[Dict[str, Any]],
     return system_prompts, repaired_bedrock_formatted_messages
 
 
-def _parse_completion_response(response_body: Dict[str, Any], model: str) -> List[ChatMessage]:
+def _parse_completion_response(response_body: dict[str, Any], model: str) -> list[ChatMessage]:
     """
     Parse a Bedrock API response into Haystack ChatMessage objects.
 
@@ -346,7 +346,7 @@ def _parse_completion_response(response_body: Dict[str, Any], model: str) -> Lis
 
 
 def _convert_event_to_streaming_chunk(
-    event: Dict[str, Any], model: str, component_info: ComponentInfo
+    event: dict[str, Any], model: str, component_info: ComponentInfo
 ) -> StreamingChunk:
     """
     Convert a Bedrock streaming event to a Haystack StreamingChunk.
@@ -433,7 +433,7 @@ def _convert_event_to_streaming_chunk(
 
     elif "metadata" in event:
         event_meta = event["metadata"]
-        chunk_meta: Dict[str, Any] = {**base_meta}
+        chunk_meta: dict[str, Any] = {**base_meta}
 
         if "usage" in event_meta:
             usage = event_meta["usage"]
@@ -454,7 +454,7 @@ def _convert_event_to_streaming_chunk(
     return streaming_chunk
 
 
-def _process_reasoning_contents(chunks: List[StreamingChunk]) -> Optional[ReasoningContent]:
+def _process_reasoning_contents(chunks: list[StreamingChunk]) -> Optional[ReasoningContent]:
     """
     Process reasoning contents from a list of StreamingChunk objects into the Bedrock expected format.
 
@@ -536,7 +536,7 @@ def _parse_streaming_response(
     streaming_callback: SyncStreamingCallbackT,
     model: str,
     component_info: ComponentInfo,
-) -> List[ChatMessage]:
+) -> list[ChatMessage]:
     """
     Parse a streaming response from Bedrock.
 
@@ -547,7 +547,7 @@ def _parse_streaming_response(
     :return: List of ChatMessage objects
     """
     content_block_idxs = set()
-    chunks: List[StreamingChunk] = []
+    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
@@ -581,7 +581,7 @@ async def _parse_streaming_response_async(
     streaming_callback: AsyncStreamingCallbackT,
     model: str,
     component_info: ComponentInfo,
-) -> List[ChatMessage]:
+) -> list[ChatMessage]:
     """
     Parse a streaming response from Bedrock.
 
@@ -592,7 +592,7 @@ async def _parse_streaming_response_async(
     :return: List of ChatMessage objects
     """
     content_block_idxs = set()
-    chunks: List[StreamingChunk] = []
+    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
@@ -613,7 +613,7 @@ async def _parse_streaming_response_async(
     return [reply]
 
 
-def _validate_guardrail_config(guardrail_config: Optional[Dict[str, str]] = None, streaming: bool = False) -> None:
+def _validate_guardrail_config(guardrail_config: Optional[dict[str, str]] = None, streaming: bool = False) -> None:
     """
     Validate the guardrail configuration.
 

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, Union, get_args
+from typing import Any, Callable, ClassVar, Literal, Optional, Union, get_args
 
 from botocore.config import Config
 from botocore.exceptions import ClientError
@@ -62,7 +62,7 @@ class AmazonBedrockGenerator:
     supports Amazon Bedrock.
     """
 
-    SUPPORTED_MODEL_PATTERNS: ClassVar[Dict[str, Type[BedrockModelAdapter]]] = {
+    SUPPORTED_MODEL_PATTERNS: ClassVar[dict[str, type[BedrockModelAdapter]]] = {
         r"([a-z]{2}\.)?amazon.titan-text.*": AmazonTitanAdapter,
         r"([a-z]{2}\.)?ai21.j2.*": AI21LabsJurassic2Adapter,
         r"([a-z]{2}\.)?cohere.command-[^r].*": CohereCommandAdapter,
@@ -72,7 +72,7 @@ class AmazonBedrockGenerator:
         r"([a-z]{2}\.)?mistral.*": MistralAdapter,
     }
 
-    SUPPORTED_MODEL_FAMILIES: ClassVar[Dict[str, Type[BedrockModelAdapter]]] = {
+    SUPPORTED_MODEL_FAMILIES: ClassVar[dict[str, type[BedrockModelAdapter]]] = {
         "amazon.titan-text": AmazonTitanAdapter,
         "ai21.j2": AI21LabsJurassic2Adapter,
         "cohere.command": CohereCommandAdapter,
@@ -105,7 +105,7 @@ class AmazonBedrockGenerator:
         max_length: Optional[int] = None,
         truncate: Optional[bool] = None,
         streaming_callback: Optional[Callable[[StreamingChunk], None]] = None,
-        boto3_config: Optional[Dict[str, Any]] = None,
+        boto3_config: Optional[dict[str, Any]] = None,
         model_family: Optional[MODEL_FAMILIES] = None,
         **kwargs: Any,
     ) -> None:
@@ -183,13 +183,13 @@ class AmazonBedrockGenerator:
         model_adapter_cls = self.get_model_adapter(model=model, model_family=model_family)
         self.model_adapter = model_adapter_cls(model_kwargs=model_input_kwargs, max_length=self.max_length)
 
-    @component.output_types(replies=List[str], meta=Dict[str, Any])
+    @component.output_types(replies=list[str], meta=dict[str, Any])
     def run(
         self,
         prompt: str,
         streaming_callback: Optional[Callable[[StreamingChunk], None]] = None,
-        generation_kwargs: Optional[Dict[str, Any]] = None,
-    ) -> Dict[str, Union[List[str], Dict[str, Any]]]:
+        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.
 
@@ -240,7 +240,7 @@ class AmazonBedrockGenerator:
         return {"replies": replies, "meta": metadata}
 
     @classmethod
-    def get_model_adapter(cls, model: str, model_family: Optional[str] = None) -> Type[BedrockModelAdapter]:
+    def get_model_adapter(cls, model: str, model_family: Optional[str] = None) -> type[BedrockModelAdapter]:
         """
         Gets the model adapter for the given model.
 
@@ -273,7 +273,7 @@ class AmazonBedrockGenerator:
         )
         raise AmazonBedrockConfigurationError(msg)
 
-    def to_dict(self) -> Dict[str, Any]:
+    def to_dict(self) -> dict[str, Any]:
         """
         Serializes the component to a dictionary.
 
@@ -297,7 +297,7 @@ class AmazonBedrockGenerator:
         )
 
     @classmethod
-    def from_dict(cls, data: Dict[str, Any]) -> "AmazonBedrockGenerator":
+    def from_dict(cls, data: dict[str, Any]) -> "AmazonBedrockGenerator":
         """
         Deserializes the component from a dictionary.
 

haystack_integrations/components/rankers/amazon_bedrock/ranker.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, List, Optional
+from typing import Any, Optional
 
 from botocore.exceptions import ClientError
 from haystack import Document, component, default_from_dict, default_to_dict, logging
@@ -68,7 +68,7 @@ class AmazonBedrockRanker:
         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
         max_chunks_per_doc: Optional[int] = None,
-        meta_fields_to_embed: Optional[List[str]] = None,
+        meta_fields_to_embed: Optional[list[str]] = None,
         meta_data_separator: str = "\n",
     ) -> None:
         if not model:
@@ -122,7 +122,7 @@ class AmazonBedrockRanker:
             )
             raise AmazonBedrockConfigurationError(msg) from exception
 
-    def to_dict(self) -> Dict[str, Any]:
+    def to_dict(self) -> dict[str, Any]:
         """
         Serializes the component to a dictionary.
 
@@ -144,7 +144,7 @@ class AmazonBedrockRanker:
         )
 
     @classmethod
-    def from_dict(cls, data: Dict[str, Any]) -> "AmazonBedrockRanker":
+    def from_dict(cls, data: dict[str, Any]) -> "AmazonBedrockRanker":
         """
         Deserializes the component from a dictionary.
 
@@ -159,7 +159,7 @@ class AmazonBedrockRanker:
         )
         return default_from_dict(cls, data)
 
-    def _prepare_bedrock_input_docs(self, documents: List[Document]) -> List[str]:
+    def _prepare_bedrock_input_docs(self, documents: list[Document]) -> list[str]:
         """
         Prepare the input by concatenating the document text with the metadata fields specified.
         :param documents: The list of Document objects.
@@ -176,8 +176,8 @@ class AmazonBedrockRanker:
 
         return concatenated_input_list
 
-    @component.output_types(documents=List[Document])
-    def run(self, query: str, documents: List[Document], top_k: Optional[int] = None) -> Dict[str, List[Document]]:
+    @component.output_types(documents=list[Document])
+    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.