From f671407132aa7002c25ae6bdcd53c4677452e092 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 16:39:22 -0700 Subject: [PATCH] fix(guardrails): keep Model Armor's non-streamed Responses scan on message text A masking Model Armor answered 400 on a Responses body with an email in forced tool call arguments, where main answers 200. Scan only message text on that path, since one sanitized string cannot be written back into tool call arguments, and spread that string over every part so a body with several parts is masked instead of blocked --- .../guardrail_translation/handler.py | 35 ++++-- .../model_armor/model_armor.py | 113 ++++++------------ ...test_openai_responses_guardrail_handler.py | 28 +++++ .../guardrail_hooks/test_model_armor.py | 82 +++++++++++-- 4 files changed, 163 insertions(+), 95 deletions(-) diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 1dc8a77db21..d7b5b57285e 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -1439,21 +1439,40 @@ class OpenAIResponsesHandler(BaseTranslation): rewrites_by_position=MappingProxyType(dict(zip(placeable_positions, chain((rewritten_text,), repeat(""))))), ) - async def deliver_text_rewrite( + def message_part_texts(self, response: "ResponsesAPIResponse") -> Mapping[tuple[int, int], str]: + """Text of every message content part, keyed by ``(output_index, content_index)``.""" + 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 + for output_idx, output_item in enumerate(response.output or ()): + self._extract_output_text_and_images( + output_item=output_item, + output_idx=output_idx, + texts_to_check=texts, + images_to_check=[], + task_mappings=positions, + ) + return MappingProxyType( + {position: text for position, text in zip(positions, texts) if position[1] != _COMPACTION_SUMMARY_POSITION} + ) + + async def spread_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.""" + """Replace all message text in ``response`` with one rewrite, for a guardrail that answers + one string for everything it scanned: the whole rewrite goes on the first part and every + later part is blanked, in the body and in any buffered events spelling out the same text.""" + rewrites_by_position: Final = MappingProxyType( + dict(zip(self.message_part_texts(response), chain((rewritten_text,), repeat("")))) + ) 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}) + response=response, + responses=list(rewrites_by_position.values()), + task_mappings=list(rewrites_by_position), ) + self._sync_stream_events_with_rewrites(stream_events=stream_events, rewrites_by_position=rewrites_by_position) @staticmethod def _write_event_field(event: object, field: str, value: object) -> None: 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 38c11dab063..503089fe790 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, NamedTuple, Protocol, runtime_checkable +from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Protocol, runtime_checkable import httpx from fastapi import HTTPException @@ -10,6 +10,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.llms.openai.responses.guardrail_translation.handler import OpenAIResponsesHandler from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel import json @@ -174,12 +175,6 @@ 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. @@ -548,7 +543,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): return False, None if isinstance(assembled_response, ModelResponse) or ( isinstance(assembled_response, ResponsesAPIResponse) - and self._sole_text_position(assembled_response, content) is not None + and self._responses_message_text(assembled_response) == content ): return False, sanitized_content verbose_proxy_logger.warning( @@ -919,7 +914,11 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): return # Extract content from response - content: Final = self._extract_scannable_content(response) + content: Final = ( + self._responses_message_text(response) + if isinstance(response, ResponsesAPIResponse) + else self._extract_content_from_response(response) + ) if not content: verbose_proxy_logger.debug("Model Armor: No text content to process in response, skipping guardrail") return @@ -973,13 +972,8 @@ 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) 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), - ) + if isinstance(response, ResponsesAPIResponse): + await self._responses_handler().spread_text_rewrite(response, sanitized_content) if isinstance(response, litellm.ModelResponse): self._apply_sanitized_content(response, sanitized_content) @@ -1039,26 +1033,17 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): ) @staticmethod - def _scan_responses_api_response(response: ResponsesAPIResponse) -> _ResponsesBodyScan: - """Texts of a Responses API response, where each one sits, and its tool calls.""" + def _responses_handler() -> "OpenAIResponsesHandler": 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 ()): - handler._extract_output_text_and_images( # pyright: ignore[reportPrivateUsage] # the shared Responses output extractor; forking it would duplicate per-item parsing - output_item=output_item, - output_idx=output_idx, - texts_to_check=texts, - images_to_check=[], - task_mappings=positions, - tool_calls_to_check=tool_calls, - ) - return _ResponsesBodyScan(texts=tuple(texts), positions=tuple(positions), tool_calls=tuple(tool_calls)) + return OpenAIResponsesHandler() + + @staticmethod + def _responses_message_text(response: ResponsesAPIResponse) -> str: + """Message text of a Responses API response, every part joined as Model Armor scans it.""" + return "".join(ModelArmorGuardrail._responses_handler().message_part_texts(response).values()) @staticmethod def _responses_api_response_text(response: ResponsesAPIResponse) -> str: @@ -1067,47 +1052,25 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): 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))) + texts: Final[list[str]] = [] # mutable-ok: the shared extractor below appends into caller-owned lists + tool_calls: Final[list[ChatCompletionToolCallChunk]] = [] # mutable-ok: the same extractor's tool-call sink + handler: Final = ModelArmorGuardrail._responses_handler() + for output_idx, output_item in enumerate(response.output or ()): + handler._extract_output_text_and_images( # pyright: ignore[reportPrivateUsage] # the shared Responses output extractor; forking it would duplicate per-item parsing + output_item=output_item, + output_idx=output_idx, + texts_to_check=texts, + images_to_check=[], + task_mappings=[], + tool_calls_to_check=tool_calls, + ) + return "".join((*texts, *(json.dumps(tool_call) for tool_call in 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.""" - if isinstance(response, ResponsesAPIResponse): - return self._responses_api_response_text(response) - return self._extract_content_from_response(response) + def _extract_streaming_content(self, assembled_response: object) -> str: + """Text to scan from an assembled stream, for every endpoint shape this hook serves.""" + if isinstance(assembled_response, ResponsesAPIResponse): + return self._responses_api_response_text(assembled_response) + return self._extract_content_from_response(assembled_response) @staticmethod def _responses_delta_field(chunk: object) -> tuple[str, ...]: @@ -1145,7 +1108,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): delta events that its terminal body never repeats, so every delta field the body does not already carry is scanned after it. """ - content: Final = self._extract_scannable_content(assembled_response) + content: Final = self._extract_streaming_content(assembled_response) if surface is not _StreamSurface.RESPONSES: return content unscanned: Final = tuple(text for text in self._responses_delta_field_texts(all_chunks) if text not in content) @@ -1316,8 +1279,8 @@ 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 + await self._responses_handler().spread_text_rewrite( + assembled_response, sanitized_content, stream_events=all_chunks ) if sanitized_content is not None and isinstance(assembled_response, ModelResponse): diff --git a/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py index e3089f69413..abcb48c8c02 100644 --- a/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -4966,6 +4966,34 @@ class TestCompactionOutputSummaryReachesOutputGuardrails: assert self._summary(result.output[0]) == "The user lives at 1 [MASKED]." + @pytest.mark.asyncio + async def test_spreading_one_rewrite_over_message_text_leaves_the_summary_alone(self) -> None: + response: Final = ResponsesAPIResponse( + id="resp_1", + created_at=1, + model="claude-sonnet-5-5", + object="response", + status="completed", + output=[ + self._compaction_item(), + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [ + {"type": "output_text", "text": "first", "annotations": []}, + {"type": "refusal", "refusal": "second"}, + ], + }, + ], + ) + + await OpenAIResponsesHandler().spread_text_rewrite(response, "rewritten") + + assert self._summary(response.output[0]) == self.SUMMARY + assert [response.output[1].content[0].text, response.output[1].content[1].refusal] == ["rewritten", ""] + @pytest.mark.asyncio async def test_ended_stream_summary_is_masked_in_the_done_item_and_the_envelope(self) -> None: events: Final = [ 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 caa6ee34372..94e990c1005 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py @@ -3892,9 +3892,17 @@ async def _anthropic_sse_stream(): _SECOND_TEXT_PART: Final = {"type": "output_text", "text": " and more"} +_TOOL_CALL_ITEM: Final = { + "type": "function_call", + "id": "fc_1", + "call_id": "call_1", + "name": "save_card", + "arguments": '{"card": "4111-1111-1111-1111"}', + "status": "completed", +} -def _responses_api_events(*extra_parts: dict): +def _responses_api_events(extra_parts: tuple[dict, ...] = (), extra_items: tuple[dict, ...] = ()): from litellm.types.llms.openai import ( OutputTextDeltaEvent, ResponseCompletedEvent, @@ -3914,7 +3922,8 @@ def _responses_api_events(*extra_parts: dict): "status": "completed", "role": "assistant", "content": [{"type": "output_text", "text": "my card is 4111-1111-1111-1111"}, *extra_parts], - } + }, + *extra_items, ], parallel_tool_calls=False, tool_choice="auto", @@ -4050,9 +4059,9 @@ async def test_streaming_masking_re_emits_anthropic_sse_with_sanitized_text(): @pytest.mark.asyncio -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.""" +async def test_streaming_masking_blocks_a_responses_stream_whose_scan_includes_a_tool_call(): + """Tool-call arguments are scanned on a stream but one sanitized string has no way back into + them, so releasing the stream would ship the content the guardrail just rewrote.""" from litellm.types.llms.openai import ErrorEvent guardrail = _surface_guardrail(mask_response_content=True) @@ -4060,7 +4069,7 @@ async def test_streaming_masking_blocks_a_responses_stream_with_no_single_part_t with patch.object( guardrail.async_handler, "post", _armor_post_mock(_MODEL_ARMOR_SANITIZED) ): - delivered = await _drain_surface_hook(guardrail, _responses_api_events(_SECOND_TEXT_PART)) + delivered = await _drain_surface_hook(guardrail, _responses_api_events(extra_items=(_TOOL_CALL_ITEM,))) assert len(delivered) == 1 assert isinstance(delivered[0], ErrorEvent) @@ -4538,9 +4547,9 @@ 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(): - """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.""" + """A Responses stream that scanned a tool call has nowhere to put the 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) request_data = { "model": "gpt-4o-mini", @@ -4550,7 +4559,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(_SECOND_TEXT_PART), request_data=request_data + guardrail, _responses_api_events(extra_items=(_TOOL_CALL_ITEM,)), request_data=request_data ) body = _delivered_bytes(delivered) @@ -5477,7 +5486,7 @@ async def test_post_call_masks_plain_text_in_a_responses_body(): @pytest.mark.asyncio -async def test_post_call_blocks_a_responses_body_whose_parts_one_redaction_cannot_split_across(): +async def test_post_call_masks_every_part_of_a_responses_body_through_the_first(): 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( @@ -5485,11 +5494,38 @@ async def test_post_call_blocks_a_responses_body_whose_parts_one_redaction_canno {"type": "output_text", "text": " and more", "annotations": []}, ) - with patch.object(guardrail.async_handler, "post", post), pytest.raises(HTTPException): + 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 [part.text for part in response.output[0].content] == ["card [REDACTED] and more", ""] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("mask_response_content", "flagged"), + [(True, _armor_redaction("saving [REDACTED]")), (False, _armor_block())], +) +async def test_post_call_answers_a_tool_call_beside_responses_text_unchanged(mask_response_content, flagged): + from litellm.types.llms.openai import ResponsesAPIResponse + + async def flag_the_card(**kwargs): + verdict = flagged if "4111" in str(kwargs["json"]) else {"filterMatchState": "NO_MATCH_FOUND"} + return _armor_http_response(verdict) + + guardrail = _model_armor_post_call_guardrail(mask_response_content=mask_response_content) + body = _responses_body({"type": "output_text", "text": "saving it", "annotations": []}).model_dump() + response = ResponsesAPIResponse.model_validate({**body, "output": [*body["output"], _TOOL_CALL_ITEM]}) + + with patch.object(guardrail.async_handler, "post", AsyncMock(side_effect=flag_the_card)): + 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 == "saving it" + assert response.output[1].arguments == _TOOL_CALL_ITEM["arguments"] + @pytest.mark.asyncio @pytest.mark.parametrize( @@ -5540,3 +5576,25 @@ async def test_streaming_masks_every_event_that_spells_out_a_responses_part(part 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]" + + +@pytest.mark.asyncio +async def test_streaming_masks_a_responses_body_with_several_parts(): + from litellm.types.llms.openai import OutputTextDeltaEvent, ResponsesAPIStreamEvents + + guardrail = _surface_guardrail(mask_response_content=True) + first_delta, terminal = _responses_api_events(extra_parts=(_SECOND_TEXT_PART,)) + second_delta = OutputTextDeltaEvent( + type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, + item_id="msg_1", + output_index=0, + content_index=1, + delta=_SECOND_TEXT_PART["text"], + ) + + with patch.object(guardrail.async_handler, "post", _armor_post_mock(_MODEL_ARMOR_DEIDENTIFIED)): + delivered = await _drain_surface_hook(guardrail, (first_delta, second_delta, terminal)) + + assert b"4111-1111-1111-1111" not in _delivered_bytes(delivered) + assert [delivered[0].delta, delivered[1].delta] == ["my card is [REDACTED]", ""] + assert [part["text"] for part in delivered[-1].response.output[0]["content"]] == ["my card is [REDACTED]", ""]