From f0853b2564ebe52663580c168431327ea87cb438 Mon Sep 17 00:00:00 2001 From: amirzaushnizer Date: Sun, 1 Feb 2026 18:04:08 +0200 Subject: [PATCH] feat: enhance Cohere embedding support with additional parameters and model version --- litellm/llms/bedrock/embed/cohere_transformation.py | 4 +++- litellm/types/llms/bedrock.py | 1 + litellm/utils.py | 2 +- 3 files changed, 5 insertions(+), 2 deletions(-) diff --git a/litellm/llms/bedrock/embed/cohere_transformation.py b/litellm/llms/bedrock/embed/cohere_transformation.py index 490cd71b79..d00cb74aae 100644 --- a/litellm/llms/bedrock/embed/cohere_transformation.py +++ b/litellm/llms/bedrock/embed/cohere_transformation.py @@ -15,7 +15,7 @@ def __init__(self) -> None: pass def get_supported_openai_params(self) -> List[str]: - return ["encoding_format"] + return ["encoding_format", "dimensions"] def map_openai_params( self, non_default_params: dict, optional_params: dict @@ -23,6 +23,8 @@ def map_openai_params( for k, v in non_default_params.items(): if k == "encoding_format": optional_params["embedding_types"] = v + elif k == "dimensions": + optional_params["output_dimension"] = v return optional_params def _is_v3_model(self, model: str) -> bool: diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py index a85aaafe23..6293efe9e0 100644 --- a/litellm/types/llms/bedrock.py +++ b/litellm/types/llms/bedrock.py @@ -397,6 +397,7 @@ class CohereEmbeddingRequest(TypedDict, total=False): input_type: Required[COHERE_EMBEDDING_INPUT_TYPES] truncate: Literal["NONE", "START", "END"] embedding_types: Literal["float", "int8", "uint8", "binary", "ubinary"] + output_dimension: int class CohereEmbeddingRequestWithModel(CohereEmbeddingRequest): diff --git a/litellm/utils.py b/litellm/utils.py index 7c4eec7ba3..a5df9381fc 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -3242,7 +3242,7 @@ def _check_valid_arg(supported_params: Optional[list]): object = litellm.AmazonTitanMultimodalEmbeddingG1Config() elif "amazon.titan-embed-text-v2:0" in model: object = litellm.AmazonTitanV2Config() - elif "cohere.embed-multilingual-v3" in model: + elif "cohere.embed-multilingual-v3" in model or "cohere.embed-v4" in model: object = litellm.BedrockCohereEmbeddingConfig() elif "twelvelabs" in model or "marengo" in model: object = litellm.TwelveLabsMarengoEmbeddingConfig()