diff --git a/litellm/proxy/health_check.py b/litellm/proxy/health_check.py index 7d67750c78f..c59c61a9aa4 100644 --- a/litellm/proxy/health_check.py +++ b/litellm/proxy/health_check.py @@ -32,6 +32,30 @@ ILLEGAL_DISPLAY_PARAMS = [ MINIMAL_DISPLAY_PARAMS = ["model", "mode_error"] +# Modes whose health-check probe is a chat-style completion call and +# therefore accept `max_tokens`. Other modes (embedding, image_generation, +# audio_*, rerank, video_generation, ocr, search, moderation, ...) hit +# 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"}) + + +def _should_inject_health_check_max_tokens(model_info: dict) -> 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` + 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 + def _get_process_rss_mb() -> Optional[float]: """ @@ -362,14 +386,22 @@ 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` when the deployment is a chat-style mode + (`chat`, `completion`, `responses`) or the operator explicitly opts in + via `model_info.health_check_supports_max_tokens`. Non-chat endpoints + (image, embedding, audio_*, rerank, video, ocr, search, moderation, ...) + 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 """ 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 _should_inject_health_check_max_tokens(model_info): + _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..292ac771369 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,129 @@ 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_health_check_max_tokens_value_is_ignored_for_non_chat_modes(): + """A configured `health_check_max_tokens` *value* (the int that controls + how many tokens to inject) is still skipped when the mode is outside the + allow-list — the inject decision runs before value resolution, so the + value never reaches `_resolve_health_check_max_tokens`. Note this is + distinct from `health_check_supports_max_tokens` (the bool that toggles + injection on/off per deployment).""" + 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 + + +# --------------------------------------------------------------------------- +# Allow-list behavior: only chat-style modes (chat / completion / responses) +# receive max_tokens. Every other mode is skipped by default. +# +# Per-deployment override via `health_check_supports_max_tokens` lets the +# operator force injection on (e.g. a non-listed but max_tokens-capable +# endpoint where they want to bound probe token usage) or off (e.g. a +# chat-style provider with a strict schema that rejects unknown fields). +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("mode", ["chat", "completion", "responses"]) +def test_chat_style_modes_inject_max_tokens(mode): + updated = _update_litellm_params_for_health_check( + {"mode": mode}, {"model": f"openai/dummy-{mode}"} + ) + + assert updated["max_tokens"] == 5 + + +@pytest.mark.parametrize( + "mode", + [ + "embedding", + "image_generation", + "image_edit", + "audio_speech", + "audio_transcription", + "rerank", + "video_generation", + "ocr", + "search", + "moderation", + ], +) +def test_non_chat_modes_skip_max_tokens(mode): + updated = _update_litellm_params_for_health_check( + {"mode": mode}, {"model": f"openai/dummy-{mode}"} + ) + + assert "max_tokens" not in updated + + +def test_explicit_override_true_forces_injection_outside_allowlist(): + """Operator opts a non-listed deployment in to bound probe token usage.""" + model_info = { + "mode": "image_generation", + "health_check_supports_max_tokens": True, + } + litellm_params = {"model": "openai/some-future-image-model"} + + updated = _update_litellm_params_for_health_check(model_info, litellm_params) + + assert updated["max_tokens"] == 5 + + +def test_explicit_override_false_suppresses_injection_inside_allowlist(): + """Operator opts a chat-style deployment out (strict-schema provider).""" + model_info = {"mode": "chat", "health_check_supports_max_tokens": False} + litellm_params = {"model": "openai/strict-schema-chat"} + + updated = _update_litellm_params_for_health_check(model_info, litellm_params) + + assert "max_tokens" not in updated