amazon-bedrock-haystack 5.2.0__tar.gz → 5.3.1__tar.gz

{amazon_bedrock_haystack-5.2.0 → amazon_bedrock_haystack-5.3.1}/CHANGELOG.md

@@ -1,32 +1,49 @@
 # Changelog
 
-## [integrations/amazon_bedrock-v5.1.0] - 2025-09-29
+## [integrations/amazon_bedrock-v5.3.0] - 2025-12-17
 
 ### 🚀 Features
 
-- S3Downloader - add `s3_key_generation_function` param to customize S3 key generation (#2343)
+- `AmazonBedrockChatGenerator` update tools param to ToolsType (#2415)
+- Cohere Embed v4 support in Bedrock (#2612)
 
-## [integrations/amazon_bedrock-v5.0.0] - 2025-09-22
+### 📚 Documentation
+
+- Add pydoc configurations for Docusaurus (#2411)
+
+### ⚙️ CI
+
+- Change pytest command (#2475)
 
 ### 🧹 Chores
 
-- [**breaking**] Remove deprecated `BedrockRanker` (use `AmazonBedrockRanker` instead) (#2287)
+- Remove Readme API CI workflow and configs (#2573)
+
+### 🌀 Miscellaneous
 
-## [integrations/amazon_bedrock-v4.2.0] - 2025-09-19
+- Adopt PEP 585 type hinting (part 2) (#2508)
+
+## [integrations/amazon_bedrock-v5.1.0] - 2025-09-29
 
 ### 🚀 Features
 
-- Add a new `S3Downloader` component (#2192)
+- S3Downloader - add `s3_key_generation_function` param to customize S3 key generation (#2343)
+
 
-## [integrations/amazon_bedrock-v4.1.0] - 2025-09-19
+## [integrations/amazon_bedrock-v5.0.0] - 2025-09-22
 
 ### 🚀 Features
 
 - Support AWS Bedrock Guardrails in `AmazonBedrockChatGenerator` (#2284)
+- Add a new `S3Downloader` component (#2192)
+
+### 📚 Documentation
+
 
 ### 🧹 Chores
 
 - Bedrock - remove unused `stop_words` init parameter (#2275)
+- [**breaking**] Remove deprecated `BedrockRanker` (use `AmazonBedrockRanker` instead) (#2287)
 
 ### 🌀 Miscellaneous
 

{amazon_bedrock_haystack-5.2.0 → amazon_bedrock_haystack-5.3.1}/PKG-INFO

@@ -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 → amazon_bedrock_haystack-5.3.1}/pyproject.toml

@@ -46,7 +46,7 @@ installer = "uv"
 dependencies = ["haystack-pydoc-tools", "ruff"]
 
 [tool.hatch.envs.default.scripts]
-docs = ["pydoc-markdown pydoc/config.yml"]
+docs = ["pydoc-markdown pydoc/config_docusaurus.yml"]
 fmt = "ruff check --fix {args} && ruff format {args}"
 fmt-check = "ruff check {args} && ruff format --check {args}"
 
@@ -66,7 +66,7 @@ dependencies = [
 unit = 'pytest -m "not integration" {args:tests}'
 integration = 'pytest -m "integration" {args:tests}'
 all = 'pytest {args:tests}'
-cov-retry = 'all --cov=haystack_integrations --reruns 3 --reruns-delay 30 -x'
+cov-retry = 'pytest --cov=haystack_integrations --reruns 3 --reruns-delay 30 -x {args:tests}'
 
 types = """mypy -p haystack_integrations.common.amazon_bedrock \
 -p haystack_integrations.components.embedders.amazon_bedrock \
@@ -92,7 +92,7 @@ module = [
 ignore_missing_imports = true
 
 [tool.ruff]
-target-version = "py38"
+target-version = "py39"
 line-length = 120
 
 [tool.ruff.lint]

{amazon_bedrock_haystack-5.2.0 → amazon_bedrock_haystack-5.3.1}/src/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:

{amazon_bedrock_haystack-5.2.0 → amazon_bedrock_haystack-5.3.1}/src/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.
 

{amazon_bedrock_haystack-5.2.0 → amazon_bedrock_haystack-5.3.1}/src/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

{amazon_bedrock_haystack-5.2.0 → amazon_bedrock_haystack-5.3.1}/src/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.