fix(health): run during-call and post-call hooks on /health/test_model turns, match 72 as a whole number in Fahrenheit

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

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"(?<!\d)72(?!\d)(?!\s*°?\s*C\b)")
+
+
 def _answer_failure(turn: RepeatedTurn, purpose: str, expectation: str) -> 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

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"] = "<redacted by a guardrail>"
         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 = "<masked by an output guardrail>"  # 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]

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"]