scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-15 09:12:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 311

fix(guardrails): scan Responses API input in the Model Armor hooks

Browse source 
Responses API requests carry `input` instead of `messages`. The pre-call
and during-call hooks returned early when `messages` was absent, so those
requests were never sent to Model Armor. Convert `input` to chat messages
with the existing Responses helper and scan it like any other request. With
mask_request_content on, the sanitized text is written back into `input`.

Fixes #49909
This commit is contained in:
mayank 2026-11-04 23:12:03 +05:30
parent f5e740e69a
commit f19ab94049
2 changed files with 196 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -45,6 +45,7 @@ from litellm.proxy.guardrails.guardrail_hooks.model_armor.file_scanning import (
MODEL_ARMOR_MAX_FILE_SIZE_BYTES,
plan_file_scans,
)
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.types.guardrails import GuardrailEventHooks, LitellmParams
from litellm.types.llms.openai import (
AllMessageValues,
@ -701,6 +702,40 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
detail=self._build_block_error_detail("Content blocked by Model Armor", armor_response),
)
@staticmethod
def _responses_input_messages(responses_input: object) -> list[AllMessageValues]:
"""The role-carrying items of a Responses API `input`, in the chat message shape the scan reads."""
if not isinstance(responses_input, (str, list)):
return []
return ResponsesAPIRequestUtils.responses_input_to_chat_messages(responses_input)
@staticmethod
def _mask_responses_input(responses_input: object, sanitized_content: str) -> str | list[object]:
"""Replace the last block of user messages in a Responses API `input` with the sanitized text.
Mirrors set_last_user_message for chat messages: the trailing run of user messages becomes one
user message, and items without a role (tool calls, tool outputs) keep their place.
"""
if isinstance(responses_input, str):
return sanitized_content
items: Final = list(responses_input) if isinstance(responses_input, list) else []
trailing_user: Final[list[int]] = []
for index in range(len(items) - 1, -1, -1):
item = items[index]
if not isinstance(item, dict) or "role" not in item:
continue
if item["role"] != "user":
break
trailing_user.append(index)
if not trailing_user:
return items
first: Final = min(trailing_user)
return [
{"role": "user", "content": sanitized_content} if index == first else item
for index, item in enumerate(items)
if index not in trailing_user or index == first
]
@log_guardrail_information
async def async_pre_call_hook(
self,
@ -722,7 +757,9 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
if self.should_run_guardrail(data=data, event_type=event_type) is not True:
return data
messages: Final = data.get("messages")
# Responses API requests carry `input` instead of `messages`; scan it as chat messages
responses_input: Final = data.get("input") if not data.get("messages") else None
messages: Final = data.get("messages") or self._responses_input_messages(responses_input)
if not messages:
verbose_proxy_logger.warning("Model Armor: not running guardrail. No messages in data")
return data
@ -787,7 +824,10 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
set_last_user_message,
)
data["messages"] = set_last_user_message(messages, sanitized_content)
if responses_input is not None:
data["input"] = self._mask_responses_input(responses_input, sanitized_content)
else:
data["messages"] = set_last_user_message(messages, sanitized_content)
except ModelArmorAPIError as e:
self._raise_if_fail_closed(e)
@ -821,7 +861,9 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
if self.should_run_guardrail(data=data, event_type=event_type) is not True:
return data
messages: Final = data.get("messages")
# Responses API requests carry `input` instead of `messages`; scan it as chat messages
responses_input: Final = data.get("input") if not data.get("messages") else None
messages: Final = data.get("messages") or self._responses_input_messages(responses_input)
if not messages:
verbose_proxy_logger.warning("Model Armor: not running guardrail. No messages in data")
return data
@ -878,7 +920,10 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
set_last_user_message,
)
data["messages"] = set_last_user_message(messages, sanitized_content)
if responses_input is not None:
data["input"] = self._mask_responses_input(responses_input, sanitized_content)
else:
data["messages"] = set_last_user_message(messages, sanitized_content)
except ModelArmorAPIError as e:
self._raise_if_fail_closed(e)

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

