From 1dee0d075046630fa5b3e785de2679386520d0a3 Mon Sep 17 00:00:00 2001 From: xianren Date: Tue, 28 Apr 2026 12:23:49 +0800 Subject: [PATCH] fix(health-check): keep max_tokens for completion mode, add unit tests Greptile P1: litellm.atext_completion (completion mode) accepts max_tokens, so unconditionally stripping it via _filter_model_params silently dropped the BACKGROUND_HEALTH_CHECK_MAX_TOKENS cost cap for completion-mode deployments. Make _filter_model_params mode-aware via a keep_max_tokens flag and use it in the completion-mode lambda. Strict non-chat handlers (image_generation, embedding, etc.) still strip max_tokens as before. Adds tests/litellm/litellm_core_utils/test_health_check_utils.py covering both code paths and addressing the codecov/patch failure. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../health_check_helpers.py | 4 +- .../litellm_core_utils/health_check_utils.py | 42 ++++++----- .../test_health_check_utils.py | 71 +++++++++++++++++++ 3 files changed, 100 insertions(+), 17 deletions(-) create mode 100644 tests/litellm/litellm_core_utils/test_health_check_utils.py diff --git a/litellm/litellm_core_utils/health_check_helpers.py b/litellm/litellm_core_utils/health_check_helpers.py index 9e972f1910b..626705c7267 100644 --- a/litellm/litellm_core_utils/health_check_helpers.py +++ b/litellm/litellm_core_utils/health_check_helpers.py @@ -157,7 +157,9 @@ class HealthCheckHelpers: **model_params, ), "completion": lambda: litellm.atext_completion( - **_filter_model_params(model_params=model_params), + **_filter_model_params( + model_params=model_params, keep_max_tokens=True + ), prompt=prompt or "test", ), "embedding": lambda: litellm.aembedding( diff --git a/litellm/litellm_core_utils/health_check_utils.py b/litellm/litellm_core_utils/health_check_utils.py index 874d260eba7..b0e71896728 100644 --- a/litellm/litellm_core_utils/health_check_utils.py +++ b/litellm/litellm_core_utils/health_check_utils.py @@ -2,29 +2,39 @@ Utils used for litellm.ahealth_check() """ -#: litellm params that ``_update_litellm_params_for_health_check`` injects for -#: chat-completion health checks but that are invalid (or rejected) on every -#: other handler in :class:`HealthCheckHelpers.get_mode_handlers` — image and -#: video generation, embeddings, audio speech / transcription, rerank, ocr, -#: responses, batch, etc. ``messages`` is always chat-only; ``max_tokens`` is -#: chat/completion-only and is rejected by strict providers (e.g. OpenAI's -#: image-generation endpoints return 400 ``Unknown parameter: 'max_tokens'``). -_NON_CHAT_HEALTH_CHECK_STRIP_KEYS = {"messages", "max_tokens"} +#: ``messages`` is injected by ``_update_litellm_params_for_health_check`` for +#: chat-style health checks; every non-chat handler in +#: :class:`HealthCheckHelpers.get_mode_handlers` rejects it, and +#: ``litellm.atext_completion`` takes ``prompt`` instead. +_COMPLETION_HEALTH_CHECK_STRIP_KEYS = {"messages"} + +#: Strict non-chat handlers (OpenAI image generation, embeddings, rerank, etc.) +#: also reject ``max_tokens`` — e.g. ``dall-e-3`` returns +#: ``400 Unknown parameter: 'max_tokens'``. ``atext_completion`` does accept +#: ``max_tokens``, so the ``completion`` mode keeps it (preserves the +#: ``BACKGROUND_HEALTH_CHECK_MAX_TOKENS`` cost cap). +_NON_CHAT_HEALTH_CHECK_STRIP_KEYS = _COMPLETION_HEALTH_CHECK_STRIP_KEYS | { + "max_tokens" +} -def _filter_model_params(model_params: dict) -> dict: +def _filter_model_params( + model_params: dict, *, keep_max_tokens: bool = False +) -> dict: """Strip chat-only params before invoking a non-chat health check handler. ``litellm.acompletion`` is the only mode handler that consumes ``model_params`` unfiltered; every other handler routes through this - helper, so removing chat-completion-only keys here keeps strict providers - (OpenAI image generation, etc.) from rejecting the request. + helper. ``litellm.atext_completion`` (the ``completion`` mode) accepts + ``max_tokens`` and should pass ``keep_max_tokens=True`` to preserve the + cost-control cap. """ - return { - k: v - for k, v in model_params.items() - if k not in _NON_CHAT_HEALTH_CHECK_STRIP_KEYS - } + strip_keys = ( + _COMPLETION_HEALTH_CHECK_STRIP_KEYS + if keep_max_tokens + else _NON_CHAT_HEALTH_CHECK_STRIP_KEYS + ) + return {k: v for k, v in model_params.items() if k not in strip_keys} def _create_health_check_response(response_headers: dict) -> dict: diff --git a/tests/litellm/litellm_core_utils/test_health_check_utils.py b/tests/litellm/litellm_core_utils/test_health_check_utils.py new file mode 100644 index 00000000000..944f4908523 --- /dev/null +++ b/tests/litellm/litellm_core_utils/test_health_check_utils.py @@ -0,0 +1,71 @@ +""" +Tests for ``litellm.litellm_core_utils.health_check_utils._filter_model_params``. + +Regression coverage for #26406 / #26604: chat-only params (``messages``, +``max_tokens``) must be stripped before non-chat health check handlers, but +``max_tokens`` must be preserved for the ``completion`` mode handler +(``litellm.atext_completion``) so the ``BACKGROUND_HEALTH_CHECK_MAX_TOKENS`` +cost cap still applies. +""" + +from litellm.litellm_core_utils.health_check_utils import ( + _COMPLETION_HEALTH_CHECK_STRIP_KEYS, + _NON_CHAT_HEALTH_CHECK_STRIP_KEYS, + _filter_model_params, +) + + +def _sample_params() -> dict: + return { + "model": "openai/dall-e-3", + "messages": [{"role": "user", "content": "test"}], + "max_tokens": 5, + "api_key": "sk-fake", + } + + +def test_filter_strips_messages_and_max_tokens_for_non_chat_handlers(): + """Default behavior: strip both ``messages`` and ``max_tokens``. + + Reproduces the original failure (OpenAI image-generation 400 on + ``max_tokens``) — these keys must not reach non-chat handlers. + """ + filtered = _filter_model_params(model_params=_sample_params()) + + assert "messages" not in filtered + assert "max_tokens" not in filtered + assert filtered == {"model": "openai/dall-e-3", "api_key": "sk-fake"} + + +def test_filter_keeps_max_tokens_for_completion_mode(): + """``completion`` mode passes ``keep_max_tokens=True``. + + ``litellm.atext_completion`` accepts ``max_tokens``; stripping it would + silently remove the ``BACKGROUND_HEALTH_CHECK_MAX_TOKENS`` cost cap. + """ + filtered = _filter_model_params( + model_params=_sample_params(), keep_max_tokens=True + ) + + assert "messages" not in filtered + assert filtered.get("max_tokens") == 5 + assert filtered["model"] == "openai/dall-e-3" + assert filtered["api_key"] == "sk-fake" + + +def test_filter_does_not_mutate_input(): + params = _sample_params() + snapshot = dict(params) + + _filter_model_params(model_params=params) + _filter_model_params(model_params=params, keep_max_tokens=True) + + assert params == snapshot + + +def test_strip_key_sets_are_consistent(): + """``_NON_CHAT`` is a strict superset of ``_COMPLETION`` — only + ``max_tokens`` differs between the two modes.""" + assert _COMPLETION_HEALTH_CHECK_STRIP_KEYS == {"messages"} + assert _NON_CHAT_HEALTH_CHECK_STRIP_KEYS == {"messages", "max_tokens"} + assert _COMPLETION_HEALTH_CHECK_STRIP_KEYS < _NON_CHAT_HEALTH_CHECK_STRIP_KEYS