diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index e8050b94343..1dc8a77db21 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -1439,6 +1439,22 @@ class OpenAIResponsesHandler(BaseTranslation): rewrites_by_position=MappingProxyType(dict(zip(placeable_positions, chain((rewritten_text,), repeat(""))))), ) + async def deliver_text_rewrite( + self, + response: "ResponsesAPIResponse", + position: tuple[int, int], + rewritten_text: str, + stream_events: Sequence[object] = (), + ) -> None: + """Write one rewritten text into the response at ``position`` and into any buffered events + that spell out the same text, so a client reading deltas sees what the body says.""" + await self._apply_guardrail_responses_to_output( + response=response, responses=[rewritten_text], task_mappings=[position] + ) + self._sync_stream_events_with_rewrites( + stream_events=stream_events, rewrites_by_position=MappingProxyType({position: rewritten_text}) + ) + @staticmethod def _write_event_field(event: object, field: str, value: object) -> None: if isinstance(event, dict): diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py index 738682a9a29..38c11dab063 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -1,7 +1,7 @@ import time from collections.abc import AsyncGenerator, Awaitable, Callable, Generator, Mapping, Sequence from enum import Enum, auto -from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Protocol, runtime_checkable +from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, NamedTuple, Protocol, runtime_checkable import httpx from fastapi import HTTPException @@ -174,6 +174,12 @@ def _redact_scanned_content(payload: object, depth: int = 0) -> object: return payload +class _ResponsesBodyScan(NamedTuple): + texts: tuple[str, ...] + positions: tuple[tuple[int, int], ...] + tool_calls: tuple[ChatCompletionToolCallChunk, ...] + + class ModelArmorGuardrail(CustomGuardrail, VertexBase): """ Google Cloud Model Armor Guardrail integration for LiteLLM. @@ -540,13 +546,16 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): return self._has_deidentify_match(armor_response), None if sanitized_content == content: return False, None - if not isinstance(assembled_response, ModelResponse): - verbose_proxy_logger.warning( - "Model Armor: sanitized content cannot be re-emitted on this streaming endpoint, " - "blocking the response instead" - ) - return True, None - return False, sanitized_content + if isinstance(assembled_response, ModelResponse) or ( + isinstance(assembled_response, ResponsesAPIResponse) + and self._sole_text_position(assembled_response, content) is not None + ): + return False, sanitized_content + verbose_proxy_logger.warning( + "Model Armor: sanitized content cannot be re-emitted on this streaming endpoint, " + "blocking the response instead" + ) + return True, None @staticmethod def _append_armor_response(existing: object, armor_response: object) -> object: @@ -964,7 +973,9 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): if self.mask_response_content: sanitized_content: Final = self._get_sanitized_content(armor_response) if sanitized_content and sanitized_content != content: - if isinstance(response, ResponsesAPIResponse): + if isinstance(response, ResponsesAPIResponse) and not await self._apply_sanitized_responses_content( + response, content, sanitized_content + ): raise HTTPException( status_code=400, detail=self._build_block_error_detail("Response blocked by Model Armor", armor_response), @@ -1028,17 +1039,14 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): ) @staticmethod - def _responses_api_response_text(response: ResponsesAPIResponse) -> str: - """Text to scan in a Responses API response, tool-call arguments included. - - Tool calls are folded in because ``get_content_from_model_response`` folds them into what - the chat surface scans, and a Responses turn can carry its whole payload in them. - """ + def _scan_responses_api_response(response: ResponsesAPIResponse) -> _ResponsesBodyScan: + """Texts of a Responses API response, where each one sits, and its tool calls.""" from litellm.llms.openai.responses.guardrail_translation.handler import ( OpenAIResponsesHandler, ) texts: Final[list[str]] = [] # mutable-ok: the shared extractor below appends into caller-owned lists + positions: Final[list[tuple[int, int]]] = [] # mutable-ok: the same extractor's position sink tool_calls: Final[list[ChatCompletionToolCallChunk]] = [] # mutable-ok: the same extractor's tool-call sink handler: Final = OpenAIResponsesHandler() for output_idx, output_item in enumerate(response.output or ()): @@ -1047,10 +1055,53 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): output_idx=output_idx, texts_to_check=texts, images_to_check=[], - task_mappings=[], + task_mappings=positions, tool_calls_to_check=tool_calls, ) - return "".join((*texts, *(json.dumps(tool_call) for tool_call in tool_calls))) + return _ResponsesBodyScan(texts=tuple(texts), positions=tuple(positions), tool_calls=tuple(tool_calls)) + + @staticmethod + def _responses_api_response_text(response: ResponsesAPIResponse) -> str: + """Text to scan in a Responses API response, tool-call arguments included. + + Tool calls are folded in because ``get_content_from_model_response`` folds them into what + the chat surface scans, and a Responses turn can carry its whole payload in them. + """ + scan: Final = ModelArmorGuardrail._scan_responses_api_response(response) + return "".join((*scan.texts, *(json.dumps(tool_call) for tool_call in scan.tool_calls))) + + @staticmethod + def _sole_text_position(response: ResponsesAPIResponse, content: str) -> tuple[int, int] | None: + """Where one sanitized string can land: the part ``content`` came from, when that is one part. + + A scan that also joined a second part, tool-call arguments or reasoning deltas comes back as + one string with no boundaries, so it has no single place to go. + """ + scan: Final = ModelArmorGuardrail._scan_responses_api_response(response) + return scan.positions[0] if scan.texts == (content,) else None + + async def _apply_sanitized_responses_content( + self, + response: ResponsesAPIResponse, + content: str, + sanitized_content: str, + stream_events: Sequence[object] = (), + ) -> bool: + """Write the sanitized text into the response and its buffered events, if it has one place to go.""" + from litellm.llms.openai.responses.guardrail_translation.handler import ( + OpenAIResponsesHandler, + ) + + position: Final = self._sole_text_position(response, content) + if position is None: + return False + await OpenAIResponsesHandler().deliver_text_rewrite( + response=response, + position=position, + rewritten_text=sanitized_content, + stream_events=stream_events, + ) + return True def _extract_scannable_content(self, response: object) -> str: """Text to scan from a response, for every endpoint shape this hook serves.""" @@ -1264,6 +1315,11 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): ), ) + if sanitized_content is not None and isinstance(assembled_response, ResponsesAPIResponse): + await self._apply_sanitized_responses_content( + assembled_response, content, sanitized_content, stream_events=all_chunks + ) + if sanitized_content is not None and isinstance(assembled_response, ModelResponse): self._apply_sanitized_content(assembled_response, sanitized_content) diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py index d9cfa6187dc..caa6ee34372 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py @@ -3891,7 +3891,10 @@ async def _anthropic_sse_stream(): yield chunk -def _responses_api_events(): +_SECOND_TEXT_PART: Final = {"type": "output_text", "text": " and more"} + + +def _responses_api_events(*extra_parts: dict): from litellm.types.llms.openai import ( OutputTextDeltaEvent, ResponseCompletedEvent, @@ -3910,7 +3913,7 @@ def _responses_api_events(): "id": "msg_1", "status": "completed", "role": "assistant", - "content": [{"type": "output_text", "text": "my card is 4111-1111-1111-1111"}], + "content": [{"type": "output_text", "text": "my card is 4111-1111-1111-1111"}, *extra_parts], } ], parallel_tool_calls=False, @@ -4047,8 +4050,8 @@ async def test_streaming_masking_re_emits_anthropic_sse_with_sanitized_text(): @pytest.mark.asyncio -async def test_streaming_masking_blocks_responses_api_stream(): - """A Responses event stream cannot be rebuilt from sanitized text, so releasing it would +async def test_streaming_masking_blocks_a_responses_stream_with_no_single_part_to_rewrite(): + """One sanitized string cannot be split back over several parts, so releasing the stream would ship the content the guardrail just rewrote. It is blocked instead.""" from litellm.types.llms.openai import ErrorEvent @@ -4057,7 +4060,7 @@ async def test_streaming_masking_blocks_responses_api_stream(): with patch.object( guardrail.async_handler, "post", _armor_post_mock(_MODEL_ARMOR_SANITIZED) ): - delivered = await _drain_surface_hook(guardrail, _responses_api_events()) + delivered = await _drain_surface_hook(guardrail, _responses_api_events(_SECOND_TEXT_PART)) assert len(delivered) == 1 assert isinstance(delivered[0], ErrorEvent) @@ -4535,7 +4538,7 @@ async def test_streaming_deidentify_match_without_a_rewrite_fails_closed(chunks, @pytest.mark.asyncio async def test_streaming_status_records_a_surface_that_cannot_carry_the_rewrite_as_blocked(): - """The Responses surface has no assembled body to rewrite, so a de-identify match ends as a + """A multi-part Responses body has no single part to rewrite, so a de-identify match ends as a refusal. The status stamped on metadata feeds the spend log and has to say so rather than reporting the success the block check alone would have implied.""" guardrail = _surface_guardrail(mask_response_content=True) @@ -4547,7 +4550,7 @@ async def test_streaming_status_records_a_surface_that_cannot_carry_the_rewrite_ with patch.object(guardrail.async_handler, "post", _armor_post_mock(_MODEL_ARMOR_DEIDENTIFIED)): delivered = await _drain_surface_hook( - guardrail, _responses_api_events(), request_data=request_data + guardrail, _responses_api_events(_SECOND_TEXT_PART), request_data=request_data ) body = _delivered_bytes(delivered) @@ -5354,6 +5357,32 @@ def _armor_block() -> dict: } +def _responses_body(*parts: dict): + from litellm.types.llms.openai import ResponsesAPIResponse + + return ResponsesAPIResponse.model_validate( + { + "id": "resp_1", + "created_at": 0, + "model": "gpt-4o", + "object": "response", + "status": "completed", + "output": [ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": list(parts), + } + ], + "parallel_tool_calls": False, + "tool_choice": "auto", + "tools": [], + } + ) + + def _responses_refusal_response(): from litellm.types.llms.openai import ResponsesAPIResponse @@ -5380,6 +5409,8 @@ def _responses_refusal_response(): ) +_RAW_REFUSAL: Final = "I cannot help card 4111-1111-1111-1111" + _POST_CALL_REQUEST: Final = { "model": "gpt-4", "messages": [{"role": "user", "content": "What's my credit card?"}], @@ -5417,11 +5448,95 @@ async def test_post_call_blocks_on_a_flagged_responses_refusal(): @pytest.mark.asyncio -async def test_post_call_blocks_a_responses_body_it_cannot_rewrite_instead_of_releasing_it(): +async def test_post_call_masks_a_responses_refusal(): guardrail = _model_armor_post_call_guardrail(mask_response_content=True) post = AsyncMock(return_value=_armor_http_response(_armor_redaction("I cannot help card [REDACTED]"))) + response = _responses_refusal_response() + + with patch.object(guardrail.async_handler, "post", post): + await guardrail.async_post_call_success_hook( + data=dict(_POST_CALL_REQUEST), user_api_key_dict=UserAPIKeyAuth(), response=response + ) + + assert "4111-1111-1111-1111" in str(post.call_args) + assert response.output[0].content[0].refusal == "I cannot help card [REDACTED]" + + +@pytest.mark.asyncio +async def test_post_call_masks_plain_text_in_a_responses_body(): + guardrail = _model_armor_post_call_guardrail(mask_response_content=True) + post = AsyncMock(return_value=_armor_http_response(_armor_redaction("card [REDACTED]"))) + response = _responses_body({"type": "output_text", "text": "card 4111-1111-1111-1111", "annotations": []}) + + with patch.object(guardrail.async_handler, "post", post): + await guardrail.async_post_call_success_hook( + data=dict(_POST_CALL_REQUEST), user_api_key_dict=UserAPIKeyAuth(), response=response + ) + + assert response.output[0].content[0].text == "card [REDACTED]" + + +@pytest.mark.asyncio +async def test_post_call_blocks_a_responses_body_whose_parts_one_redaction_cannot_split_across(): + guardrail = _model_armor_post_call_guardrail(mask_response_content=True) + post = AsyncMock(return_value=_armor_http_response(_armor_redaction("card [REDACTED] and more"))) + response = _responses_body( + {"type": "output_text", "text": "card 4111-1111-1111-1111", "annotations": []}, + {"type": "output_text", "text": " and more", "annotations": []}, + ) with patch.object(guardrail.async_handler, "post", post), pytest.raises(HTTPException): await guardrail.async_post_call_success_hook( - data=dict(_POST_CALL_REQUEST), user_api_key_dict=UserAPIKeyAuth(), response=_responses_refusal_response() + data=dict(_POST_CALL_REQUEST), user_api_key_dict=UserAPIKeyAuth(), response=response ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("part", "field"), + [ + ({"type": "refusal", "refusal": _RAW_REFUSAL}, "refusal"), + ({"type": "output_text", "text": _RAW_REFUSAL, "annotations": []}, "text"), + ], +) +async def test_streaming_masks_every_event_that_spells_out_a_responses_part(part, field): + from litellm.types.llms.openai import ( + OutputTextDeltaEvent, + OutputTextDoneEvent, + RefusalDeltaEvent, + RefusalDoneEvent, + ResponseCompletedEvent, + ) + + delta_event, done_event, event_name = ( + (RefusalDeltaEvent, RefusalDoneEvent, "refusal") + if field == "refusal" + else (OutputTextDeltaEvent, OutputTextDoneEvent, "output_text") + ) + position = {"item_id": "msg_1", "output_index": 0, "content_index": 0} + events = ( + delta_event(type=f"response.{event_name}.delta", delta=_RAW_REFUSAL[:18], **position), + delta_event(type=f"response.{event_name}.delta", delta=_RAW_REFUSAL[18:], **position), + done_event(type=f"response.{event_name}.done", **{field: _RAW_REFUSAL}, **position), + ResponseCompletedEvent(type="response.completed", response=_responses_body(part)), + ) + guardrail = _model_armor_post_call_guardrail(mask_response_content=True) + post = AsyncMock(return_value=_armor_http_response(_armor_redaction("I cannot help card [REDACTED]"))) + request_data = {**_POST_CALL_REQUEST, "metadata": dict(_POST_CALL_REQUEST["metadata"])} + + async def upstream(): + for event in events: + yield event + + with patch.object(guardrail.async_handler, "post", post): + released = [ + chunk + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(), response=upstream(), request_data=request_data + ) + ] + + assert request_data["metadata"]["_model_armor_status"] == "success" + assert "4111" not in str(released) + assert "".join(event.delta for event in released[:2]) == "I cannot help card [REDACTED]" + assert getattr(released[2], field) == "I cannot help card [REDACTED]"