@ -5318,3 +5318,150 @@ async def test_apply_guardrail_raises_on_flagged_when_not_logging_only():
entries = request_data["metadata"]["standard_logging_guardrail_information"]
flagged = [e for e in entries if e["guardrail_status"] == "guardrail_flagged"]
assert len(flagged) == 1
_RAI_MATCH = {"sanitizationResult": {"filterResults": {"rai": {"raiFilterResult": {"matchState": "MATCH_FOUND"}}}}}
_NO_MATCH = {"sanitizationResult": {"filterMatchState": "NO_MATCH_FOUND", "filterResults": {}}}
_MASKED = {
"sanitizationResult": {
"filterResults": {
"sdp": {
"sdpFilterResult": {
"deidentifyResult": {"matchState": "MATCH_FOUND", "data": {"text": "call me at [REDACTED]"}}
}
}
}
}
}
def _responses_guardrail(armor_response: dict, **kwargs) -> ModelArmorGuardrail:
guardrail = ModelArmorGuardrail(
template_id="test-template",
project_id="test-project",
location="us-central1",
guardrail_name="model-armor-test",
**kwargs,
)
guardrail.should_run_guardrail = Mock(return_value=True)
guardrail.make_model_armor_request = AsyncMock(return_value=armor_response)
return guardrail
async def _run_pre_call(guardrail: ModelArmorGuardrail, data: dict) -> dict:
return await guardrail.async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(),
cache=MagicMock(spec=DualCache),
data=data,
call_type="responses",
)
@pytest.mark.asyncio
async def test_model_armor_responses_string_input_is_scanned_and_blocked():
"""A Responses API request has `input`, not `messages`, and must not skip the scan."""
guardrail = _responses_guardrail(_RAI_MATCH)
data = {"model": "gpt-4", "input": "synthetic harmful prompt", "metadata": {}}
with pytest.raises(HTTPException) as exc_info:
await _run_pre_call(guardrail, data)
assert exc_info.value.status_code == 400
assert guardrail.make_model_armor_request.await_args.kwargs["content"] == "synthetic harmful prompt"
@pytest.mark.asyncio
async def test_model_armor_responses_list_input_scans_last_user_text():
guardrail = _responses_guardrail(_NO_MATCH)
data = {
"model": "gpt-4",
"input": [
{"role": "user", "content": "first question"},
{"role": "assistant", "content": "an answer"},
{"type": "function_call_output", "call_id": "call_1", "output": "tool text"},
{"role": "user", "content": [{"type": "input_text", "text": "follow-up question"}]},
],
"metadata": {},
}
result = await _run_pre_call(guardrail, data)
assert guardrail.make_model_armor_request.await_args.kwargs["content"] == "follow-up question"
assert result["input"] == data["input"]
@pytest.mark.asyncio
async def test_model_armor_responses_masks_string_input():
guardrail = _responses_guardrail(_MASKED, mask_request_content=True)
data = {"model": "gpt-4", "input": "call me at 555-0100", "metadata": {}}
result = await _run_pre_call(guardrail, data)
assert result["input"] == "call me at [REDACTED]"
assert "messages" not in result
@pytest.mark.asyncio
async def test_model_armor_responses_masks_trailing_user_items_and_keeps_other_items():
guardrail = _responses_guardrail(_MASKED, mask_request_content=True)
tool_output = {"type": "function_call_output", "call_id": "call_1", "output": "tool text"}
data = {
"model": "gpt-4",
"input": [
{"role": "user", "content": "earlier question"},
{"role": "assistant", "content": "an answer"},
{"role": "user", "content": "call me"},
tool_output,
{"role": "user", "content": "at 555-0100"},
],
"metadata": {},
}
result = await _run_pre_call(guardrail, data)
assert result["input"] == [
{"role": "user", "content": "earlier question"},
{"role": "assistant", "content": "an answer"},
{"role": "user", "content": "call me at [REDACTED]"},
tool_output,
]
@pytest.mark.asyncio
async def test_model_armor_messages_take_precedence_over_input():
guardrail = _responses_guardrail(_NO_MATCH)
data = {
"model": "gpt-4",
"messages": [{"role": "user", "content": "chat message"}],
"input": "responses input",
"metadata": {},
}
await _run_pre_call(guardrail, data)
assert guardrail.make_model_armor_request.await_args.kwargs["content"] == "chat message"
@pytest.mark.asyncio
async def test_model_armor_no_messages_and_no_usable_input_skips_scan():
guardrail = _responses_guardrail(_NO_MATCH)
for data in ({"model": "gpt-4", "metadata": {}}, {"model": "gpt-4", "input": None, "metadata": {}}):
assert await _run_pre_call(guardrail, data) is data
guardrail.make_model_armor_request.assert_not_awaited()
@pytest.mark.asyncio
async def test_model_armor_moderation_hook_scans_responses_input():
guardrail = _responses_guardrail(_RAI_MATCH)
data = {"model": "gpt-4", "input": "synthetic harmful prompt", "metadata": {}}
with pytest.raises(HTTPException) as exc_info:
await guardrail.async_moderation_hook(
data=data,
user_api_key_dict=UserAPIKeyAuth(),
call_type="responses",
)
assert exc_info.value.status_code == 400