From c51ba3429400e0c2fa953b7ecab0fa92f915e8ab Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 17 Jun 2026 14:34:09 -0700 Subject: [PATCH] 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 | 17 +-- litellm/proxy/health_check.py | 59 ++++++-- .../proxy/test_health_check_max_tokens.py | 126 ++++++++++++++++++ 3 files changed, 176 insertions(+), 26 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index 80176cc8b16..63c5798e70a 100644 --- a/litellm/main.py +++ b/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, ): diff --git a/litellm/proxy/health_check.py b/litellm/proxy/health_check.py index 4a28143e617..be51234e7bc 100644 --- a/litellm/proxy/health_check.py +++ b/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 diff --git a/tests/test_litellm/proxy/test_health_check_max_tokens.py b/tests/test_litellm/proxy/test_health_check_max_tokens.py index 5cb7cdacc60..a1c0b5ee450 100644 --- a/tests/test_litellm/proxy/test_health_check_max_tokens.py +++ b/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