diff --git a/litellm/litellm_core_utils/health_check_helpers.py b/litellm/litellm_core_utils/health_check_helpers.py index 0e03d513b83..ae4fff6dd75 100644 --- a/litellm/litellm_core_utils/health_check_helpers.py +++ b/litellm/litellm_core_utils/health_check_helpers.py @@ -35,6 +35,10 @@ IMAGE_EDIT_HEALTH_CHECK_PROMPT: Final = ( ANTHROPIC_MESSAGES_HEALTH_CHECK_MAX_TOKENS: Final = 16 WILDCARD_HEALTH_CHECK_SERVED_MODEL_COUNT: Final = 2 WILDCARD_HEALTH_CHECK_PROVIDER_MODEL_COUNT: Final = 3 +WILDCARD_HEALTH_CHECK_MAX_TOKENS: Final = 16 # GPT-5 models require max_output_tokens >= 16 +_WILDCARD_PROBE_TOKEN_CAP_PARAM: Final = MappingProxyType( + {"completion": "max_tokens", "responses": "max_output_tokens"} +) BEDROCK_ANY_MODEL_ROUTES: Final = frozenset({"converse", "invoke"}) EXTRACTION_HEALTH_CHECK_FIELDS: Final[Mapping[str, Mapping[str, str]]] = MappingProxyType( {"reachable": MappingProxyType({"type": "text", "instructions": "Any word from the text"})} @@ -187,6 +191,20 @@ def _wildcard_served_models_of_mode(wildcard: str, custom_llm_provider: str, han ) +def _with_probe_token_cap(model_params: Mapping[str, object], handler_mode: str) -> dict[str, object]: + cap_param: Final = _WILDCARD_PROBE_TOKEN_CAP_PARAM.get(handler_mode) + if cap_param is None: + return dict(model_params) + return { + **model_params, + cap_param: model_params.get(cap_param, model_params.get("max_tokens", WILDCARD_HEALTH_CHECK_MAX_TOKENS)), + } + + +def _optional_text(value: object) -> str | None: + return None if value is None else str(value) + + async def _first_answer(probes: Sequence[Callable[[], Awaitable[object]]]) -> object: try: return await probes[0]() @@ -221,7 +239,7 @@ class HealthCheckHelpers: handler := HealthCheckHelpers.get_mode_handlers( model=served_model.removeprefix(provider_prefix), custom_llm_provider=custom_llm_provider, - model_params={**model_params, "model": served_model}, + model_params={**_with_probe_token_cap(model_params, handler_mode), "model": served_model}, prompt=prompt, input=None if input is None else list(input), ).get(handler_mode) @@ -241,7 +259,7 @@ class HealthCheckHelpers: model_params["model"] = probe_models[0] model_params["litellm_logging_obj"] = litellm_logging_obj model_params["fallbacks"] = list(probe_models[1:]) or None - model_params["max_tokens"] = model_params.get("max_tokens", 16) # GPT-5 models require max_output_tokens >= 16 + model_params["max_tokens"] = model_params.get("max_tokens", WILDCARD_HEALTH_CHECK_MAX_TOKENS) await acompletion(**model_params) return {} @@ -356,7 +374,6 @@ class HealthCheckHelpers: from litellm.litellm_core_utils.health_check_utils import ( DECISIONS_CALL_PARAMS, EXTRACTION_CALL_PARAMS, - OPTIONAL_STR, _filter_model_params, ) from litellm.realtime_api.main import realtime_health_check @@ -407,9 +424,9 @@ class HealthCheckHelpers: "realtime": lambda: realtime_health_check( model=model, custom_llm_provider=custom_llm_provider, - api_base=OPTIONAL_STR.validate_python(model_params.get("api_base")), - api_key=OPTIONAL_STR.validate_python(model_params.get("api_key")), - api_version=OPTIONAL_STR.validate_python(model_params.get("api_version")), + api_base=_optional_text(model_params.get("api_base")), + api_key=_optional_text(model_params.get("api_key")), + api_version=_optional_text(model_params.get("api_version")), model_params=model_params, ), "batch": lambda: HealthCheckHelpers._batch_health_check( diff --git a/litellm/main.py b/litellm/main.py index 04b5b8c8a8c..a032edbf3f8 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -104,7 +104,10 @@ from litellm.litellm_core_utils.get_llm_provider_logic import ( from litellm.litellm_core_utils.get_provider_specific_headers import ( ProviderSpecificHeaderUtils, ) -from litellm.litellm_core_utils.health_check_utils import _filter_model_params +from litellm.litellm_core_utils.health_check_utils import ( + _filter_model_params, + create_health_check_response, +) from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, get_hidden_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.mock_functions import ( @@ -172,6 +175,7 @@ from litellm.types.utils import ( ProviderRoutingHints, RawRequestTypedDict, StreamingChoices, + is_decisions_model_mode, ) from litellm.types.workload_identity import ANTHROPIC_WIF_KWARGS_KEYS, OPENAI_WIF_KWARGS_KEYS from litellm.utils import ( @@ -9388,7 +9392,7 @@ async def ahealth_check( model_params: dict, mode: str | None = None, prompt: str | None = None, - input: list[str] | None = None, + input: list | None = None, ): """ Support health checks for different providers. Return remaining rate limit, etc. @@ -9401,12 +9405,7 @@ async def ahealth_check( } """ from litellm.litellm_core_utils.cached_imports import get_litellm_logging_class - from litellm.litellm_core_utils.health_check_helpers import ( - HealthCheckHelpers, - default_health_check_mode, - health_check_handler_mode, - health_check_response, - ) + from litellm.litellm_core_utils.health_check_helpers import HealthCheckHelpers, default_health_check_mode from litellm.litellm_core_utils.health_check_utils import OPTIONAL_STR # Use cached import helper to lazy-load Logging class (only loads when function is called) @@ -9471,9 +9470,13 @@ async def ahealth_check( input=input, ) - handler_mode: Final = health_check_handler_mode(mode) + handler_mode: Final = "decisions" if is_decisions_model_mode(mode) else mode if handler_mode in mode_handlers: - return health_check_response(await mode_handlers[handler_mode]()) + _response: Final = await mode_handlers[handler_mode]() + _response_headers: Final = cast( # cast-ok: provider headers are stored as a string-keyed mapping + Mapping[str, object], (get_hidden_params(_response) or {}).get("headers", {}) or {} + ) + return create_health_check_response(_response_headers) else: raise Exception(f"Mode {mode} not supported. See modes here: https://docs.litellm.ai/docs/proxy/health") except Exception as e: diff --git a/tests/unit/litellm_core_utils/test_health_check_helpers.py b/tests/unit/litellm_core_utils/test_health_check_helpers.py index 00ce0d9ec70..f1c7e6e3bf1 100644 --- a/tests/unit/litellm_core_utils/test_health_check_helpers.py +++ b/tests/unit/litellm_core_utils/test_health_check_helpers.py @@ -2161,3 +2161,61 @@ async def test_embedding_mode_wildcard_health_check_without_a_served_embedding_m assert [json.loads(call.request.content)["model"] for call in upstream.calls] == ["databricks-claude-fixture-a"] assert response == {} + + +_RESPONSES_API_BODY: Final = MappingProxyType( + { + "id": "resp_health", + "object": "response", + "created_at": 1700000000, + "status": "completed", + "model": "gpt-fixture-mini", + "output": [ + { + "type": "message", + "id": "msg_health", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "2", "annotations": []}], + } + ], + "usage": {"input_tokens": 5, "output_tokens": 1, "total_tokens": 6}, + } +) +_TEXT_COMPLETION_BODY: Final = MappingProxyType( + { + "id": "cmpl-health", + "object": "text_completion", + "created": 1700000000, + "model": "gpt-fixture-mini", + "choices": [{"text": "2", "index": 0, "logprobs": None, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 5, "completion_tokens": 1, "total_tokens": 6}, + } +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("mode", "url", "body", "token_cap_param"), + [ + ("responses", "https://api.openai.com/v1/responses", _RESPONSES_API_BODY, "max_output_tokens"), + ("completion", "https://api.openai.com/v1/completions", _TEXT_COMPLETION_BODY, "max_tokens"), + ], +) +async def test_text_generating_mode_wildcard_health_check_caps_the_probe_output_tokens( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + mode: str, + url: str, + body: Mapping[str, object], + token_cap_param: str, +) -> None: + _register_catalog(monkeypatch, "openai", {"openai/gpt-fixture-mini": (mode, 1e-9)}) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + upstream: Final = respx_mock.post(url).respond(json=dict(body)) + + response: Final = await ahealth_check({"model": "openai/gpt-fixture-*", "api_key": "fake-key"}, mode=mode) + + assert "error" not in response, response + assert json.loads(upstream.calls.last.request.content)[token_cap_param] == 16