diff --git a/litellm/llms/bedrock/embed/cohere_transformation.py b/litellm/llms/bedrock/embed/cohere_transformation.py index 490cd71b793..d00cb74aae0 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 a85aaafe23d..6293efe9e09 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 7c4eec7ba32..a5df9381fc7 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()