mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
Merge fff74b9de2 into 08f2a98250
This commit is contained in:
mac146
GitHub
commit
f1932708dc
3 changed files with 327 additions and 7 deletions
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
@ -154,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"})
|
||||
|
||||
|
||||
|
|
@ -701,6 +704,47 @@ 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."""
|
||||
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 +766,13 @@ 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. 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
|
||||
|
|
@ -787,7 +837,10 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
|
|||
set_last_user_message,
|
||||
)
|
||||
|
||||
data["messages"] = set_last_user_message(messages, sanitized_content)
|
||||
if scans_responses_input:
|
||||
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 +874,13 @@ 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. 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
|
||||
|
|
@ -878,7 +937,10 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
|
|||
set_last_user_message,
|
||||
)
|
||||
|
||||
data["messages"] = set_last_user_message(messages, sanitized_content)
|
||||
if scans_responses_input:
|
||||
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)
|
||||
|
|
|
|||
|
|
@ -5318,3 +5318,242 @@ 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, 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=call_type,
|
||||
)
|
||||
|
||||
|
||||
@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_chat_call_scans_messages_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, call_type="completion")
|
||||
|
||||
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
|
||||
|
||||
|
||||
@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()
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue