diff --git a/litellm/proxy/health_check.py b/litellm/proxy/health_check.py index 7d67750c78f..723d251457f 100644 --- a/litellm/proxy/health_check.py +++ b/litellm/proxy/health_check.py @@ -362,14 +362,21 @@ def _update_litellm_params_for_health_check( Update the litellm params for health check. - gets a short `messages` param for health check + - adds a bounded `max_tokens`, except when + `model_info.mode == "image_generation"` — OpenAI + `/v1/images/generations` rejects 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 """ litellm_params["messages"] = _get_random_llm_message() - _resolved_max_tokens = _resolve_health_check_max_tokens(model_info, litellm_params) - if _resolved_max_tokens is not None: - litellm_params["max_tokens"] = _resolved_max_tokens + if model_info.get("mode", None) != "image_generation": + _resolved_max_tokens = _resolve_health_check_max_tokens( + model_info, litellm_params + ) + if _resolved_max_tokens is not None: + litellm_params["max_tokens"] = _resolved_max_tokens _health_check_model = model_info.get("health_check_model", None) if _health_check_model is not None: 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 09211b72c3e..ffb210a149a 100644 --- a/tests/test_litellm/proxy/test_health_check_max_tokens.py +++ b/tests/test_litellm/proxy/test_health_check_max_tokens.py @@ -225,3 +225,59 @@ def test_wildcard_ignores_reasoning_split_model_info(monkeypatch): litellm_params = {"model": "openai/*"} assert _resolve_health_check_max_tokens(model_info, litellm_params) is None + + +# --------------------------------------------------------------------------- +# image_generation must not receive max_tokens. +# +# _update_litellm_params_for_health_check injected `max_tokens` for every +# deployment. For `mode: image_generation` that leaked into OpenAI +# `/v1/images/generations`, which strictly rejects unknown fields with +# `400 "Unknown parameter: 'max_tokens'"`, marking dall-e-* and +# gpt-image-1 as permanently unhealthy even though their actual image +# calls succeed. `messages` still gets injected (downstream +# `_filter_model_params` already strips it for non-chat handlers). +# --------------------------------------------------------------------------- + + +def test_image_generation_mode_skips_max_tokens(): + """image_generation must not receive max_tokens.""" + model_info = {"mode": "image_generation"} + litellm_params = {"model": "openai/dall-e-3", "api_key": "sk-test"} + + updated = _update_litellm_params_for_health_check(model_info, litellm_params) + + assert "max_tokens" not in updated + # connection-level params must still pass through unchanged + assert updated["api_key"] == "sk-test" + + +def test_image_generation_ignores_explicit_health_check_max_tokens(): + """Even an explicit `health_check_max_tokens` override must not be + forwarded to image endpoints — the field is invalid there.""" + model_info = {"mode": "image_generation", "health_check_max_tokens": 50} + litellm_params = {"model": "openai/dall-e-3"} + + updated = _update_litellm_params_for_health_check(model_info, litellm_params) + + assert "max_tokens" not in updated + + +def test_chat_mode_still_injects_max_tokens(): + """Regression guard: the chat-style probe payload is unchanged.""" + model_info = {"mode": "chat"} + litellm_params = {"model": "gpt-4"} + + updated = _update_litellm_params_for_health_check(model_info, litellm_params) + + assert updated["max_tokens"] == 5 + + +def test_no_mode_still_injects_max_tokens(): + """Regression guard: model_info without `mode` keeps the legacy path.""" + model_info: dict = {} + litellm_params = {"model": "gpt-4"} + + updated = _update_litellm_params_for_health_check(model_info, litellm_params) + + assert updated["max_tokens"] == 5