From 9600fda2cc94024182ce395093f21854d43a1aba Mon Sep 17 00:00:00 2001 From: milan-berri Date: Fri, 22 May 2026 22:00:42 +0300 Subject: [PATCH] fix(sagemaker): send native Cohere embed payload to Cohere SageMaker endpoints (#28613) * fix(sagemaker): use Cohere embed payload for Marketplace endpoints SageMaker embedding only special-cased Voyage; every other endpoint received HuggingFace TGI `{"inputs": [...]}`. AWS Marketplace Cohere containers expect the native Cohere embed payload (`texts`, `input_type`) and reject the HF shape with `422 EmbedReqV2.inputs is of type string but should be of type Object`. Add `SagemakerCohereEmbeddingConfig` that reuses Bedrock/Cohere request and response transforms, and route SageMaker endpoint names containing `cohere` or a Cohere embed model fragment (`embed-multilingual`, `embed-english`, `embed-v3`, `embed-v4`) to it. Supports `input_type`, `dimensions`, and `encoding_format`. Voyage and HuggingFace SageMaker endpoints are unchanged. Co-authored-by: Cursor * refactor(sagemaker): simplify cohere detection and align with file conventions - Detect Cohere SageMaker endpoints with a single `"cohere" in model.lower()` check, mirroring the existing Voyage branch instead of a separate helper function and marker constant. - Drop instance caches of sub-configs; instantiate `BedrockCohereEmbeddingConfig` / `CohereEmbeddingConfig` per call to match the existing pattern in `BedrockCohereEmbeddingConfig._transform_request`. - Match `SagemakerEmbeddingConfig`'s signatures, defaults, and `Any` typing for `logging_obj`; collapse the input-normalization helper inline. - Inline `transform_embedding_response` input lookup; no behavior change. Co-authored-by: Cursor * fix(sagemaker): restore provider-supported embedding params after map Cohere input_type is advertised in get_supported_openai_params but was filtered out of non_default_params by OPENAI_EMBEDDING_PARAMS before map_openai_params ran. Merge supported params from passed_params after map (same path Greptile flagged). Handle input_type explicitly in SagemakerCohereEmbeddingConfig.map_openai_params and add an integration test through get_optional_params_embeddings. Co-authored-by: Cursor * fix(embeddings): only restore non-OpenAI supported params after map The post-map restore loop must skip OPENAI_EMBEDDING_PARAMS so mapped fields (e.g. dimensions -> output_dimension) are not duplicated under their OpenAI names. Align SageMaker embedding import order with sibling files and add a regression test for dimensions mapping. Co-authored-by: Cursor * fix(sagemaker): avoid double post_call on Cohere embedding response Greptile review on #28613 caught that `CohereEmbeddingConfig._transform_response` calls `logging_obj.post_call` internally. The SageMaker embedding handler already calls `post_call` once before invoking the transform, so the Cohere SageMaker path fired callbacks, cost calculators, and log handlers twice per request. Extract the parsing body of `_transform_response` into `_populate_embedding_response` (pure extract-method, no behavior change for existing Cohere direct or Bedrock Cohere paths, which keep calling `_transform_response`). Have `SagemakerCohereEmbeddingConfig` call the new helper directly so it parses the response without re-logging. Add a regression test asserting `logging_obj.post_call` is not invoked by the SageMaker Cohere transform. Co-authored-by: Cursor --------- Co-authored-by: Cursor --- .../llms/cohere/embed/v1_transformation.py | 35 +++- litellm/llms/sagemaker/completion/handler.py | 2 +- .../embedding/cohere_transformation.py | 141 +++++++++++++++ .../sagemaker/embedding/transformation.py | 22 ++- litellm/utils.py | 15 ++ .../test_sagemaker_embedding_voyage.py | 169 ++++++++++++++++++ 6 files changed, 365 insertions(+), 19 deletions(-) create mode 100644 litellm/llms/sagemaker/embedding/cohere_transformation.py diff --git a/litellm/llms/cohere/embed/v1_transformation.py b/litellm/llms/cohere/embed/v1_transformation.py index feca9cb5b88..82c901e7eca 100644 --- a/litellm/llms/cohere/embed/v1_transformation.py +++ b/litellm/llms/cohere/embed/v1_transformation.py @@ -110,15 +110,35 @@ class CohereEmbeddingConfig: additional_args={"complete_input_dict": data}, original_response=response_json, ) + return self._populate_embedding_response( + response_json=response_json, + model_response=model_response, + model=model, + encoding=encoding, + input=input, + ) + + def _populate_embedding_response( + self, + response_json: dict, + model_response: EmbeddingResponse, + model: str, + encoding: Any, + input: list, + ) -> EmbeddingResponse: """ - response + Parse a Cohere embed response body into an OpenAI-style EmbeddingResponse. + + Split out from `_transform_response` so callers that already log + `post_call` themselves (e.g. SageMaker's embedding handler) can reuse + the parsing without triggering a second `post_call`. + + Response shape: { 'object': "list", - 'data': [ - - ] - 'model', - 'usage' + 'data': [...], + 'model', + 'usage', } """ embeddings = response_json["embeddings"] @@ -149,9 +169,6 @@ class CohereEmbeddingConfig: model_response.object = "list" model_response.data = output_data model_response.model = model - input_tokens = 0 - for text in input: - input_tokens += len(encoding.encode(text)) setattr( model_response, diff --git a/litellm/llms/sagemaker/completion/handler.py b/litellm/llms/sagemaker/completion/handler.py index efbb218f575..de7be18e8ba 100644 --- a/litellm/llms/sagemaker/completion/handler.py +++ b/litellm/llms/sagemaker/completion/handler.py @@ -578,7 +578,7 @@ class SagemakerLLM(BaseAWSLLM): logger_fn=None, ): """ - Supports both Huggingface Jumpstart embeddings and Voyage models + Supports Hugging Face (TGI), Voyage, and Cohere embedding endpoints """ ### BOTO3 INIT import boto3 diff --git a/litellm/llms/sagemaker/embedding/cohere_transformation.py b/litellm/llms/sagemaker/embedding/cohere_transformation.py new file mode 100644 index 00000000000..fdb67202ebb --- /dev/null +++ b/litellm/llms/sagemaker/embedding/cohere_transformation.py @@ -0,0 +1,141 @@ +""" +Translate from OpenAI's `/v1/embeddings` to Sagemaker's `/invoke` + +In the native Cohere embed format for self-hosted Cohere endpoints +(AWS Marketplace / JumpStart). Cohere containers expect +`{"texts": [...], "input_type": "..."}` and reject the HuggingFace TGI shape +`{"inputs": [...]}` with `422 EmbedReqV2.inputs is of type string but should +be of type Object`. + +Reference: https://docs.cohere.com/v2/reference/embed +""" + +from typing import TYPE_CHECKING, Any, List, Optional, Union, cast + +if TYPE_CHECKING: + from litellm.types.llms.openai import AllEmbeddingInputValues + +from httpx._models import Headers, Response + +import litellm +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig +from litellm.llms.bedrock.embed.cohere_transformation import ( + BedrockCohereEmbeddingConfig, +) +from litellm.llms.cohere.embed.v1_transformation import CohereEmbeddingConfig +from litellm.types.utils import EmbeddingResponse + +from ..common_utils import SagemakerError + + +class SagemakerCohereEmbeddingConfig(BaseEmbeddingConfig): + """ + SageMaker invoke payload for self-hosted Cohere embed models. + """ + + def __init__(self) -> None: + pass + + def get_supported_openai_params(self, model: str) -> List[str]: + return ["encoding_format", "dimensions", "input_type"] + + def map_openai_params( + self, + non_default_params: dict, + optional_params: dict, + model: str, + drop_params: bool, + ) -> dict: + optional_params = BedrockCohereEmbeddingConfig().map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + ) + if "input_type" in non_default_params: + optional_params["input_type"] = non_default_params["input_type"] + return optional_params + + def get_error_class( + self, error_message: str, status_code: int, headers: Union[dict, Headers] + ) -> BaseLLMException: + return SagemakerError( + message=error_message, status_code=status_code, headers=headers + ) + + def transform_embedding_request( + self, + model: str, + input: "AllEmbeddingInputValues", + optional_params: dict, + headers: dict, + ) -> dict: + """ + Transform embedding request for Cohere models on SageMaker + """ + if isinstance(input, str): + input_list: List[str] = [input] + elif isinstance(input, list): + if input and (isinstance(input[0], list) or isinstance(input[0], int)): + raise ValueError("Input must be a list of strings") + input_list = cast(List[str], input) + else: + input_list = [str(input)] + + return dict( + BedrockCohereEmbeddingConfig()._transform_request( + model=model, + input=input_list, + inference_params=optional_params, + ) + ) + + def transform_embedding_response( + self, + model: str, + raw_response: Response, + model_response: "EmbeddingResponse", + logging_obj: Any, + api_key: Optional[str] = None, + request_data: dict = {}, + optional_params: dict = {}, + litellm_params: dict = {}, + ) -> "EmbeddingResponse": + """ + Transform embedding response for Cohere models on SageMaker. + + Uses `CohereEmbeddingConfig._populate_embedding_response` (not + `_transform_response`) so we do not log `post_call` a second time + — the SageMaker embedding handler already logs `post_call` before + invoking this transform. + """ + input_value = ( + logging_obj.model_call_details.get("input") + or request_data.get("texts") + or request_data.get("images") + or [] + ) + if isinstance(input_value, str): + input_value = [input_value] + + return CohereEmbeddingConfig()._populate_embedding_response( + response_json=raw_response.json(), + model_response=model_response, + model=model, + encoding=litellm.encoding, + input=input_value, + ) + + def validate_environment( + self, + headers: dict, + model: str, + messages: List[Any], + optional_params: dict, + litellm_params: dict, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + ) -> dict: + """ + Validate environment for SageMaker Cohere embeddings + """ + return {"Content-Type": "application/json"} diff --git a/litellm/llms/sagemaker/embedding/transformation.py b/litellm/llms/sagemaker/embedding/transformation.py index 09bdb9295e7..5e2aa99534f 100644 --- a/litellm/llms/sagemaker/embedding/transformation.py +++ b/litellm/llms/sagemaker/embedding/transformation.py @@ -11,12 +11,13 @@ if TYPE_CHECKING: from httpx._models import Headers, Response -from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.llms.base_llm.chat.transformation import BaseLLMException -from litellm.types.utils import Usage, EmbeddingResponse +from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig from litellm.llms.voyage.embedding.transformation import VoyageEmbeddingConfig +from litellm.types.utils import EmbeddingResponse, Usage from ..common_utils import SagemakerError +from .cohere_transformation import SagemakerCohereEmbeddingConfig class SagemakerEmbeddingConfig(BaseEmbeddingConfig): @@ -38,17 +39,20 @@ class SagemakerEmbeddingConfig(BaseEmbeddingConfig): Returns: Appropriate embedding config instance """ - if "voyage" in model.lower(): + model_lower = model.lower() + if "voyage" in model_lower: return VoyageEmbeddingConfig() - else: - return cls() + if "cohere" in model_lower: + return SagemakerCohereEmbeddingConfig() + return cls() def get_supported_openai_params(self, model: str) -> List[str]: - # Check if this is an embedding model - if "voyage" in model.lower(): + model_lower = model.lower() + if "voyage" in model_lower: return VoyageEmbeddingConfig().get_supported_openai_params(model) - else: - return [] + if "cohere" in model_lower: + return SagemakerCohereEmbeddingConfig().get_supported_openai_params(model) + return [] def map_openai_params( self, diff --git a/litellm/utils.py b/litellm/utils.py index 18ee811f0f1..c28a88e0f1c 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -3350,6 +3350,21 @@ def get_optional_params_embeddings( # noqa: PLR0915 model=model, drop_params=drop_params if drop_params is not None else False, ) + # Provider-only params (e.g. Cohere input_type) are not in + # OPENAI_EMBEDDING_PARAMS, so embedding_pre_process drops them from + # non_default_params before map_openai_params. Restore only those extras + # from passed_params — skip OPENAI_EMBEDDING_PARAMS to avoid duplicating + # values already mapped (e.g. dimensions -> output_dimension). + if supported_params: + for param in supported_params: + if param in OPENAI_EMBEDDING_PARAMS: + continue + if ( + param in passed_params + and passed_params[param] is not None + and param not in optional_params + ): + optional_params[param] = passed_params[param] ## raise exception if non-default value passed for non-openai/azure embedding calls elif custom_llm_provider == "openai": # 'dimensions` is only supported in `text-embedding-3` and later models diff --git a/tests/test_litellm/llms/sagemaker/test_sagemaker_embedding_voyage.py b/tests/test_litellm/llms/sagemaker/test_sagemaker_embedding_voyage.py index a36aec32d13..943a3160bb7 100644 --- a/tests/test_litellm/llms/sagemaker/test_sagemaker_embedding_voyage.py +++ b/tests/test_litellm/llms/sagemaker/test_sagemaker_embedding_voyage.py @@ -17,6 +17,9 @@ import pytest sys.path.insert(0, os.path.abspath("../../../../..")) from litellm import embedding +from litellm.llms.sagemaker.embedding.cohere_transformation import ( + SagemakerCohereEmbeddingConfig, +) from litellm.llms.sagemaker.embedding.transformation import SagemakerEmbeddingConfig from litellm.llms.voyage.embedding.transformation import VoyageEmbeddingConfig from litellm.types.utils import EmbeddingResponse, Usage @@ -54,6 +57,172 @@ class TestSagemakerEmbeddingFactory: assert isinstance(config2, VoyageEmbeddingConfig) assert isinstance(config3, VoyageEmbeddingConfig) + def test_get_model_config_cohere_model(self): + """Cohere SageMaker endpoints route to SagemakerCohereEmbeddingConfig""" + for endpoint_name in ( + "cohere.embed-multilingual-v3", + "cohere-embed-english-v3-prod", + "my-cohere-marketplace-endpoint", + "COHERE-EMBED-V4", + ): + config = SagemakerEmbeddingConfig.get_model_config(endpoint_name) + assert isinstance(config, SagemakerCohereEmbeddingConfig), endpoint_name + + +class TestSagemakerCohereEmbeddingConfig: + """Cohere-specific SageMaker embedding request/response transforms""" + + def setup_method(self): + self.config = SagemakerCohereEmbeddingConfig() + + MODEL = "cohere.embed-multilingual-v3" + + def test_transform_request_uses_cohere_payload(self): + """Bug repro: request must use `texts` + `input_type`, not HF `inputs`""" + result = self.config.transform_embedding_request( + model=self.MODEL, + input=["hello"], + optional_params={"input_type": "search_query"}, + headers={}, + ) + assert "inputs" not in result + assert result["texts"] == ["hello"] + assert result["input_type"] == "search_query" + + def test_transform_request_default_input_type(self): + result = self.config.transform_embedding_request( + model=self.MODEL, + input=["hello"], + optional_params={}, + headers={}, + ) + assert result["texts"] == ["hello"] + assert result["input_type"] == "search_document" + + def test_transform_request_normalizes_string_input(self): + result = self.config.transform_embedding_request( + model=self.MODEL, + input="hello", + optional_params={}, + headers={}, + ) + assert result["texts"] == ["hello"] + + def test_map_openai_params_dimensions_to_output_dimension(self): + params = self.config.map_openai_params( + non_default_params={"dimensions": 512, "encoding_format": "float"}, + optional_params={}, + model=self.MODEL, + drop_params=False, + ) + assert params["output_dimension"] == 512 + assert params["embedding_types"] == ["float"] + + def test_map_openai_params_input_type_from_non_default_params(self): + params = self.config.map_openai_params( + non_default_params={"input_type": "search_query"}, + optional_params={}, + model=self.MODEL, + drop_params=False, + ) + assert params["input_type"] == "search_query" + + def test_get_optional_params_embeddings_preserves_input_type(self): + """Exercises get_optional_params_embeddings, not transform in isolation.""" + from litellm.utils import get_optional_params_embeddings + + optional_params = get_optional_params_embeddings( + model=self.MODEL, + custom_llm_provider="sagemaker", + input_type="search_query", + ) + assert optional_params.get("input_type") == "search_query" + + body = self.config.transform_embedding_request( + model=self.MODEL, + input=["hello"], + optional_params=optional_params, + headers={}, + ) + assert body["texts"] == ["hello"] + assert body["input_type"] == "search_query" + + def test_get_optional_params_embeddings_maps_dimensions_without_duplicate(self): + """dimensions must map to output_dimension only, not also stay as dimensions.""" + from litellm.utils import get_optional_params_embeddings + + optional_params = get_optional_params_embeddings( + model=self.MODEL, + custom_llm_provider="sagemaker", + dimensions=512, + input_type="search_query", + ) + assert optional_params.get("output_dimension") == 512 + assert "dimensions" not in optional_params + assert optional_params.get("input_type") == "search_query" + + def test_transform_response_parses_cohere_payload(self): + cohere_response = { + "embeddings": [[0.1, 0.2, 0.3]], + "meta": {"billed_units": {"input_tokens": 2}}, + } + mock_response = httpx.Response( + status_code=200, + content=json.dumps(cohere_response).encode("utf-8"), + headers={"content-type": "application/json"}, + ) + logging_obj = MagicMock() + logging_obj.model_call_details = {"input": ["hello"]} + + result = self.config.transform_embedding_response( + model=self.MODEL, + raw_response=mock_response, + model_response=EmbeddingResponse(), + logging_obj=logging_obj, + api_key=None, + request_data={"texts": ["hello"], "input_type": "search_query"}, + optional_params={}, + litellm_params={}, + ) + + assert result.object == "list" + assert len(result.data) == 1 + assert result.data[0]["embedding"] == [0.1, 0.2, 0.3] + assert result.usage.prompt_tokens == 2 + + def test_transform_response_does_not_double_call_post_call(self): + """ + Greptile review fix: SageMaker handler already calls + `logging_obj.post_call` once before invoking + `transform_embedding_response`. The transform must NOT call it again, + otherwise callbacks, cost calculators, and log handlers double-fire + for every Cohere SageMaker embedding call. + """ + cohere_response = { + "embeddings": [[0.1, 0.2, 0.3]], + "meta": {"billed_units": {"input_tokens": 2}}, + } + mock_response = httpx.Response( + status_code=200, + content=json.dumps(cohere_response).encode("utf-8"), + headers={"content-type": "application/json"}, + ) + logging_obj = MagicMock() + logging_obj.model_call_details = {"input": ["hello"]} + + self.config.transform_embedding_response( + model=self.MODEL, + raw_response=mock_response, + model_response=EmbeddingResponse(), + logging_obj=logging_obj, + api_key=None, + request_data={"texts": ["hello"], "input_type": "search_query"}, + optional_params={}, + litellm_params={}, + ) + + logging_obj.post_call.assert_not_called() + class TestVoyageEmbeddingConfig: """Test Voyage-specific embedding configuration"""