mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-02 02:11:58 +00:00
fix black formatting; add router-resolved path test
Black reformats the long ternary in _update_litellm_params_for_health_check call to a multi-line form. Also add a second regression test covering the deployment-resolved-by-id path where server-side model_info carries health_check_supports_max_tokens: False.
This commit is contained in:
parent
4e286a4b8b
commit
2e5e78c35d
2 changed files with 77 additions and 1 deletions
|
|
@ -1849,7 +1849,11 @@ async def test_model_connection(
|
|||
)
|
||||
# Include health_check_params if provided
|
||||
litellm_params = _update_litellm_params_for_health_check(
|
||||
model_info=resolved_model_info if resolved_model_info is not None else (model_info or {}),
|
||||
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)
|
||||
|
|
|
|||
|
|
@ -761,6 +761,78 @@ async def test_test_model_connection_respects_health_check_supports_max_tokens()
|
|||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_test_model_connection_uses_server_model_info_when_deployment_resolved():
|
||||
"""
|
||||
When the router resolves a deployment by id, the server-side model_info
|
||||
must be forwarded to the health-check helper, not the caller-supplied one.
|
||||
|
||||
This covers the router-resolved path for the same bug as
|
||||
test_test_model_connection_respects_health_check_supports_max_tokens.
|
||||
"""
|
||||
from litellm.types.router import Deployment, LiteLLM_Params
|
||||
|
||||
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_can_user_make_model_call = AsyncMock()
|
||||
mock_ahealth_check = AsyncMock(return_value={"status": "healthy"})
|
||||
mock_run_with_timeout = AsyncMock(return_value={"status": "healthy"})
|
||||
|
||||
server_deployment = Deployment(
|
||||
model_name="my-model",
|
||||
litellm_params=LiteLLM_Params(
|
||||
model="azure/gpt-4o",
|
||||
api_key="server-key",
|
||||
api_base="https://server.invalid/v1",
|
||||
),
|
||||
model_info={"id": "dep-1", "health_check_supports_max_tokens": False},
|
||||
)
|
||||
mock_router = MagicMock()
|
||||
mock_router.get_deployment.return_value = server_deployment
|
||||
|
||||
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"},
|
||||
model_info={"id": "dep-1"},
|
||||
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 server model_info has "
|
||||
"health_check_supports_max_tokens: 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