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 only Responses input on Responses calls and cover input_file attachments

Browse source 
On a Responses call the provider receives `input`, so a `messages` field sent
next to it is no longer scanned in its place. Native Responses `input_file`
items are now classified like chat file blocks: inline file_data is scanned,
and a file_id or file_url reference fails closed.
This commit is contained in:
mayank 2026-11-05 00:05:33 +05:30
parent f19ab94049
commit fff74b9de2
3 changed files with 143 additions and 15 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

25
litellm/proxy/guardrails/guardrail_hooks/model_armor/file_scanning.py
View file

@ -86,6 +86,17 @@ class _FileBlock(LiteLLMBaseModel):
file: _FileData
class _ResponsesFileBlock(LiteLLMBaseModel):
"""A Responses API `input_file` item: inline `file_data`, or a `file_id` / `file_url` reference."""
model_config = ConfigDict(extra="ignore")
type: Literal["input_file"]
file_data: str | None = None
file_url: str | None = None
file_id: str | None = None
filename: str | None = None
class _DocumentSource(LiteLLMBaseModel):
model_config = ConfigDict(extra="ignore")
data: str | None = None
@ -98,8 +109,8 @@ class _DocumentBlock(LiteLLMBaseModel):
source: _DocumentSource
_AttachmentBlock = Annotated[_FileBlock | _DocumentBlock, Field(discriminator="type")]
_BLOCK_ADAPTER: Final[TypeAdapter[_FileBlock | _DocumentBlock]] = TypeAdapter(_AttachmentBlock)
_AttachmentBlock = Annotated[_FileBlock | _DocumentBlock | _ResponsesFileBlock, Field(discriminator="type")]
_BLOCK_ADAPTER: Final[TypeAdapter[_FileBlock | _DocumentBlock | _ResponsesFileBlock]] = TypeAdapter(_AttachmentBlock)
def plan_file_scans(messages: Sequence[AllMessageValues]) -> FileScanPlan:
@ -152,9 +163,17 @@ def _is_reference(block: _FileBlock | _DocumentBlock) -> bool:
def _parse_block(block: object) -> _FileBlock | _DocumentBlock | None:
try:
return _BLOCK_ADAPTER.validate_python(block)
parsed: Final = _BLOCK_ADAPTER.validate_python(block)
except ValidationError:
return None
if isinstance(parsed, _ResponsesFileBlock):
# Same shape as a chat `file` block; a file_url or file_id alone has no inline bytes, so it is
# classified as a reference and fails closed like its chat equivalent.
return _FileBlock(
type="file",
file=_FileData(file_data=parsed.file_data or parsed.file_url, filename=parsed.filename),
)
return parsed
def _block_byte_data_type_and_data(

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

@ -155,6 +155,8 @@ class ModelArmorAPIError(Exception):
self.detail = detail
_RESPONSES_CALL_TYPES: Final = frozenset({CallTypes.responses.value, CallTypes.aresponses.value})
_SCANNED_CONTENT_KEYS: Final = frozenset({"text", "sanitizedText", "findings", "maliciousUriMatchedItems"})
@ -702,6 +704,13 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
detail=self._build_block_error_detail("Content blocked by Model Armor", armor_response),
)
@staticmethod
def _scans_responses_input(data: dict, call_type: CallTypesLiteral) -> bool:
"""Whether the scan reads the Responses API `input` rather than `messages`."""
if call_type in _RESPONSES_CALL_TYPES:
return True
return not data.get("messages") and data.get("input") is not None
@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."""
@ -757,9 +766,13 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
if self.should_run_guardrail(data=data, event_type=event_type) is not True:
return data
# 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)
# Responses API requests carry `input` instead of `messages`; scan it as chat messages. For those
# call types `input` is what reaches the provider, so a `messages` field sent alongside it is ignored.
scans_responses_input: Final = self._scans_responses_input(data, call_type)
responses_input: Final = data.get("input") if scans_responses_input else None
messages: Final = (
self._responses_input_messages(responses_input) if scans_responses_input else data.get("messages")
)
if not messages:
verbose_proxy_logger.warning("Model Armor: not running guardrail. No messages in data")
return data
@ -824,7 +837,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
set_last_user_message,
)
if responses_input is not None:
if scans_responses_input:
data["input"] = self._mask_responses_input(responses_input, sanitized_content)
else:
data["messages"] = set_last_user_message(messages, sanitized_content)
@ -861,9 +874,13 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
if self.should_run_guardrail(data=data, event_type=event_type) is not True:
return data
# 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)
# Responses API requests carry `input` instead of `messages`; scan it as chat messages. For those
# call types `input` is what reaches the provider, so a `messages` field sent alongside it is ignored.
scans_responses_input: Final = self._scans_responses_input(data, call_type)
responses_input: Final = data.get("input") if scans_responses_input else None
messages: Final = (
self._responses_input_messages(responses_input) if scans_responses_input else data.get("messages")
)
if not messages:
verbose_proxy_logger.warning("Model Armor: not running guardrail. No messages in data")
return data
@ -920,7 +937,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
set_last_user_message,
)
if responses_input is not None:
if scans_responses_input:
data["input"] = self._mask_responses_input(responses_input, sanitized_content)
else:
data["messages"] = set_last_user_message(messages, sanitized_content)

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

@ -5348,12 +5348,12 @@ def _responses_guardrail(armor_response: dict, **kwargs) -> ModelArmorGuardrail:
return guardrail
async def _run_pre_call(guardrail: ModelArmorGuardrail, data: dict) -> dict:
async def _run_pre_call(guardrail: ModelArmorGuardrail, data: dict, call_type: str = "responses") -> dict:
return await guardrail.async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(),
cache=MagicMock(spec=DualCache),
data=data,
call_type="responses",
call_type=call_type,
)
@ -5428,7 +5428,7 @@ async def test_model_armor_responses_masks_trailing_user_items_and_keeps_other_i
@pytest.mark.asyncio
async def test_model_armor_messages_take_precedence_over_input():
async def test_model_armor_chat_call_scans_messages_over_input():
guardrail = _responses_guardrail(_NO_MATCH)
data = {
"model": "gpt-4",
@ -5437,7 +5437,7 @@ async def test_model_armor_messages_take_precedence_over_input():
"metadata": {},
}
await _run_pre_call(guardrail, data)
await _run_pre_call(guardrail, data, call_type="completion")
assert guardrail.make_model_armor_request.await_args.kwargs["content"] == "chat message"
@ -5465,3 +5465,95 @@ async def test_model_armor_moderation_hook_scans_responses_input():
)
assert exc_info.value.status_code == 400
@pytest.mark.asyncio
@pytest.mark.parametrize("call_type", ["responses", "aresponses"])
async def test_model_armor_responses_call_scans_input_not_a_decoy_messages_field(call_type: str):
"""`input` is what reaches the provider on a Responses call, so a benign `messages` sent next to it is ignored."""
guardrail = _responses_guardrail(_RAI_MATCH)
data = {
"model": "gpt-4",
"messages": [{"role": "user", "content": "hello"}],
"input": "synthetic harmful prompt",
"metadata": {},
}
with pytest.raises(HTTPException) as exc_info:
await _run_pre_call(guardrail, data, call_type=call_type)
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_masking_replaces_input_even_with_a_decoy_messages_field():
guardrail = _responses_guardrail(_MASKED, mask_request_content=True)
data = {
"model": "gpt-4",
"messages": [{"role": "user", "content": "hello"}],
"input": "call me at 555-0100",
"metadata": {},
}
result = await _run_pre_call(guardrail, data)
assert result["input"] == "call me at [REDACTED]"
assert result["messages"] == [{"role": "user", "content": "hello"}]
@pytest.mark.asyncio
async def test_model_armor_responses_inline_input_file_is_scanned():
import base64
guardrail = _responses_guardrail(_RAI_MATCH)
document = b"synthetic document text"
data = {
"model": "gpt-4",
"input": [
{
"role": "user",
"content": [
{"type": "input_text", "text": "summarize this"},
{
"type": "input_file",
"filename": "notes.txt",
"file_data": "data:text/plain;base64," + base64.b64encode(document).decode(),
},
],
}
],
"metadata": {},
}
with pytest.raises(HTTPException) as exc_info:
await _run_pre_call(guardrail, data)
assert exc_info.value.status_code == 400
first_call = guardrail.make_model_armor_request.await_args_list[0].kwargs
assert first_call["file_bytes"] == document
assert first_call["file_type"] == "TXT"
@pytest.mark.asyncio
@pytest.mark.parametrize(
"file_item",
[
{"type": "input_file", "file_id": "file-abc123"},
{"type": "input_file", "file_url": "https://example.test/report.pdf"},
],
)
async def test_model_armor_responses_input_file_reference_fails_closed(file_item: dict):
guardrail = _responses_guardrail(_NO_MATCH)
data = {
"model": "gpt-4",
"input": [{"role": "user", "content": [{"type": "input_text", "text": "summarize this"}, file_item]}],
"metadata": {},
}
with pytest.raises(HTTPException) as exc_info:
await _run_pre_call(guardrail, data)
assert exc_info.value.status_code == 400
guardrail.make_model_armor_request.assert_not_awaited()