fix(sagemaker): send native Cohere embed payload to Cohere SageMaker endpoints

[milan-berri]

Relevant issues

Customer report: SageMaker + AWS Marketplace Cohere (cohere.embed-multilingual-v3) fails via LiteLLM with 422 EmbedReqV2.inputs is of type string but should be of type Object; same endpoint works with direct boto3 and Bedrock. Reproduced on LiteLLM 1.81.15 and 1.84.0.

Root cause: SagemakerEmbeddingConfig only special-cases Voyage; all other endpoints get HuggingFace TGI {"inputs": [...]}. Marketplace Cohere containers expect {"texts": [...], "input_type": "..."}.

Fix: Add SagemakerCohereEmbeddingConfig (reuses Bedrock/Cohere request + response transforms). Route when "cohere" in model.lower(), same pattern as Voyage. HF and Voyage SageMaker paths unchanged.

Linear ticket

Fixes LIT-3289

Pre-Submission checklist

Please complete all items before asking a LiteLLM maintainer to review your PR

• I have Added testing in the tests/test_litellm/ directory, Adding at least 1 test is a hard requirement - see details
• My PR passes all unit tests on make test-unit (CI on PR)
• My PR's scope is as isolated as possible, it only solves 1 specific problem
• I have requested a Greptile review by commenting @greptileai and received a Confidence Score of at least 4/5 before requesting a maintainer review

Delays in PR merge?

If you're seeing a delay in your PR being merged, ping the LiteLLM Team on Slack (#pr-review).

CI (LiteLLM team)

CI status guideline:

• 50-55 passing tests: main is stable with minor issues.
• 45-49 passing tests: acceptable but needs attention
• <= 40 passing tests: unstable; be careful with your merges and assess the risk.

• Branch creation CI run
  Link: (fill after PR opened)

• CI run for the last commit
  Link: (fill after PR opened)

• Merge / cherry-pick CI run
  Links: (fill after PR opened)

Screenshots / Proof of Fix

What we validated

End-to-end litellm.embedding() on the SageMaker provider path — the same code path the customer hits (SagemakerLLM.embedding → transform_embedding_request → boto3 invoke_endpoint). Not a transform-only unit test in isolation.

Check	What it proves
Request body	LiteLLM sends the JSON the container actually expects before any AWS call is made
Success path	Cohere-shaped mock response is parsed into OpenAI-style EmbeddingResponse
Factory routing	Only Cohere-named endpoints pick up the new config; Voyage and HF defaults stay the same

Method: boto3.Session is patched so invoke_endpoint never runs against AWS. The mock records Body, then either returns a minimal Cohere embeddings payload or raises ClientError with the customer's 422 text. Input: model="sagemaker/cohere-embed-multilingual-v3-prod", input=["hello"], input_type="search_query".

Request body (customer failure mode)

	Staging (bug)	This PR
Body sent to SageMaker	{"inputs": ["hello"], "input_type": "search_query"}	{"texts": ["hello"], "input_type": "search_query"}
Matches HF TGI	yes — wrong for Cohere	no
Matches Cohere Marketplace API	no — causes customer 422	yes

On staging, a mock that returns the customer's 422 reproduces BadRequestError with EmbedReqV2.inputs is of type string but should be of type Object — same message as their log.

On this branch, the body is correct in both mock modes below; the 422 no longer reflects a LiteLLM formatting bug.

Mock scenarios (runtime)

A — Happy path: mock returns {"embeddings": [[0.1, 0.2, 0.3]], "meta": {...}}. Call completes; data[0].embedding populated. Confirms response transform works with the new request shape.

B — Server error path: mock raises the exact 422 string from the ticket. LiteLLM still surfaces BadRequestError, but the captured body is already texts + input_type — i.e. we are not sending inputs anymore; a real Cohere endpoint would not reject the payload for this reason.

Factory routing (no boto3)

SagemakerEmbeddingConfig.get_model_config(endpoint_name) for representative names:

Endpoint name	Config	Expected
cohere.embed-multilingual-v3	SagemakerCohereEmbeddingConfig	new
cohere-marketplace-endpoint	SagemakerCohereEmbeddingConfig	new
COHERE-EMBED-V4	SagemakerCohereEmbeddingConfig	new (case-insensitive)
voyage-3-5-embedding	VoyageEmbeddingConfig	unchanged
sentence-transformers-mpnet	SagemakerEmbeddingConfig	unchanged (HF)
embed-multilingual-v3-prod	SagemakerEmbeddingConfig	unchanged*

* Names without cohere still use HF inputs; intentional — customer ticket used a cohere.* endpoint name.

Type

🐛 Bug Fix

Changes

• litellm/llms/sagemaker/embedding/cohere_transformation.py — new SagemakerCohereEmbeddingConfig
• litellm/llms/sagemaker/embedding/transformation.py — factory routes "cohere" endpoints to Cohere config
• litellm/llms/sagemaker/completion/handler.py — docstring only
• tests/test_litellm/llms/sagemaker/test_sagemaker_embedding_voyage.py — Cohere request/response + factory tests

[greptile-apps]

Greptile Summary

