scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-28 01:32:17 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 42

fix(proxy): accept an inferred probe mode and drop the added commentary

Browse source 
A deployment without model_info.mode gets its probe mode from
litellm.model_cost inside ahealth_check, so compare a non-admin's explicit
mode against that same lookup instead of rejecting it. Use Mapping types
and mark the two unavoidable mutable constructions, and remove the
explanatory comments and docstrings per the repository's comment policy.
This commit is contained in:
Ishkirat-Singh 2026-09-09 16:21:09 +05:30
parent 9b72cf720b
commit f7a5fb8467
2 changed files with 55 additions and 51 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

47
litellm/proxy/health_endpoints/_health_endpoints.py
View file

@ -2001,17 +2001,26 @@ async def health_liveliness_options():
return Response(headers=response_headers, status_code=200)
# What a caller admitted through the non-admin path gets to see of the probe:
# the model and the outcome, none of the deployment's routing configuration.
_NON_ADMIN_TEST_CONNECTION_RESULT_KEYS: Final[frozenset[str]] = frozenset(("model", "error", "mode_error"))
def _test_connection_result_for_display(endpoint_data: dict, *, outcome_only: bool) -> dict:
"""Clean the probe result for display; ``outcome_only`` hides the deployment's configuration."""
cleaned: Final = _clean_endpoint_data(endpoint_data, details=True)
def _test_connection_result_for_display(
litellm_params: Mapping[str, object], result: Mapping[str, object], *, outcome_only: bool
) -> Mapping[str, object]:
cleaned: Final = _clean_endpoint_data({**litellm_params, **result}, details=True)
if not outcome_only:
return cleaned
return {k: v for k, v in cleaned.items() if k in _NON_ADMIN_TEST_CONNECTION_RESULT_KEYS}
return { # mutable-ok: fresh filtered copy handed to the caller
k: v for k, v in cleaned.items() if k in _NON_ADMIN_TEST_CONNECTION_RESULT_KEYS
}
def _configured_probe_mode(model_info: Mapping[str, object] | None, model: object) -> str | None:
configured: Final = model_info.get("mode") if model_info else None
if configured is not None:
return str(configured)
cost_entry: Final = litellm.model_cost.get(model) if isinstance(model, str) else None
return cost_entry.get("mode") if cost_entry else None
def _probe_is_configured_deployment(
@ -2022,11 +2031,6 @@ def _probe_is_configured_deployment(
request_litellm_params: Mapping[str, object],
requested_mode: str | None,
) -> bool:
"""Whether a probe targets a team-less configured deployment exactly as configured.
A request value that differs from the configuration (model, provider,
endpoint, credentials, mode, ...) describes a different probe.
"""
if getattr(model_params.model_info, "team_id", None) is not None:
return False
if any(configured_litellm_params.get(key) != value for key, value in request_litellm_params.items()):
@ -2041,7 +2045,6 @@ async def _assert_caller_can_call_model(
prisma_client: "PrismaClient",
llm_router: "Router | None",
) -> None:
"""Raise 403 unless the caller's key and user may call ``model``."""
from litellm.proxy.auth.auth_checks import (
UserNotFoundError,
can_key_call_model,
@ -2068,7 +2071,10 @@ async def _assert_caller_can_call_model(
user_object = None
await can_user_call_model(model=model, llm_router=llm_router, user_object=user_object)
except ProxyException as e:
raise HTTPException(status_code=403, detail={"error": str(e.message)}) from e
raise HTTPException(
status_code=403,
detail={"error": str(e.message)}, # mutable-ok: same 403 payload shape as the rest of this endpoint
) from e
async def _authorize_test_connection(
@ -2084,16 +2090,7 @@ async def _authorize_test_connection(
request_litellm_params: Mapping[str, object],
requested_mode: str | None,
) -> bool:
"""Decide whether the caller may probe this model.
Proxy admins and team admins may probe any model they manage, as before.
Any other user may probe a configured deployment that has no team, exactly
as configured, when their key and user are allowed to call its model.
Anything else stays a management operation.
Returns True when the caller was admitted through that non-admin path, so
the response can be limited to the outcome of the probe.
"""
"""Returns True when a non-admin was admitted to probe a team-less deployment as configured."""
from litellm.proxy.management_endpoints.model_management_endpoints import (
ModelManagementAuthChecks,
)
@ -2321,7 +2318,7 @@ async def test_model_connection(
llm_router=llm_router,
configured_model_name=configured_model_name,
configured_litellm_params=config_litellm_params,
configured_mode=(loaded_model_info or {}).get("mode"),
configured_mode=_configured_probe_mode(loaded_model_info, config_litellm_params.get("model")),
request_litellm_params=request_litellm_params,
requested_mode=mode or request_litellm_params.get("mode"),
)
@ -2339,7 +2336,7 @@ async def test_model_connection(
# Clean the result for display
cleaned_result: Final = _test_connection_result_for_display(
{**litellm_params, **result}, outcome_only=admitted_as_caller
litellm_params, result, outcome_only=admitted_as_caller
)
return {

59
tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py
View file

@ -4205,10 +4205,6 @@ def _internal_user():
@pytest.mark.asyncio
async def test_test_model_connection_allows_internal_user_to_probe_configured_model_they_can_call():
"""
An internal user may test a configured, non-team model they are allowed
to call, and the probe runs with the configured credentials.
"""
mock_router = MagicMock()
mock_router.get_deployment.return_value = _configured_non_team_deployment()
@ -4233,7 +4229,6 @@ async def test_test_model_connection_allows_internal_user_to_probe_configured_mo
assert key_check.await_args.kwargs["model"] == "gpt-4o"
assert user_check.await_args.kwargs["model"] == "gpt-4o"
assert health_check.await_args.kwargs["model_params"]["api_key"] == "CONFIGURED-API-KEY"
# the caller learns the outcome, not the deployment's routing configuration
assert "api_base" not in result["result"]
assert "api_key" not in result["result"]
@ -4278,11 +4273,6 @@ async def test_test_model_connection_denies_internal_user_without_model_access()
async def test_test_model_connection_keeps_connection_overrides_admin_only():
from fastapi import HTTPException
"""
A request that sets its own connection fields describes a different
endpoint than the configured one; probing that stays a management
operation, so an internal user is denied before any access check runs.
"""
mock_router = MagicMock()
mock_router.get_deployment.return_value = _configured_non_team_deployment()
@ -4312,11 +4302,6 @@ async def test_test_model_connection_keeps_connection_overrides_admin_only():
@pytest.mark.asyncio
async def test_test_model_connection_keeps_model_overrides_admin_only():
"""
The probe runs whatever `litellm_params.model` the request sends, so a
non-admin may only send the configured deployment's own model: the access
check and the probe must target the same model.
"""
from fastapi import HTTPException
mock_router = MagicMock()
@ -4345,11 +4330,6 @@ async def test_test_model_connection_keeps_model_overrides_admin_only():
@pytest.mark.asyncio
async def test_test_model_connection_keeps_team_deployments_admin_only_for_non_admins():
"""
A team deployment keeps its team-admin policy: a non-admin whose key can
call the same model name must not reach the non-admin path with the
team's configured credentials.
"""
from fastapi import HTTPException
from litellm.proxy._types import LiteLLM_TeamTable
@ -4393,7 +4373,6 @@ async def test_test_model_connection_keeps_team_deployments_admin_only_for_non_a
@pytest.mark.asyncio
async def test_test_model_connection_tolerates_missing_user_record_for_non_admin():
"""A key whose user record is gone is judged on the key's own model access."""
from litellm.proxy.auth.auth_checks import UserNotFoundError
mock_router = MagicMock()
@ -4425,11 +4404,6 @@ async def test_test_model_connection_tolerates_missing_user_record_for_non_admin
@pytest.mark.asyncio
async def test_test_model_connection_keeps_mode_overrides_admin_only():
"""
`mode` selects which provider operation the probe performs, so a non-admin
may only probe with the deployment's configured mode (or none, which
auto-detects it).
"""
from fastapi import HTTPException
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
@ -4461,3 +4435,36 @@ async def test_test_model_connection_keeps_mode_overrides_admin_only():
assert exc_info.value.status_code == 403
key_check.assert_not_awaited()
health_check.assert_not_awaited()
@pytest.mark.asyncio
async def test_test_model_connection_accepts_mode_the_probe_would_infer():
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
deployment = Deployment(
model_name="gpt-4o",
litellm_params=LiteLLM_Params(model="gpt-4o", api_key="CONFIGURED-API-KEY"),
model_info=ModelInfo(id="non-team-deployment-id"),
)
mock_router = MagicMock()
mock_router.get_deployment.return_value = deployment
with (
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
patch("litellm.proxy.proxy_server.llm_router", mock_router),
patch("litellm.proxy.proxy_server.premium_user", True),
patch("litellm.proxy.auth.auth_checks.can_key_call_model", AsyncMock(return_value=True)),
patch("litellm.proxy.auth.auth_checks.get_user_object", AsyncMock(return_value=None)),
patch("litellm.proxy.auth.auth_checks.can_user_call_model", AsyncMock(return_value=True)),
patch("litellm.ahealth_check", AsyncMock(return_value={"status": "healthy"})) as health_check,
):
result = await health_test_model_connection(
request=MagicMock(),
mode="chat",
litellm_params={"model": "gpt-4o"},
model_info={"id": "non-team-deployment-id"},
user_api_key_dict=_internal_user(),
)
assert result["status"] == "success"
assert health_check.await_args.kwargs["mode"] == "chat"