diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 10b6ac5c76c..8866e7d6ded 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -117,6 +117,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.router import Router from litellm.router_api import deployment_usable_by_team from litellm.secret_managers.main import get_secret_bool @@ -2738,33 +2739,50 @@ def _turn_caller( pinned_model_id: str | None, call_params: Mapping[str, object], proxy_metadata: Mapping[str, object], + user_api_key_dict: UserAPIKeyAuth, + proxy_logging_obj: ProxyLogging, ) -> 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 call carries the proxy's own logging object, the way every proxied request does, 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 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.""" 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 - logging_obj, prepared = litellm.utils.function_setup( - original_function=script.route_type, - rules_obj=litellm.utils.Rules(), - start_time=datetime.now(), # noqa: DTZ005 # Logging pairs it with the naive end_time it records - **{ # pyright: ignore[reportArgumentType] # takes the untyped request dict, like the proxy - **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}, - }, - ) - prepared["litellm_logging_obj"] = logging_obj - target: Final = llm_router if pinned_model_id is not None and llm_router is not None else litellm - response: Final[object] = await asyncio.wait_for( - getattr(target, script.route_type)(**prepared), HEALTH_CHECK_TIMEOUT_SECONDS - ) + 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}, + } + 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 + ) + logging_obj, prepared = litellm.utils.function_setup( + original_function=script.route_type, + rules_obj=litellm.utils.Rules(), + start_time=datetime.now(), # noqa: DTZ005 # Logging pairs it with the naive end_time it records + **admitted, # pyright: ignore[reportArgumentType] # takes the untyped request dict, like the proxy + ) + prepared["litellm_logging_obj"] = logging_obj + target: Final = llm_router if pinned_model_id is not None and llm_router is not None else litellm + response: Final[object] = await asyncio.wait_for( + getattr(target, script.route_type)(**prepared), HEALTH_CHECK_TIMEOUT_SECONDS + ) + except Exception as error: + await proxy_logging_obj.post_call_failure_hook( + request_data=request, + original_exception=error, + user_api_key_dict=user_api_key_dict, + route="/health/test_model", + ) + raise return CalledTurn(response=response, cost=await _logged_cost(logging_obj, response_cost(response))) return call @@ -2842,6 +2860,7 @@ async def test_model( llm_router, premium_user, prisma_client, + proxy_logging_obj, ) from litellm.types.router import Deployment, LiteLLM_Params @@ -2915,6 +2934,8 @@ async def test_model( ), }, proxy_metadata=proxy_metadata, + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, ) 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/integration/management/test_model_health_check.py b/tests/integration/management/test_model_health_check.py index 7264737ecee..abae9251c02 100644 --- a/tests/integration/management/test_model_health_check.py +++ b/tests/integration/management/test_model_health_check.py @@ -1645,3 +1645,34 @@ def test_health_test_model_with_a_different_model_than_the_saved_deployment_call assert [request.headers["authorization"] for request in requests] == ["Bearer synthetic-override-key"] * 3, ( response.text ) + + +def test_health_test_model_turns_count_against_the_calling_keys_rate_limit(gateway: Gateway) -> None: + with ( + gateway.scenario() as scenario, + wire_server(lambda request: _test_model_chat_reply(request, cached_tokens=4800)) as wire, + ): + model_name: Final = scenario.model(api_base=f"{wire.url}/v1", api_key="synthetic-test-model-key") + model_id: Final = next( + string_value(object_value(info["model_info"])["id"]) + for info in (object_value(entry) for entry in list_value(gateway.get("/v1/model/info")["data"])) + if info["model_name"] == model_name + ) + admin: Final = scenario.user(user_role="proxy_admin") + one_per_minute: Final = scenario.key(user_id=admin, rpm_limit=1) + + response: Final = gateway.request( + "POST", + "/health/test_model", + {"model_id": model_id, "endpoints": ["/v1/chat/completions"]}, + key=one_per_minute, + ) + + assert response.status_code == 200, response.text + payload: Final = JSON_OBJECT.validate_json(response.content) + assert payload["status"] == "error", response.text + checks: Final = _check_statuses(object_value(list_value(payload["endpoints"])[0])) + assert checks["conversation"][0] == "fail", response.text + assert checks["conversation"][1].startswith("turn 2 request with the tool result failed"), response.text + assert "Rate limit exceeded" in checks["conversation"][1], response.text + assert len(wire.drain()) == 1, response.text