fix(health): correct bedrock embedding health checks (#30583)

* fix(health): correct bedrock embedding health checks

Health checks for Bedrock embedding deployments failed in two ways. A
deployment configured without an explicit model_info.mode was probed as
chat, so max_tokens was injected and Bedrock embeddings rejected it with
400 "extraneous key [max_tokens]". Separately, stripping the bedrock/
routing prefix dropped the provider, so a cross-region inference-profile
id like us.cohere.embed-v4:0 failed downstream with "LLM Provider NOT
provided".

Resolve the deployment mode from the model cost map (which understands
the bedrock/ and us./eu./apac. prefixes) before deciding whether to
inject max_tokens, and pin custom_llm_provider to bedrock when stripping
the prefix so the bare model id still resolves. ahealth_check now accepts
any string mode so the resolved embedding mode routes the probe to the
embedding handler.

* fix(health): preserve explicit custom_llm_provider on bedrock probe

The bedrock prefix-strip pinned custom_llm_provider to bedrock
unconditionally, so a deployment that set custom_llm_provider:
bedrock_converse had it overwritten at health-check time and the probe
hit the Invoke endpoint instead of Converse, a different request format
that can report a spurious failure. Only fill in bedrock when the
deployment left the provider blank, which still resolves bare
cross-region ids like us.cohere.embed-v4:0 while leaving an explicit
provider untouched.

* test(health): assert resolved mode reaches the ahealth_check probe

The existing tests check _resolve_health_check_mode and the params builder
in isolation, but nothing verified that _run_model_health_check actually
threads the resolved mode into litellm.ahealth_check. Without that, a
refactor that probed with model_info.get("mode") again would reintroduce
the chat fallback for embedding deployments while every test stayed green.
This drives _run_model_health_check with a bedrock embedding deployment and
asserts the probe is called with mode=embedding and the embedding params.

* fix(health): resolve probe mode once for reasoning_effort and audio_speech

The reasoning_effort and audio_speech branches read model_info.mode
directly, so an embedding deployment declared without an explicit mode (the
case this PR targets) was still treated as chat-like: a configured
health_check_reasoning_effort got injected into the embedding probe, which
embeddings reject as an unknown field, and an auto-detected audio_speech
deployment never had its voice set. Resolve the effective mode once from the
cost map and reuse it for the max_tokens, reasoning_effort, and audio_speech
decisions so they all agree with the mode threaded into ahealth_check.

litellm/main.py

@@ -7434,22 +7434,7 @@ def speech(
 
 async def ahealth_check(
     model_params: dict,
-    mode: Optional[
-        Literal[
-            "chat",
-            "completion",
-            "embedding",
-            "audio_speech",
-            "audio_transcription",
-            "image_generation",
-            "video_generation",
-            "batch",
-            "rerank",
-            "realtime",
-            "responses",
-            "ocr",
-        ]
-    ] = "chat",
+    mode: str | None = "chat",
     prompt: Optional[str] = None,
     input: Optional[List] = None,
 ):

litellm/proxy/health_check.py

@@ -6,6 +6,7 @@ import random
 import sys
 import threading
 import time
+from collections.abc import Mapping
 from typing import List, Optional
 
 import litellm
@@ -42,23 +43,50 @@ MINIMAL_DISPLAY_PARAMS = ["model", "mode_error"]
 # endpoints that reject unknown fields with 400 "Unknown parameter:
 # 'max_tokens'". Allow-list so new modes are safe by default.
 # Per-deployment override: `model_info.health_check_supports_max_tokens`.
-_MAX_TOKEN_SUPPORT_MODES: frozenset = frozenset({"chat", "completion", "responses"})
+_MAX_TOKEN_SUPPORT_MODES: frozenset[str] = frozenset(
+    {"chat", "completion", "responses"}
+)
 
 
-def _should_inject_health_check_max_tokens(model_info: dict) -> bool:
+def _resolve_health_check_mode(
+    model_info: Mapping[str, object], litellm_params: Mapping[str, object]
+) -> str | None:
+    """
+    Effective mode for a deployment's health-check probe.
+
+    Prefers operator-set `model_info.mode`; otherwise resolves it from the model
+    cost map, which understands `bedrock/` and cross-region inference-profile
+    prefixes (`us.`, `eu.`, `apac.`). Without this, non-chat Bedrock deployments
+    (e.g. embeddings) are probed as chat, so `max_tokens` is injected and the
+    request 400s on "extraneous key [max_tokens]".
+    """
+    explicit_mode = model_info.get("mode")
+    if isinstance(explicit_mode, str):
+        return explicit_mode
+    model = litellm_params.get("model")
+    if not isinstance(model, str):
+        return None
+    try:
+        return litellm.get_model_info(model=model).get("mode")
+    except Exception:
+        return None
+
+
+def _should_inject_health_check_max_tokens(
+    model_info: Mapping[str, object], mode: str | None
+) -> bool:
     """
     Whether the health-check probe should include `max_tokens`.
 
     Order:
       1. `model_info.health_check_supports_max_tokens` (operator override).
-      2. `_MAX_TOKEN_SUPPORT_MODES`. Missing `mode` is treated as `chat`
+      2. `_MAX_TOKEN_SUPPORT_MODES`. An unresolvable mode is treated as `chat`
          for backward compatibility.
     """
     explicit = model_info.get("health_check_supports_max_tokens")
     if explicit is not None:
         return bool(explicit)
-    mode = model_info.get("mode") or "chat"
-    return mode in _MAX_TOKEN_SUPPORT_MODES
+    return (mode or "chat") in _MAX_TOKEN_SUPPORT_MODES
 
 
 # Health-check modes that forward `reasoning_effort` to the provider (chat-style calls).
@@ -165,7 +193,9 @@ async def run_with_timeout(task, timeout):
 async def _run_model_health_check(model: dict):
     litellm_params = model["litellm_params"]
     model_info = model.get("model_info", {})
-    mode = model_info.get("mode", None)
+    mode = _resolve_health_check_mode(
+        model_info, litellm_params  # any-ok: untyped router config dict
+    )
     litellm_params = _update_litellm_params_for_health_check(model_info, litellm_params)
     timeout = model_info.get("health_check_timeout") or HEALTH_CHECK_TIMEOUT_SECONDS
 
@@ -421,10 +451,15 @@ def _update_litellm_params_for_health_check(
       reject unknown fields with 400 "Unknown parameter: 'max_tokens'".
     - updates the `model` param with the `health_check_model` if it exists Doc: https://docs.litellm.ai/docs/proxy/health#wildcard-routes
     - updates the `voice` param with the `health_check_voice` for `audio_speech` mode if it exists Doc: https://docs.litellm.ai/docs/proxy/health#text-to-speech-models
-    - for Bedrock models with region routing (bedrock/region/model), strips the litellm routing prefix but preserves the model ID
+    - for Bedrock models with region routing (bedrock/region/model), strips the litellm routing prefix but preserves the model ID, and pins `custom_llm_provider` to `bedrock` (only when the deployment hasn't already set one, so an explicit `bedrock_converse` survives) so the bare model id still resolves to the provider (e.g. cross-region ids like `us.cohere.embed-v4:0`)
     """
+    mode = _resolve_health_check_mode(
+        model_info, litellm_params  # any-ok: untyped router config dict
+    )
     litellm_params["messages"] = _get_random_llm_message()
-    if _should_inject_health_check_max_tokens(model_info):
+    if _should_inject_health_check_max_tokens(
+        model_info, mode  # any-ok: untyped router config dict
+    ):
         _resolved_max_tokens = _resolve_health_check_max_tokens(
             model_info, litellm_params
         )
@@ -432,7 +467,7 @@ def _update_litellm_params_for_health_check(
             litellm_params["max_tokens"] = _resolved_max_tokens
 
     # Per-model reasoning effort for health checks only (e.g. reasoning_effort=none).
-    if model_info.get("mode", None) in _HEALTH_CHECK_MODES_SUPPORTING_REASONING_EFFORT:
+    if mode in _HEALTH_CHECK_MODES_SUPPORTING_REASONING_EFFORT:
         _hc_reasoning_effort = model_info.get("health_check_reasoning_effort", None)
         if _hc_reasoning_effort is not None:
             litellm_params["reasoning_effort"] = _hc_reasoning_effort
@@ -440,7 +475,7 @@ def _update_litellm_params_for_health_check(
     _health_check_model = model_info.get("health_check_model", None)
     if _health_check_model is not None:
         litellm_params["model"] = _health_check_model
-    if model_info.get("mode", None) == "audio_speech":
+    if mode == "audio_speech":
         litellm_params["voice"] = model_info.get("health_check_voice", "alloy")
 
     # Handle Bedrock region routing format: bedrock/region/model
@@ -477,6 +512,10 @@ def _update_litellm_params_for_health_check(
 
         model = "/".join(filtered_parts)
         litellm_params["model"] = model
+        if not litellm_params.get("custom_llm_provider"):  # any-ok: untyped router dict
+            litellm_params["custom_llm_provider"] = (  # any-ok: untyped router dict
+                "bedrock"
+            )
 
     return litellm_params
 

tests/test_litellm/proxy/test_health_check_max_tokens.py

@@ -6,6 +6,7 @@ from litellm.litellm_core_utils.health_check_helpers import HealthCheckHelpers
 from litellm.proxy import health_check as hc_module
 from litellm.proxy.health_check import (
     _resolve_health_check_max_tokens,
+    _resolve_health_check_mode,
     _update_litellm_params_for_health_check,
 )
 
@@ -391,3 +392,128 @@ def test_update_litellm_params_health_check_reasoning_effort():
         model_info, {"model": "openai/gpt-4o", "api_key": "x"}
     )
     assert "reasoning_effort" not in out
+
+
+# ---------------------------------------------------------------------------
+# Bedrock embedding deployments declared without an explicit `model_info.mode`.
+#
+# The health-check builder used to treat a missing mode as `chat`, so it
+# injected `max_tokens` into the embedding probe. Bedrock embeddings reject it
+# with 400 "extraneous key [max_tokens]". It also stripped the `bedrock/`
+# routing prefix without pinning the provider, so a cross-region id like
+# `us.cohere.embed-v4:0` failed downstream with "LLM Provider NOT provided".
+# Mode is now resolved from the model cost map (which understands `bedrock/`
+# and `us.`/`eu.`/`apac.` prefixes) and the provider is pinned to `bedrock`.
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    "deployment_model, expected_request_model",
+    [
+        ("bedrock/amazon.titan-embed-text-v2:0", "amazon.titan-embed-text-v2:0"),
+        ("bedrock/us.cohere.embed-v4:0", "us.cohere.embed-v4:0"),
+    ],
+)
+def test_bedrock_embedding_without_explicit_mode_skips_max_tokens(
+    deployment_model, expected_request_model
+):
+    """Embedding mode auto-detected from model cost map -> no max_tokens, provider pinned."""
+    assert _resolve_health_check_mode({}, {"model": deployment_model}) == "embedding"
+
+    updated = _update_litellm_params_for_health_check({}, {"model": deployment_model})
+
+    assert "max_tokens" not in updated
+    assert updated["custom_llm_provider"] == "bedrock"
+    assert updated["model"] == expected_request_model
+
+
+def test_resolve_health_check_mode_prefers_explicit_model_info_mode():
+    """An operator-set mode wins over model-cost lookup."""
+    assert (
+        _resolve_health_check_mode(
+            {"mode": "chat"}, {"model": "bedrock/amazon.titan-embed-text-v2:0"}
+        )
+        == "chat"
+    )
+
+
+def test_resolve_health_check_mode_unknown_model_returns_none():
+    assert (
+        _resolve_health_check_mode({}, {"model": "bedrock/not-a-real-model-xyz"})
+        is None
+    )
+    assert _resolve_health_check_mode({}, {}) is None
+
+
+def test_bedrock_chat_without_mode_still_injects_max_tokens_and_pins_provider():
+    """Regression guard: chat-style Bedrock deployments keep max_tokens and get the provider pin."""
+    updated = _update_litellm_params_for_health_check(
+        {}, {"model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0"}
+    )
+
+    assert updated["max_tokens"] == 5
+    assert updated["custom_llm_provider"] == "bedrock"
+    assert updated["model"] == "us.anthropic.claude-haiku-4-5-20251001-v1:0"
+
+
+def test_bedrock_prefix_strip_preserves_explicit_custom_llm_provider():
+    """An operator-set provider (e.g. bedrock_converse) must survive the prefix strip.
+
+    The pin only fills in a provider when the deployment left it blank; it must
+    not clobber a more specific one, otherwise a converse deployment would be
+    probed against the Invoke endpoint and report a spurious failure.
+    """
+    updated = _update_litellm_params_for_health_check(
+        {},
+        {
+            "model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0",
+            "custom_llm_provider": "bedrock_converse",
+        },
+    )
+
+    assert updated["custom_llm_provider"] == "bedrock_converse"
+    assert updated["model"] == "us.anthropic.claude-haiku-4-5-20251001-v1:0"
+
+
+@pytest.mark.asyncio
+async def test_run_model_health_check_threads_resolved_mode_to_ahealth_check():
+    """The resolved mode must reach `ahealth_check`, not just the params builder.
+
+    A Bedrock embedding deployment declared without an explicit `model_info.mode`
+    has to be probed with `mode="embedding"` so the call routes to the embedding
+    handler; if the resolution were dropped it would fall back to `chat`. This
+    also guards that the embedding params (no `max_tokens`, provider pinned) are
+    the ones actually handed to the probe.
+    """
+    fake_ahealth_check = AsyncMock(return_value={})
+    model = {
+        "litellm_params": {"model": "bedrock/amazon.titan-embed-text-v2:0"},
+        "model_info": {},
+    }
+
+    with patch.object(hc_module.litellm, "ahealth_check", fake_ahealth_check):
+        await hc_module._run_model_health_check(model)
+
+    assert fake_ahealth_check.call_args.kwargs["mode"] == "embedding"
+    probed_params = fake_ahealth_check.call_args.args[0]
+    assert "max_tokens" not in probed_params
+    assert probed_params["custom_llm_provider"] == "bedrock"
+    assert probed_params["model"] == "amazon.titan-embed-text-v2:0"
+
+
+def test_autodetected_embedding_skips_reasoning_effort():
+    """reasoning_effort must not leak into an embedding probe whose mode is auto-detected.
+
+    Same bug class as the max_tokens fix: with no explicit `model_info.mode`, the
+    reasoning-effort gate used to read the raw (missing) mode and treat it as
+    chat-like, so a configured `health_check_reasoning_effort` was injected into a
+    Bedrock embedding probe, which embeddings reject as an unknown field. The mode
+    is now resolved from the cost map, so embeddings are excluded.
+    """
+    updated = _update_litellm_params_for_health_check(
+        {"health_check_reasoning_effort": "low"},
+        {"model": "bedrock/amazon.titan-embed-text-v2:0"},
+    )
+
+    assert "reasoning_effort" not in updated
+    assert "max_tokens" not in updated