feat: enhance Cohere embedding support with additional parameters and model version

This commit is contained in:
amirzaushnizer 2026-02-01 18:04:08 +02:00
parent df387c39f6
commit f0853b2564
3 changed files with 5 additions and 2 deletions

View file

@ -15,7 +15,7 @@ class BedrockCohereEmbeddingConfig:
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 @@ class BedrockCohereEmbeddingConfig:
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:

View file

@ -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):

View file

@ -3242,7 +3242,7 @@ def get_optional_params_embeddings( # noqa: PLR0915
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()