This PR fixes a 422 EmbedReqV2.inputs is of type string error on AWS Marketplace Cohere SageMaker endpoints by routing cohere-named endpoints to a new SagemakerCohereEmbeddingConfig that sends {"texts": [...], "input_type": "..."} instead of the HuggingFace TGI {"inputs": [...]} shape. It also includes a targeted extract-method refactor in CohereEmbeddingConfig to prevent double post_call logging on the SageMaker path, and a utils.py fix that restores non-OpenAI embedding params (like input_type) which were previously silently dropped by embedding_pre_process_non_default_params.

• SagemakerCohereEmbeddingConfig (new file): delegates request building to BedrockCohereEmbeddingConfig._transform_request (which correctly produces texts/input_type and strips the model field), calls CohereEmbeddingConfig._populate_embedding_response to avoid re-invoking post_call.
• CohereEmbeddingConfig._populate_embedding_response (extract-method): splits parse logic out of _transform_response; zero behavior change for existing Cohere direct and Bedrock callers.
• utils.py fallback loop: after map_openai_params, restores supported non-OpenAI params from the original passed_params dict; works correctly but implicitly relies on passed_params being mutated in-place by the pre-process step rather than reading directly from special_params.

Confidence Score: 5/5

The change is well-scoped and additive: new SageMaker Cohere routing does not affect existing Voyage or HF paths, and the Cohere v1 refactor is a pure extract-method with no logic change for existing callers.

All three code paths changed (new routing, extract-method, utils.py param restoration) are covered by targeted unit tests. The double-post_call fix is verified with an explicit assertion. The utils.py loop correctly gates on supported_params, OPENAI_EMBEDDING_PARAMS exclusion, and param not in optional_params, making unintended side-effects on other providers unlikely.

The utils.py fallback loop merits a second look to confirm it should read from special_params rather than the mutated passed_params dict.

Important Files Changed

Filename	Overview
litellm/llms/sagemaker/embedding/cohere_transformation.py	New config class routing Cohere SageMaker endpoints to the correct {"texts": [...]} payload format; correctly delegates to BedrockCohereEmbeddingConfig for request building and avoids double post_call logging.
litellm/utils.py	Adds a post-map loop to restore non-OpenAI embedding params (e.g. input_type) that are stripped by embedding_pre_process_non_default_params; fix works correctly but relies on an implicit side-effect mutation of passed_params.
litellm/llms/cohere/embed/v1_transformation.py	Extract-method refactor splits _transform_response into _populate_embedding_response to allow SageMaker to reuse parse logic without triggering a second post_call; correctly removes dead token-count loop in the process.
litellm/llms/sagemaker/embedding/transformation.py	Factory routing updated to dispatch Cohere-named endpoints to the new config; Voyage and HF paths unchanged.
tests/test_litellm/llms/sagemaker/test_sagemaker_embedding_voyage.py	New test coverage for Cohere request/response transform, factory routing, input_type pass-through, and no-double-post_call invariant; all mocked, no real network calls.
litellm/llms/sagemaker/completion/handler.py	Docstring-only change updating supported endpoint types list.

_{Reviews (3): Last reviewed commit: "fix(sagemaker): avoid double post_call o..." | Re-trigger Greptile}

[codecov]

Codecov Report

❌ Patch coverage is 90.32258% with 6 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
.../llms/sagemaker/embedding/cohere_transformation.py	86.48%	5 Missing ⚠️
litellm/llms/sagemaker/embedding/transformation.py	92.30%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

[milan-berri]

@greptileai check again

[milan-berri]

Addressed the 4/5 Greptile concern in da5aab78.

Bug: CohereEmbeddingConfig._transform_response (v1_transformation.py) calls logging_obj.post_call internally. The SageMaker embedding handler already calls post_call once before invoking the transform (handler.py:652), so every Cohere SageMaker embedding call fired callbacks, cost calculators, and log handlers twice. Voyage and HF SageMaker paths are unaffected because their transforms don't log.

Why I couldn't just remove the inner post_call: cohere/embed/handler.py and Bedrock Cohere both rely on _transform_response doing its own logging — removing it would silently break those paths.

Fix: Pure extract-method on CohereEmbeddingConfig:

• New _populate_embedding_response(response_json, model_response, model, encoding, input) contains the parsing/usage logic.
• _transform_response keeps its signature, still calls post_call, then delegates to the new helper — zero behavior change for existing callers.
• SagemakerCohereEmbeddingConfig.transform_embedding_response now calls _populate_embedding_response directly, so the SageMaker handler's single post_call is the only one.

Verification:

• New test test_transform_response_does_not_double_call_post_call asserts logging_obj.post_call.assert_not_called() on the SageMaker Cohere transform.
• Existing tests/test_litellm/llms/cohere/embed/test_v1_transformation.py (Cohere direct + Bedrock path coverage) still passes — proves the extract-method is behavior-preserving.
• tests/test_litellm/proxy/pass_through_endpoints/.../test_cohere_passthrough_logging_handler.py (mocks _transform_response) still passes.
• Full SageMaker + Cohere v1 + passthrough suites: 72/72 passing.

@greptileai check again

[mateo-berri]

LGTM; thanks!