diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 15888d2601a..540606b7190 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -1858,7 +1858,7 @@ def _responses_ws_frame_guard( } model: Final = admitted.get("model") - # a blocked frame is billed as its own failed request; only the first turn settles the admission reservation + # blocked and refused frames bill as their own failed requests; only the first turn settles the reservation rejected_frame_auth: Final = user_api_key_dict.model_copy(update={"budget_reservation": None}) async def guard(frame: str) -> GuardedResponsesWsFrame: @@ -1906,6 +1906,7 @@ def _responses_ws_frame_guard( frame=json.dumps(rewritten), guardrail_information=_recorded_guardrail_information(guarded_data.get(metadata_key)), pii_tokens=_recorded_pii_tokens(guarded_data.get("metadata")), + bill_refusal=partial(_bill_rejected_frame, rejected_frame_auth, frame_data=guarded_data), ) return guard diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 20ddbf906b1..1e7368e15eb 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -1962,11 +1962,6 @@ def _turn_rejection_frame(rejection: Exception) -> str: return _turn_admission_error_frame(rejection) -async def _answer_rejected_turn(websocket: ResponsesClientWebSocket, rejection: Exception) -> None: - with contextlib.suppress(Exception): - await websocket.send_text(_turn_rejection_frame(rejection)) - - RESPONSES_WS_LOGGED_EVENT_TYPES: Final = [ "response.created", "response.completed", @@ -2039,16 +2034,28 @@ def _unmask_responses_event(response_str: str, pii_tokens: Mapping[str, str], un @dataclass(frozen=True, slots=True) class GuardedResponsesWsFrame: """A ``response.create`` frame after the proxy's pre_call guardrails ran on it: the frame to forward (guardrails - may rewrite ``input`` and ``instructions``), the guardrail records that bill against this turn only, and the - Presidio ``pii_tokens`` the walk masked with, which the handler needs to unmask the provider's reply.""" + may rewrite ``input`` and ``instructions``), the guardrail records that bill against this turn only, the + Presidio ``pii_tokens`` the walk masked with, which the handler needs to unmask the provider's reply, and how + to bill those guardrail records when the deployment then refuses the turn.""" frame: str guardrail_information: tuple[Mapping[str, object], ...] = () pii_tokens: Mapping[str, str] = field(default_factory=lambda: MappingProxyType({})) + bill_refusal: Callable[[Exception], Awaitable[None]] | None = None ResponsesWsFrameGuard: TypeAlias = Callable[[str], Awaitable[GuardedResponsesWsFrame]] + +async def _reject_guarded_turn( + websocket: ResponsesClientWebSocket, guarded: GuardedResponsesWsFrame, rejection: Exception +) -> None: + with contextlib.suppress(Exception): + await websocket.send_text(_turn_rejection_frame(rejection)) + if guarded.guardrail_information and guarded.bill_refusal is not None: + await guarded.bill_refusal(rejection) + + _RESPONSES_WS_TERMINAL_EVENT_TYPES: Final = _HISTORY_TERMINAL_EVENT_TYPES | _RESPONSES_WS_FAILURE_EVENT_TYPES @@ -2310,7 +2317,7 @@ class ResponsesWebSocketStreaming: self, message: str, turn_admission: Callable[[], Awaitable[None]] | None = None ) -> None: guarded: Final = await self._guard_frame(message) - if guarded is None or not await self._admit_or_reject_turn(message, turn_admission): + if guarded is None or not await self._admit_or_reject_turn(guarded, message, turn_admission): return if guarded.pii_tokens: self._remember_pii_tokens(guarded.pii_tokens) @@ -2708,13 +2715,15 @@ class ResponsesWebSocketStreaming: return False return True - async def _admit_or_reject_turn(self, message: str, turn_admission: Callable[[], Awaitable[None]] | None) -> bool: + async def _admit_or_reject_turn( + self, guarded: GuardedResponsesWsFrame, message: str, turn_admission: Callable[[], Awaitable[None]] | None + ) -> bool: rejection: Final = await _admit_frame_turn(turn_admission, message) if rejection is None: return True if self._failure is None: self._failure = rejection - await _answer_rejected_turn(self.websocket, rejection) + await _reject_guarded_turn(self.websocket, guarded, rejection) return False async def _send_turn_refusal(self, refusal: RateLimitError) -> None: @@ -3290,7 +3299,7 @@ class ManagedResponsesWebSocketHandler: if rejection is not None: if self._failure is None: self._failure = rejection - await _answer_rejected_turn(self.websocket, rejection) + await _reject_guarded_turn(self.websocket, guarded, rejection) return if guarded.pii_tokens: self._pii_tokens = {**self._pii_tokens, **guarded.pii_tokens} diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 03be269fdef..34d7fabbe2e 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -2054,6 +2054,44 @@ class TestResponsesWSFirstFrameModelAuth: assert failure["request_data"]["input"] == "hi" assert len(failure["request_data"]["litellm_call_id"]) == 36 + @pytest.mark.asyncio + async def test_frame_guard_bills_a_passed_frame_the_deployment_refuses_with_its_guardrail_records(self): + from litellm.exceptions import RateLimitError + from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + async def pre_call_hook(user_api_key_dict, data, call_type, guardrails_only): + data["litellm_metadata"]["standard_logging_guardrail_information"] = [{"guardrail_cost": 0.5}] + return {**data, "input": data["input"] + " [guarded]"} + + proxy_logging_obj = MagicMock() + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) + spend_logger = MagicMock(spec=ProxyDBLogger) + spend_logger.async_post_call_failure_hook = AsyncMock() + rate_limiter = MagicMock() + rate_limiter.async_post_call_failure_hook = AsyncMock() + user_api_key_dict = MagicMock() + refusal = RateLimitError(message="crossed budget", llm_provider="", model="") + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, user_api_key_dict, proxy_logging_obj + ) + + guarded = await guard(json.dumps({"type": "response.create", "input": "hi"})) + assert guarded.bill_refusal is not None + with patch.object(litellm, "callbacks", [rate_limiter, spend_logger]): + await guarded.bill_refusal(refusal) + + rate_limiter.async_post_call_failure_hook.assert_not_awaited() + spend_logger.async_post_call_failure_hook.assert_awaited_once() + failure = spend_logger.async_post_call_failure_hook.await_args.kwargs + user_api_key_dict.model_copy.assert_called_once_with(update={"budget_reservation": None}) + assert failure["user_api_key_dict"] is user_api_key_dict.model_copy.return_value + assert failure["original_exception"] is refusal + assert failure["request_data"]["litellm_metadata"]["standard_logging_guardrail_information"] == [ + {"guardrail_cost": 0.5} + ] + assert failure["request_data"]["input"] == "hi [guarded]" + @pytest.mark.asyncio async def test_frame_guard_guards_the_nested_response_body_and_copies_back_tools(self): from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard diff --git a/tests/unit/responses/test_responses_websocket_all_providers.py b/tests/unit/responses/test_responses_websocket_all_providers.py index 030214d95e6..405f1a5b731 100644 --- a/tests/unit/responses/test_responses_websocket_all_providers.py +++ b/tests/unit/responses/test_responses_websocket_all_providers.py @@ -1632,6 +1632,17 @@ def _guard_blocking(word: str): return guard +def _guard_recording(guardrail_information: tuple[Mapping[str, object], ...], bill_refusal: object): + from litellm.responses.streaming_iterator import GuardedResponsesWsFrame + + async def guard(frame: str) -> GuardedResponsesWsFrame: + return GuardedResponsesWsFrame( + frame=frame, guardrail_information=guardrail_information, bill_refusal=bill_refusal + ) + + return guard + + def _errors(websocket: _ClientSocket) -> list[tuple[object, object]]: return [(frame.get("status"), frame["error"]) for frame in websocket.sent if frame.get("type") == "error"] @@ -1833,6 +1844,90 @@ class TestWebSocketTurnAdmission: ] + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("guardrail_information", "billed"), [(({"guardrail_cost": 0.5},), True), ((), False)] + ) + async def test_native_relay_bills_the_guardrails_of_a_turn_the_deployment_refuses( + self, guardrail_information: tuple[Mapping[str, object], ...], billed: bool + ): + from unittest.mock import AsyncMock, call + + from litellm.exceptions import RateLimitError + + first, second, third = _turn_frames() + refusal: Final = RateLimitError(message="crossed budget", llm_provider="", model="") + bill_refusal: Final = AsyncMock() + backend_ws: Final = MagicMock() + backend_ws.send = AsyncMock() + handler: Final = _make_streaming( + websocket=_ClientSocket(second, third), + backend_ws=backend_ws, + request_data={}, + first_message=first, + frame_guard=_guard_recording(guardrail_information, bill_refusal), + turn_admission=AsyncMock(side_effect=[refusal, None]), + ) + + await handler.client_to_backend() + + assert [json.loads(sent.args[0])["input"] for sent in backend_ws.send.await_args_list] == ["one", "three"] + assert bill_refusal.await_args_list == ([call(refusal)] if billed else []) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("guardrail_information", "billed"), [(({"guardrail_cost": 0.5},), True), ((), False)] + ) + async def test_managed_relay_bills_the_guardrails_of_a_turn_the_deployment_refuses( + self, + monkeypatch: pytest.MonkeyPatch, + guardrail_information: tuple[Mapping[str, object], ...], + billed: bool, + ): + from unittest.mock import AsyncMock, call + + import litellm + from litellm.exceptions import RateLimitError + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.responses.streaming_iterator import ManagedResponsesWebSocketHandler + + answered: Final[list[object]] = [] + + async def fake_aresponses(*args: object, **kwargs: object): + answered.append(kwargs.get("input")) + + async def _empty(): + return + yield + + return _empty() + + monkeypatch.setattr(litellm, "aresponses", fake_aresponses) + first, second, third = _turn_frames() + refusal: Final = RateLimitError(message="crossed budget", llm_provider="", model="") + bill_refusal: Final = AsyncMock() + handler: Final = ManagedResponsesWebSocketHandler( + websocket=_ClientSocket(second, third), + model="test-model", + logging_obj=Logging( + model="test-model", + messages=[], + stream=True, + call_type="aresponses", + start_time=0, + litellm_call_id="test-id", + function_id="test-func", + ), + first_message=first, + frame_guard=_guard_recording(guardrail_information, bill_refusal), + turn_admission=AsyncMock(side_effect=[refusal, None]), + ) + + await handler.run() + + assert answered == ["one", "three"] + assert bill_refusal.await_args_list == ([call(refusal)] if billed else []) + @pytest.mark.asyncio async def test_native_relay_answers_a_crashed_turn_check_with_a_500_and_keeps_serving(self): from unittest.mock import AsyncMock