mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-02 02:11:58 +00:00
fix(proxy): forward model_info in /health/test_connection
test_model_connection hardcoded model_info={} when calling
_update_litellm_params_for_health_check, silently dropping flags like
health_check_supports_max_tokens: false and always injecting max_tokens: 5
for chat mode. The background /health endpoint already honored this flag;
only the UI Test Connection button was broken.
When the router resolves a deployment, use its server-side model_info so
routing-affecting fields (health_check_model, mode, health_check_voice)
cannot be overridden by the request body after authorization. When no
deployment is resolved (ad-hoc credentials), fall back to the caller-supplied
model_info.
Fixes #29266
This commit is contained in:
parent
af17400c38
commit
4e286a4b8b
2 changed files with 73 additions and 1 deletions
|
|
@ -1771,6 +1771,7 @@ async def test_model_connection(
|
|||
# Look up model configuration from router if model name is provided
|
||||
# This gets the litellm_params from proxy config (with resolved env vars)
|
||||
config_litellm_params: dict = {}
|
||||
resolved_model_info: Optional[dict] = None
|
||||
if llm_router is not None:
|
||||
# Prefer disambiguation by deployment id (`model_info.id`) when
|
||||
# the caller supplies it. This is required when multiple
|
||||
|
|
@ -1793,6 +1794,9 @@ async def test_model_connection(
|
|||
config_litellm_params = deployment_by_id.litellm_params.model_dump(
|
||||
exclude_none=True
|
||||
)
|
||||
resolved_model_info = deployment_by_id.model_info.model_dump(
|
||||
exclude_none=True
|
||||
)
|
||||
elif model_name:
|
||||
# Fall back to model_name lookup for callers (e.g. the
|
||||
# "Add Model" wizard, or curl) that don't supply an id.
|
||||
|
|
@ -1819,6 +1823,9 @@ async def test_model_connection(
|
|||
config_litellm_params = dict(
|
||||
deployments[0].get("litellm_params", {})
|
||||
)
|
||||
resolved_model_info = dict(
|
||||
deployments[0].get("model_info", {}) or {}
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.debug(
|
||||
f"Could not find model {model_name} in router: {e}. "
|
||||
|
|
@ -1842,7 +1849,7 @@ async def test_model_connection(
|
|||
)
|
||||
# Include health_check_params if provided
|
||||
litellm_params = _update_litellm_params_for_health_check(
|
||||
model_info={},
|
||||
model_info=resolved_model_info if resolved_model_info is not None else (model_info or {}),
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
mode = mode or litellm_params.pop("mode", None)
|
||||
|
|
|
|||
|
|
@ -696,6 +696,71 @@ async def test_test_model_connection_falls_back_to_deployments_zero_without_id()
|
|||
assert model_params.get("api_key") == "fake-key-A"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_test_model_connection_respects_health_check_supports_max_tokens():
|
||||
"""
|
||||
Regression test for https://github.com/BerriAI/litellm/issues/29266.
|
||||
|
||||
When model_info.health_check_supports_max_tokens is False, the outgoing
|
||||
request must not include max_tokens. Previously model_info={} was hardcoded
|
||||
in the call to _update_litellm_params_for_health_check, so the flag was
|
||||
silently dropped and max_tokens: 5 was always injected for chat mode.
|
||||
"""
|
||||
mock_request = MagicMock()
|
||||
mock_user_api_key_dict = MagicMock()
|
||||
mock_user_api_key_dict.user_id = "test-user"
|
||||
mock_user_api_key_dict.token = "test-token"
|
||||
mock_prisma_client = MagicMock()
|
||||
mock_router = MagicMock()
|
||||
mock_router.get_deployment.return_value = None
|
||||
mock_router.get_model_list.return_value = []
|
||||
mock_can_user_make_model_call = AsyncMock()
|
||||
mock_ahealth_check = AsyncMock(return_value={"status": "healthy"})
|
||||
mock_run_with_timeout = AsyncMock(return_value={"status": "healthy"})
|
||||
|
||||
def mock_reject_os_environ(params):
|
||||
return None
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client),
|
||||
patch("litellm.proxy.proxy_server.llm_router", mock_router),
|
||||
patch("litellm.proxy.proxy_server.premium_user", False),
|
||||
patch(
|
||||
"litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call",
|
||||
mock_can_user_make_model_call,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.health_endpoints._health_endpoints.litellm.ahealth_check",
|
||||
mock_ahealth_check,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.health_endpoints._health_endpoints.run_with_timeout",
|
||||
mock_run_with_timeout,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.health_endpoints._health_endpoints._reject_os_environ_references",
|
||||
mock_reject_os_environ,
|
||||
),
|
||||
):
|
||||
await health_test_model_connection(
|
||||
request=mock_request,
|
||||
mode="chat",
|
||||
litellm_params={
|
||||
"model": "azure/gpt-4o",
|
||||
"api_key": "fake-key",
|
||||
"api_base": "https://fake.invalid",
|
||||
},
|
||||
model_info={"health_check_supports_max_tokens": False},
|
||||
user_api_key_dict=mock_user_api_key_dict,
|
||||
)
|
||||
|
||||
model_params = mock_ahealth_check.call_args.kwargs.get("model_params", {})
|
||||
assert "max_tokens" not in model_params, (
|
||||
"max_tokens must not be sent when health_check_supports_max_tokens is False; "
|
||||
f"got model_params={model_params}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_health_services_endpoint_datadog_llm_observability():
|
||||
"""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue