litellm/tests/test_litellm/llms/sagemaker
milan-berri 9600fda2cc
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 <cursoragent@cursor.com>

* 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 <cursoragent@cursor.com>

* 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 <cursoragent@cursor.com>

* 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 <cursoragent@cursor.com>

* 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 <cursoragent@cursor.com>

---------

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-05-22 12:00:42 -07:00
..
test_sagemaker_common_utils.py refactor(bedrock/sagemaker): switch to lazy loading for response stre… (#28189) 2026-05-18 23:21:04 -07:00
test_sagemaker_embedding_role_assumption.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00
test_sagemaker_embedding_voyage.py fix(sagemaker): send native Cohere embed payload to Cohere SageMaker endpoints (#28613) 2026-05-22 12:00:42 -07:00
test_sagemaker_nova_transformation.py style: run black formatter on files from main merge 2026-04-17 13:02:59 -07:00