diff --git a/litellm/constants.py b/litellm/constants.py index 717103450a2..c38a000a684 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1,6 +1,6 @@ import os import sys -from typing import List, Literal +from typing import List, Literal, Optional from litellm.litellm_core_utils.env_utils import get_env_int @@ -1318,6 +1318,14 @@ BATCH_STATUS_POLL_MAX_ATTEMPTS = int( HEALTH_CHECK_TIMEOUT_SECONDS = int( os.getenv("HEALTH_CHECK_TIMEOUT_SECONDS", 60) ) # 60 seconds +_background_health_check_max_tokens_env = os.getenv( + "BACKGROUND_HEALTH_CHECK_MAX_TOKENS" +) +BACKGROUND_HEALTH_CHECK_MAX_TOKENS: Optional[int] = ( + int(_background_health_check_max_tokens_env) + if _background_health_check_max_tokens_env + else None +) LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME = "litellm-internal-health-check" LITTELM_CLI_SERVICE_ACCOUNT_NAME = "litellm-cli" LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME = "litellm_internal_jobs" diff --git a/litellm/proxy/health_check.py b/litellm/proxy/health_check.py index 3e05ee3c484..e4705719fa3 100644 --- a/litellm/proxy/health_check.py +++ b/litellm/proxy/health_check.py @@ -11,7 +11,11 @@ from typing import List, Optional import litellm logger = logging.getLogger(__name__) -from litellm.constants import DEFAULT_HEALTH_CHECK_PROMPT, HEALTH_CHECK_TIMEOUT_SECONDS +from litellm.constants import ( + BACKGROUND_HEALTH_CHECK_MAX_TOKENS, + DEFAULT_HEALTH_CHECK_PROMPT, + HEALTH_CHECK_TIMEOUT_SECONDS, +) ILLEGAL_DISPLAY_PARAMS = [ "messages", @@ -281,6 +285,8 @@ def _update_litellm_params_for_health_check( _health_check_max_tokens = model_info.get("health_check_max_tokens", None) if _health_check_max_tokens is not None: litellm_params["max_tokens"] = _health_check_max_tokens + elif BACKGROUND_HEALTH_CHECK_MAX_TOKENS is not None: + litellm_params["max_tokens"] = BACKGROUND_HEALTH_CHECK_MAX_TOKENS elif "*" not in ( model_info.get("health_check_model") or litellm_params.get("model") or "" ): 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 e26f7fb9f20..bb125f764b8 100644 --- a/tests/test_litellm/proxy/test_health_check_max_tokens.py +++ b/tests/test_litellm/proxy/test_health_check_max_tokens.py @@ -1,7 +1,10 @@ +from unittest.mock import AsyncMock, MagicMock, patch + import pytest -from litellm.proxy.health_check import _update_litellm_params_for_health_check + from litellm.litellm_core_utils.health_check_helpers import HealthCheckHelpers -from unittest.mock import AsyncMock, patch, MagicMock +from litellm.proxy import health_check as hc_module +from litellm.proxy.health_check import _update_litellm_params_for_health_check @pytest.mark.asyncio @@ -73,3 +76,50 @@ async def test_ahealth_check_wildcard_models_respects_max_tokens(): litellm_logging_obj=MagicMock(), ) assert model_params["max_tokens"] == 3 + + +@pytest.mark.asyncio +async def test_background_health_check_max_tokens_env_var(monkeypatch): + """ + Test that BACKGROUND_HEALTH_CHECK_MAX_TOKENS env var is used as global default + for explicit (non-wildcard) models. + """ + monkeypatch.setattr(hc_module, "BACKGROUND_HEALTH_CHECK_MAX_TOKENS", 10) + + model_info = {} + litellm_params = {"model": "azure/gpt-4"} + + updated_params = _update_litellm_params_for_health_check(model_info, litellm_params) + + assert updated_params["max_tokens"] == 10 + + +@pytest.mark.asyncio +async def test_per_model_overrides_global_env_var(monkeypatch): + """ + Test that per-model health_check_max_tokens takes priority over + BACKGROUND_HEALTH_CHECK_MAX_TOKENS env var. + """ + monkeypatch.setattr(hc_module, "BACKGROUND_HEALTH_CHECK_MAX_TOKENS", 10) + + model_info = {"health_check_max_tokens": 5} + litellm_params = {"model": "azure/gpt-4"} + + updated_params = _update_litellm_params_for_health_check(model_info, litellm_params) + + assert updated_params["max_tokens"] == 5 + + +@pytest.mark.asyncio +async def test_global_env_var_applies_to_wildcard_models(monkeypatch): + """ + Test that BACKGROUND_HEALTH_CHECK_MAX_TOKENS env var also applies to wildcard models. + """ + monkeypatch.setattr(hc_module, "BACKGROUND_HEALTH_CHECK_MAX_TOKENS", 15) + + model_info = {} + litellm_params = {"model": "openai/*"} + + updated_params = _update_litellm_params_for_health_check(model_info, litellm_params) + + assert updated_params["max_tokens"] == 15