scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-01 07:03:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 128

fix(guardrails): write Model Armor redactions into a single-part Responses body

Browse source 
Model Armor returns one sanitized string for everything it scanned, so a Responses
body answered 400 and a stream was blocked whenever masking had to rewrite it. When
the scan came from one content part, write the string into that part and sync the
buffered stream events. A body that joined several parts or tool-call arguments has
no single place for the string and stays blocked.
This commit is contained in:
mateo-berri 2026-10-30 15:45:37 -07:00
parent 7b13c88400
commit 1108ac2ab6
3 changed files with 213 additions and 26 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

16
litellm/llms/openai/responses/guardrail_translation/handler.py
View file

@ -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):

90
litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py
View file

@ -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)

133
tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py
View file

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