scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-13 08:52:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 174

fix(health): cap wildcard text probes and keep the public health exports

Browse source 
Completion and Responses wildcard probes send a 16 token output limit, since the proxy leaves it unset for wildcards and the Responses API ignores max_tokens. The ahealth_check handler code stays as on main, so the litellm.create_health_check_response and litellm.is_decisions_model_mode exports keep working.
This commit is contained in:
mateo-berri 2026-11-07 11:28:21 -08:00
parent efa9f4ab86
commit 2ca990c0dc
3 changed files with 94 additions and 16 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

29
litellm/litellm_core_utils/health_check_helpers.py
View file

@ -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(

23
litellm/main.py
View file

@ -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:

58
tests/unit/litellm_core_utils/test_health_check_helpers.py
View file

@ -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