diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 8866e7d6ded..25a812fedbf 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -12,7 +12,7 @@ from dataclasses import dataclass from datetime import datetime, timedelta, timezone from itertools import chain, groupby from types import MappingProxyType -from typing import Annotated, Final, Literal, TypeAlias, TypedDict +from typing import Annotated, Final, Literal, Protocol, TypeAlias, TypedDict import fastapi from fastapi import APIRouter, Depends, HTTPException, Request, Response, status @@ -105,6 +105,7 @@ from litellm.proxy.health_endpoints.model_checks import ( EndpointScript, TestModelEndpoint, TestModelResponse, + TestModelRouteType, endpoint_result_payload, overall_status, response_cost, @@ -117,7 +118,7 @@ from litellm.proxy.middleware.in_flight_requests_middleware import ( get_in_flight_requests, ) from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager -from litellm.proxy.utils import ProxyLogging +from litellm.proxy.utils import check_and_merge_model_level_guardrails from litellm.router import Router from litellm.router_api import deployment_usable_by_team from litellm.secret_managers.main import get_secret_bool @@ -2733,6 +2734,26 @@ async def _logged_cost(logging_obj: LiteLLMLoggingObj, fallback: float | None) - return float(logged) if isinstance(logged, int | float) else fallback +class TurnHooks(Protocol): + """The part of ProxyLogging a health-check turn goes through, the same two hooks a proxied request hits.""" + + async def pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + data: dict[str, object], # mutable-ok: ProxyLogging.pre_call_hook rewrites the request dict in place + call_type: TestModelRouteType, + ) -> dict[str, object]: ... # mutable-ok: ProxyLogging.pre_call_hook returns that same request dict + + async def post_call_failure_hook( + self, + request_data: dict[str, object], # mutable-ok: ProxyLogging.post_call_failure_hook takes the request dict + original_exception: Exception, + user_api_key_dict: UserAPIKeyAuth, + *, + route: str, + ) -> HTTPException | None: ... + + def _turn_caller( llm_router: Router | None, *, @@ -2740,29 +2761,31 @@ def _turn_caller( call_params: Mapping[str, object], proxy_metadata: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth, - proxy_logging_obj: ProxyLogging, + hooks: TurnHooks, + guardrail_model_name: str | None, ) -> Callable[[EndpointScript, Mapping[str, object], str], Awaitable[CalledTurn]]: """The two ways a turn reaches the provider: pinned to a saved deployment through the router, so the deployment's own credentials and pricing apply, or straight through litellm with the submitted params. - Either way the turn goes through the proxy's pre-call hooks first, so the caller's rate limits, parallel - request cap and guardrails apply to it the way they apply to every proxied request, and it carries the - proxy's own logging object, so the cost read back is the one the spend log row gets. Every turn skips - LiteLLM's response cache, so the reply, and the cache-miss retry of the same prompt, come from the provider - and not from a stored answer.""" + Either way the turn picks up the guardrails attached to the deployments under ``guardrail_model_name`` and + goes through the proxy's pre-call hooks, so the caller's rate limits, parallel request cap and those + guardrails apply to it the way they apply to every proxied request, and it carries the + proxy's own logging object, so the cost read back is the one the spend log row gets. Each turn runs in its + own task, so hooks that keep per-request state in context variables (the rate limiter's slot stash) see one + request per turn, as they do for separate proxied requests, and the hooks get a deep copy of the shared + conversation fixture so a content-rewriting guardrail cannot change the next check's prompt. Every turn + skips LiteLLM's response cache, so the reply, and the cache-miss retry of the same prompt, come from the + provider and not from a stored answer.""" - async def call(script: EndpointScript, body: Mapping[str, object], call_id: str) -> CalledTurn: - target_params: Final = {"model": pinned_model_id} if pinned_model_id is not None else call_params - request: Final = { - **target_params, - **body, - script.metadata_key: {**proxy_metadata, "tags": [_test_model_tag(call_id)]}, - "litellm_call_id": call_id, - "drop_params": True, - "cache": {"no-cache": True}, - } + async def admitted_turn( + script: EndpointScript, + request: dict[str, object], # mutable-ok: the pre-call hooks rewrite the request dict in place + ) -> tuple[LiteLLMLoggingObj, object]: try: - admitted: Final = await proxy_logging_obj.pre_call_hook( - user_api_key_dict=user_api_key_dict, data=request, call_type=script.route_type + guarded: Final = check_and_merge_model_level_guardrails( + data=request, llm_router=llm_router, trust_client_model_info=False, model_alias=guardrail_model_name + ) + admitted: Final = await hooks.pre_call_hook( + user_api_key_dict=user_api_key_dict, data=guarded, call_type=script.route_type ) logging_obj, prepared = litellm.utils.function_setup( original_function=script.route_type, @@ -2776,13 +2799,26 @@ def _turn_caller( getattr(target, script.route_type)(**prepared), HEALTH_CHECK_TIMEOUT_SECONDS ) except Exception as error: - await proxy_logging_obj.post_call_failure_hook( + await hooks.post_call_failure_hook( request_data=request, original_exception=error, user_api_key_dict=user_api_key_dict, route="/health/test_model", ) raise + return logging_obj, response + + async def call(script: EndpointScript, body: Mapping[str, object], call_id: str) -> CalledTurn: + target_params: Final = {"model": pinned_model_id} if pinned_model_id is not None else call_params + request: Final = { + **target_params, + **copy.deepcopy(dict(body)), + script.metadata_key: {**proxy_metadata, "tags": [_test_model_tag(call_id)]}, + "litellm_call_id": call_id, + "drop_params": True, + "cache": {"no-cache": True}, + } + logging_obj, response = await asyncio.create_task(admitted_turn(script, request)) return CalledTurn(response=response, cost=await _logged_cost(logging_obj, response_cost(response))) return call @@ -2914,11 +2950,12 @@ async def test_model( overrides_saved_deployment: Final = any( key != "model" or value != resolved.stored_model for key, value in resolved.request_litellm_params.items() ) - pinned_model_id: Final = ( - resolved.deployment_by_id.model_info.id + pinned: Final = ( + resolved.deployment_by_id if resolved.deployment_by_id is not None and not overrides_saved_deployment else None ) + pinned_model_id: Final = pinned.model_info.id if pinned is not None else None proxy_metadata: Final = { **LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(user_api_key_dict=user_api_key_dict), "user_api_key": LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict), @@ -2935,7 +2972,8 @@ async def test_model( }, proxy_metadata=proxy_metadata, user_api_key_dict=user_api_key_dict, - proxy_logging_obj=proxy_logging_obj, + hooks=proxy_logging_obj, + guardrail_model_name=pinned.model_name if pinned is not None else None, ) run_query: Final[_RawQuery | None] = prisma_client.db.query_raw if prisma_client is not None else None started_at: Final = datetime.now(timezone.utc) - timedelta(minutes=1) diff --git a/tests/unit/proxy/health_endpoints/test_health_endpoints.py b/tests/unit/proxy/health_endpoints/test_health_endpoints.py index 2258eb95fb8..8d974c27677 100644 --- a/tests/unit/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/unit/proxy/health_endpoints/test_health_endpoints.py @@ -35,6 +35,9 @@ from litellm.proxy.health_endpoints._health_endpoints import ( from litellm.proxy.health_endpoints._health_endpoints import ( test_model_connection as health_test_model_connection, ) +from litellm.proxy.health_endpoints.model_checks import CHAT_COMPLETIONS_SCRIPT, TestModelRouteType +from litellm.proxy.health_endpoints.test_model_requests import CHAT_COMPLETIONS +from litellm.proxy.hooks.parallel_request_limiter_v3 import RequestRateLimiterStash, claim_request_stash_for_data from litellm.router import Router from litellm.types.router import RouterModelGroupAliasItem from litellm.types.utils import oauth_token_exchange_litellm_params @@ -6199,6 +6202,7 @@ async def test_unsaved_chatgpt_account_probe_requires_proxy_admin( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, admin: bool ) -> None: from starlette.requests import Request + from litellm.proxy import proxy_server monkeypatch.setenv("CHATGPT_AUTH_FILE", str(tmp_path / "missing-default-auth.json")) @@ -6252,3 +6256,108 @@ async def test_unsaved_chatgpt_account_probe_requires_proxy_admin( await call assert error.value.code == "403" assert len(upstream.calls) == 0 + + +class _RewritingHooks: + """Stands in for ProxyLogging: a guardrail that rewrites the prompt, then a rate limiter that rejects it.""" + + def __init__(self) -> None: + self.stashes: tuple[RequestRateLimiterStash, ...] = () + self.failed_routes: tuple[str, ...] = () + + async def pre_call_hook( + self, user_api_key_dict: UserAPIKeyAuth, data: dict[str, object], call_type: TestModelRouteType + ) -> dict[str, object]: + self.stashes = (*self.stashes, claim_request_stash_for_data(data)) + messages: Final = data["messages"] + assert isinstance(messages, list) + first: Final = messages[0] + assert isinstance(first, dict) + first["content"] = "" + raise HTTPException(status_code=429, detail="Rate limit exceeded") + + async def post_call_failure_hook( + self, + request_data: dict[str, object], + original_exception: Exception, + user_api_key_dict: UserAPIKeyAuth, + *, + route: str, + ) -> HTTPException | None: + self.failed_routes = (*self.failed_routes, route) + return None + + +@pytest.mark.asyncio +async def test_test_model_turns_run_the_proxy_hooks_on_a_copy_of_the_fixture_each_in_its_own_request_context() -> None: + hooks: Final = _RewritingHooks() + call: Final = _health_endpoints_module._turn_caller( # pyright: ignore[reportPrivateUsage] # the seam under test + None, + pinned_model_id=None, + call_params={"model": "openai/synthetic"}, + proxy_metadata={}, + user_api_key_dict=UserAPIKeyAuth(token="synthetic-token"), + hooks=hooks, + guardrail_model_name=None, + ) + fixture_before: Final = copy.deepcopy(CHAT_COMPLETIONS.turn_1) + + for call_id in ("turn-a", "turn-b"): + with pytest.raises(HTTPException, match="Rate limit exceeded"): + await call(CHAT_COMPLETIONS_SCRIPT, CHAT_COMPLETIONS.turn_1, call_id) + + assert CHAT_COMPLETIONS.turn_1 == fixture_before + assert hooks.stashes[0] is not hooks.stashes[1] + assert [stash.owner_litellm_call_id for stash in hooks.stashes] == ["turn-a", "turn-b"] + assert hooks.failed_routes == ("/health/test_model", "/health/test_model") + + +class _RecordingHooks: + def __init__(self) -> None: + self.guardrails: tuple[object, ...] = () + + async def pre_call_hook( + self, user_api_key_dict: UserAPIKeyAuth, data: dict[str, object], call_type: TestModelRouteType + ) -> dict[str, object]: + metadata: Final = data.get("metadata") + assert isinstance(metadata, dict) + self.guardrails = (*self.guardrails, metadata.get("guardrails")) + raise HTTPException(status_code=400, detail="stopped before the provider") + + async def post_call_failure_hook( + self, + request_data: dict[str, object], + original_exception: Exception, + user_api_key_dict: UserAPIKeyAuth, + *, + route: str, + ) -> HTTPException | None: + return None + + +@pytest.mark.asyncio +async def test_test_model_turns_pick_up_the_guardrails_attached_to_the_saved_deployment() -> None: + router: Final = Router( + model_list=[ + { + "model_name": "guarded-model", + "litellm_params": {"model": "openai/synthetic", "api_key": "synthetic", "guardrails": ["pii-mask"]}, + "model_info": {"id": "guarded-deployment"}, + } + ] + ) + hooks: Final = _RecordingHooks() + call: Final = _health_endpoints_module._turn_caller( # pyright: ignore[reportPrivateUsage] # the seam under test + router, + pinned_model_id="guarded-deployment", + call_params={}, + proxy_metadata={}, + user_api_key_dict=UserAPIKeyAuth(token="synthetic-token"), + hooks=hooks, + guardrail_model_name="guarded-model", + ) + + with pytest.raises(HTTPException, match="stopped before the provider"): + await call(CHAT_COMPLETIONS_SCRIPT, CHAT_COMPLETIONS.turn_1, "turn-a") + + assert hooks.guardrails == (["pii-mask"],)