From 0f7ed4a16943ff14dfb8164cf0386d2dc3991cc2 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 5 Nov 2026 03:10:47 +0000 Subject: [PATCH] fix(health): run during-call and post-call hooks on /health/test_model turns, match 72 as a whole number in Fahrenheit --- .../health_endpoints/_health_endpoints.py | 58 +++++++++++-- .../proxy/health_endpoints/model_checks.py | 6 +- .../health_endpoints/test_health_endpoints.py | 87 ++++++++++++++++++- .../health_endpoints/test_model_checks.py | 20 +++++ 4 files changed, 163 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index ff79707ca89..542ae5219d1 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -16,7 +16,7 @@ from typing import Annotated, Final, Literal, Protocol, TypeAlias, TypedDict import fastapi from fastapi import APIRouter, Depends, HTTPException, Request, Response, status -from pydantic import BaseModel, ConfigDict, TypeAdapter +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError from typing_extensions import ReadOnly import litellm @@ -123,11 +123,15 @@ 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 +from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse +from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.proxy.litellm_pre_call_utils import RedactedDict, SecretFields from litellm.types.router import Deployment, reject_server_owned_wif_fields from litellm.types.utils import ( AdminApiKeySecretReference, CustomPricingLiteLLMParams, + LLMResponseTypes, + ModelResponse, admin_api_key_secret_reference, ) @@ -2735,6 +2739,22 @@ async def _logged_cost(logging_obj: LiteLLMLoggingObj, fallback: float | None) - return float(logged) if isinstance(logged, int | float) else fallback +_MESSAGES_RESPONSE: Final = TypeAdapter(AnthropicMessagesResponse) + + +def _as_llm_response(response: object) -> LLMResponseTypes | None: + """The reply as the post-call hooks type it: chat and responses replies are pydantic objects, a /v1/messages + reply is a plain dict shaped like AnthropicMessagesResponse.""" + if isinstance(response, ModelResponse | ResponsesAPIResponse): + return response + if not isinstance(response, dict): + return None + try: + return _MESSAGES_RESPONSE.validate_python(response) + except ValidationError: + return None + + class TurnHooks(Protocol): """The part of ProxyLogging a health-check turn goes through, the same two hooks a proxied request hits.""" @@ -2745,6 +2765,20 @@ class TurnHooks(Protocol): call_type: TestModelRouteType, ) -> dict[str, object]: ... # mutable-ok: ProxyLogging.pre_call_hook returns that same request dict + async def during_call_hook( + self, + data: dict[str, object], # mutable-ok: ProxyLogging.during_call_hook takes the request dict + user_api_key_dict: UserAPIKeyAuth | None, + call_type: TestModelRouteType, + ) -> object: ... + + async def post_call_success_hook( + self, + data: dict[str, object], # mutable-ok: ProxyLogging.post_call_success_hook takes the request dict + response: LLMResponseTypes, + user_api_key_dict: UserAPIKeyAuth, + ) -> object: ... + async def post_call_failure_hook( self, request_data: dict[str, object], # mutable-ok: ProxyLogging.post_call_failure_hook takes the request dict @@ -2779,8 +2813,9 @@ def _turn_caller( deployment's own credentials and pricing apply, or straight through litellm with the submitted params. Either way the turn picks up the guardrails of the calling key, team, project and policies, then the ones on the pinned deployment (or on the deployments serving the requested 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 pre-call, during-call and post-call hooks, so the caller's rate limits, parallel request cap and + those guardrails apply to it, and mask its answers, the way they do for 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 @@ -2812,8 +2847,19 @@ def _turn_caller( ) 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 + during: Final = asyncio.create_task( + hooks.during_call_hook(data=admitted, user_api_key_dict=user_api_key_dict, call_type=script.route_type) + ) + response, _ = await asyncio.gather( + asyncio.wait_for(getattr(target, script.route_type)(**prepared), HEALTH_CHECK_TIMEOUT_SECONDS), during + ) + llm_response: Final = _as_llm_response(response) + checked: Final = ( + await hooks.post_call_success_hook( + data=admitted, response=llm_response, user_api_key_dict=user_api_key_dict + ) + if llm_response is not None + else response ) except Exception as error: await hooks.post_call_failure_hook( @@ -2823,7 +2869,7 @@ def _turn_caller( route="/health/test_model", ) raise - return logging_obj, response + return logging_obj, checked 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 diff --git a/litellm/proxy/health_endpoints/model_checks.py b/litellm/proxy/health_endpoints/model_checks.py index 2c5ed1e4680..9c06e546f5c 100644 --- a/litellm/proxy/health_endpoints/model_checks.py +++ b/litellm/proxy/health_endpoints/model_checks.py @@ -17,6 +17,7 @@ from __future__ import annotations import json import math +import re import time from collections.abc import Awaitable, Callable, Mapping, Sequence from dataclasses import dataclass @@ -476,13 +477,16 @@ async def _turn(script: EndpointScript, body: Mapping[str, object], call: TurnCa return TurnOutcome(response=called.response, cost=called.cost, error=None, call_id=call_id) +_SEVENTY_TWO_FAHRENHEIT: Final = re.compile(r"(? CheckResult | None: if turn.outcome is None or turn.outcome.error is not None: reason: Final = turn.outcome.error if turn.outcome is not None else "not sent" return CheckResult("conversation", "fail", f"turn {turn.number} request with {purpose} failed: {reason}") if isinstance(turn.parsed, str): return CheckResult("conversation", "fail", f"turn {turn.number} {turn.parsed}") - if turn.parsed is None or "72" not in turn.parsed.text: + if turn.parsed is None or _SEVENTY_TWO_FAHRENHEIT.search(turn.parsed.text) is None: answer: Final = turn.parsed.text[:160] if turn.parsed is not None else "" return CheckResult("conversation", "fail", f"turn {turn.number} {expectation}: {answer!r}") return None diff --git a/tests/unit/proxy/health_endpoints/test_health_endpoints.py b/tests/unit/proxy/health_endpoints/test_health_endpoints.py index 5b1f4223e52..db20625a30b 100644 --- a/tests/unit/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/unit/proxy/health_endpoints/test_health_endpoints.py @@ -40,7 +40,7 @@ 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 +from litellm.types.utils import LLMResponseTypes, oauth_token_exchange_litellm_params from litellm.types.workload_identity import ( ANTHROPIC_WIF_KWARGS_KEYS, OPENAI_WIF_KWARGS_KEYS, @@ -6276,6 +6276,16 @@ class _RewritingHooks: first["content"] = "" raise HTTPException(status_code=429, detail="Rate limit exceeded") + async def during_call_hook( + self, data: dict[str, object], user_api_key_dict: UserAPIKeyAuth | None, call_type: TestModelRouteType + ) -> object: + return None + + async def post_call_success_hook( + self, data: dict[str, object], response: LLMResponseTypes, user_api_key_dict: UserAPIKeyAuth + ) -> object: + return response + async def post_call_failure_hook( self, request_data: dict[str, object], @@ -6323,6 +6333,16 @@ class _RecordingHooks: self.guardrails = (*self.guardrails, metadata.get("guardrails")) raise HTTPException(status_code=400, detail="stopped before the provider") + async def during_call_hook( + self, data: dict[str, object], user_api_key_dict: UserAPIKeyAuth | None, call_type: TestModelRouteType + ) -> object: + return None + + async def post_call_success_hook( + self, data: dict[str, object], response: LLMResponseTypes, user_api_key_dict: UserAPIKeyAuth + ) -> object: + return response + async def post_call_failure_hook( self, request_data: dict[str, object], @@ -6362,3 +6382,68 @@ async def test_test_model_turns_pick_up_the_keys_guardrails_and_the_pinned_team_ await call(CHAT_COMPLETIONS_SCRIPT, CHAT_COMPLETIONS.turn_1, "turn-a") assert hooks.guardrails == (["key-prompt-guard", "pii-mask"],) + + +class _MaskingHooks: + """Lets the turn through and masks the answer on the way out, like an output guardrail.""" + + def __init__(self) -> None: + self.during_calls: int = 0 + self.masked: tuple[LLMResponseTypes, ...] = () + + async def pre_call_hook( + self, user_api_key_dict: UserAPIKeyAuth, data: dict[str, object], call_type: TestModelRouteType + ) -> dict[str, object]: + return data + + async def during_call_hook( + self, data: dict[str, object], user_api_key_dict: UserAPIKeyAuth | None, call_type: TestModelRouteType + ) -> object: + self.during_calls += 1 + return None + + async def post_call_success_hook( + self, data: dict[str, object], response: LLMResponseTypes, user_api_key_dict: UserAPIKeyAuth + ) -> object: + assert isinstance(response, litellm.ModelResponse) + masked: Final = copy.deepcopy(response) + masked.choices[0].message.content = "" # pyright: ignore[reportAttributeAccessIssue] # a non-streaming choice + self.masked = (*self.masked, masked) + return masked + + 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_run_the_during_and_post_call_hooks_and_report_the_masked_answer() -> None: + router: Final = Router( + model_list=[ + { + "model_name": "mocked-model", + "litellm_params": {"model": "openai/synthetic", "api_key": "synthetic", "mock_response": "It is 72F"}, + "model_info": {"id": "mocked-deployment"}, + } + ] + ) + hooks: Final = _MaskingHooks() + call: Final = _health_endpoints_module._turn_caller( # pyright: ignore[reportPrivateUsage] # the seam under test + router, + pinned_model_id="mocked-deployment", + call_params={}, + proxy_metadata={}, + user_api_key_dict=UserAPIKeyAuth(token="synthetic-token"), + hooks=hooks, + ) + + turn: Final = await call(CHAT_COMPLETIONS_SCRIPT, CHAT_COMPLETIONS.turn_1, "turn-a") + + assert hooks.during_calls == 1 + assert turn.response is hooks.masked[0] diff --git a/tests/unit/proxy/health_endpoints/test_model_checks.py b/tests/unit/proxy/health_endpoints/test_model_checks.py index c13463eb511..4730588ce13 100644 --- a/tests/unit/proxy/health_endpoints/test_model_checks.py +++ b/tests/unit/proxy/health_endpoints/test_model_checks.py @@ -550,3 +550,23 @@ def test_the_three_request_files_send_one_system_prompt_and_one_tool_so_endpoint assert CHAT_COMPLETIONS_SCRIPT.turn_one() == chat.turn_1 assert RESPONSES_SCRIPT.turn_one() == responses.turn_1 assert MESSAGES_SCRIPT.turn_one() == messages.turn_1 + + +@pytest.mark.parametrize( + ("answer", "status"), + [ + ("It is 72°F in Boston.", "pass"), + ("Around 72 degrees.", "pass"), + ("It is 172F.", "fail"), + ("It is 72C.", "fail"), + ], +) +def test_the_answer_check_wants_seventy_two_as_a_whole_number_and_not_in_celsius(answer: str, status: str) -> None: + provider: Final = _ScriptedProvider( + CalledTurn(_chat_reply(tool_call=True), 0.004), + [CalledTurn(_chat_reply(tool_call=False, text=answer, cached_tokens=4800), 0.0005)] * 2, + ) + + checks: Final = _checks(_run(CHAT_COMPLETIONS_SCRIPT, provider)) + + assert checks["conversation"][0] == status, checks["conversation"]