scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 89

fix(model_armor): restore reference attachments via skip_unscannable_attachments and remove the attachment count cap (#33554)

Browse source 
* fix(model_armor): add skip_unscannable_attachments to allow reference-only attachments through

* fix(model_armor): wire skip_unscannable_attachments through guardrail config

* fix(model_armor): make max_file_attachments configurable and scan overflow instead of dropping

* fix(model_armor): remove the per-request attachment count cap and scan all attachments

---------

Co-authored-by: yucheng <yucheng@berri.ai>
This commit is contained in:
devin-ai-integration[bot] 2026-07-16 16:06:41 -07:00 • committed by GitHub
parent 3459956fd2
commit 0d7b0f708b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 168 additions and 34 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm/proxy/guardrails/guardrail_hooks/model_armor/__init__.py
View file

@ -26,6 +26,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"
mask_request_content=litellm_params.mask_request_content,
mask_response_content=litellm_params.mask_response_content,
fail_on_error=litellm_params.fail_on_error,
skip_unscannable_attachments=litellm_params.skip_unscannable_attachments,
)
litellm.logging_callback_manager.add_litellm_callback(_model_armor_callback)

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

@ -25,10 +25,6 @@ from litellm.types.llms.openai import AllMessageValues
MODEL_ARMOR_MAX_FILE_SIZE_BYTES = 4 * 1024 * 1024
# Hard cap on how many attachments a single request may submit to Model Armor, to bound
# per-request fan-out (latency and quota).
MAX_FILE_ATTACHMENTS_PER_REQUEST = 10
_REMOTE_URI_SCHEMES = ("gs://", "http://", "https://")
ModelArmorByteDataType = Literal["PDF", "WORD_DOCUMENT", "EXCEL_DOCUMENT", "POWERPOINT_DOCUMENT", "CSV", "TXT"]

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

@ -32,7 +32,6 @@ from litellm.llms.custom_httpx.http_handler import (
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.guardrails.guardrail_hooks.model_armor.file_scanning import (
MAX_FILE_ATTACHMENTS_PER_REQUEST,
MODEL_ARMOR_MAX_FILE_SIZE_BYTES,
plan_file_scans,
)
@ -383,10 +382,14 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
Each attachment is sent through the byte API and a MATCH_FOUND raises a 400 before the
request reaches the LLM. File scanning does not support masking (Model Armor returns
findings, not a sanitized document), so it only blocks. Anything the guardrail cannot
scan - a file_id or remote URL reference with no inline bytes, a document over the 4 MB
byte limit, or more attachments than the per-request cap - is a guardrail failure and
blocks unless the operator has opted into fail-open via fail_on_error=False.
findings, not a sanitized document), so it only blocks. A file_id or remote URL reference
with no inline bytes and a document over the 4 MB byte limit are guardrail failures that
block unless the operator has opted into fail-open via fail_on_error=False.
skip_unscannable_attachments decouples reference-only attachments from fail_on_error: when
enabled, attachments Model Armor cannot scan (file_id, gs://, or http(s) references with no
inline bytes, and inline content whose base64 will not decode) pass through instead of
blocking, while fail_on_error still governs real Model Armor API errors.
"""
from litellm.proxy.common_utils.callback_utils import (
_get_or_create_proxy_metadata_bucket,
@ -395,7 +398,14 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
plan = plan_file_scans(messages)
attachments = plan.attachments
unscannable_references = plan.unscannable_count
skip_unscannable = bool(self.optional_params.get("skip_unscannable_attachments", False))
if skip_unscannable and plan.unscannable_count > 0:
verbose_proxy_logger.warning(
"Model Armor: allowing %d unscannable attachment(s) through because "
"skip_unscannable_attachments is enabled",
plan.unscannable_count,
)
unscannable_references = 0 if skip_unscannable else plan.unscannable_count
if not attachments and unscannable_references == 0:
return
@ -415,14 +425,6 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
metadata["_model_armor_status"] = "blocked"
raise self._unscannable_block_error(reason)
if len(attachments) > MAX_FILE_ATTACHMENTS_PER_REQUEST:
reason = f"{len(attachments)} attachments exceed the per-request scan limit of {MAX_FILE_ATTACHMENTS_PER_REQUEST}"
verbose_proxy_logger.warning("Model Armor: %s", reason)
if fail_on_error:
metadata["_model_armor_status"] = "blocked"
raise self._unscannable_block_error(reason)
attachments = attachments[:MAX_FILE_ATTACHMENTS_PER_REQUEST]
for attachment in attachments:
if len(attachment.file_bytes) > MODEL_ARMOR_MAX_FILE_SIZE_BYTES:
reason = (

8
litellm/types/guardrails.py
View file

@ -800,6 +800,14 @@ class BaseLitellmParams(ContentFilterConfigModel): # works for new and patch up
"so only a valid guardrail response can block or modify it."
),
)
skip_unscannable_attachments: Optional[bool] = Field(
default=False,
description=(
"Implemented by guardrail='model_armor'. When True, attachment references that carry no "
"inline bytes (file_id, gs://, or http(s) URLs) pass through unscanned instead of blocking, "
"while fail_on_error still governs real Model Armor API errors. Default False blocks them."
),
)
additional_provider_specific_params: Optional[Dict[str, Any]] = Field(
default=None,

159
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py
View file

@ -2205,21 +2205,20 @@ async def test_pre_call_file_id_reference_skipped_when_fail_open():
@pytest.mark.asyncio
async def test_pre_call_blocks_when_attachment_count_exceeds_cap():
"""More attachments than the per-request cap fail closed by default to bound scan fan-out."""
from litellm.proxy.guardrails.guardrail_hooks.model_armor.file_scanning import (
MAX_FILE_ATTACHMENTS_PER_REQUEST,
)
guardrail = _make_guardrail()
pdf_b64 = base64.b64encode(PDF_BYTES).decode("utf-8")
block = {
"type": "file",
"file": {"file_data": f"data:application/pdf;base64,{pdf_b64}"},
}
async def test_pre_call_file_id_reference_passthrough_when_skip_unscannable_enabled():
"""skip_unscannable_attachments lets a file_id reference through even with fail_on_error=True."""
guardrail = _make_guardrail(skip_unscannable_attachments=True)
request_data = {
"model": "gpt-4",
"messages": [{"role": "user", "content": [block] * (MAX_FILE_ATTACHMENTS_PER_REQUEST + 1)}],
"messages": [
{
"role": "user",
"content": [
{"type": "text", "text": "summarize this"},
{"type": "file", "file": {"file_id": "file-abc123"}},
],
}
],
"metadata": {"guardrails": ["model-armor-test"]},
}
@ -2227,8 +2226,109 @@ async def test_pre_call_blocks_when_attachment_count_exceeds_cap():
guardrail.async_handler,
"post",
AsyncMock(return_value=_armor_response(blocked=False)),
) as mock_post:
await guardrail.async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(),
cache=MagicMock(spec=DualCache),
data=request_data,
call_type="completion",
)
assert _byte_items_sent(mock_post) == []
assert _text_payloads_sent(mock_post) == ["summarize this"]
@pytest.mark.asyncio
async def test_pre_call_gs_uri_reference_passthrough_when_skip_unscannable_enabled():
"""A gs:// document reference passes through when skip_unscannable_attachments is enabled."""
guardrail = _make_guardrail(skip_unscannable_attachments=True)
request_data = {
"model": "gpt-4",
"messages": [
{
"role": "user",
"content": [
{
"type": "file",
"file": {"file_data": "gs://my-bucket/report.pdf", "filename": "report.pdf"},
}
],
}
],
"metadata": {"guardrails": ["model-armor-test"]},
}
with patch.object(
guardrail.async_handler,
"post",
AsyncMock(return_value=_armor_response(blocked=False)),
) as mock_post:
await guardrail.async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(),
cache=MagicMock(spec=DualCache),
data=request_data,
call_type="completion",
)
assert _byte_items_sent(mock_post) == []
def test_initialize_guardrail_forwards_skip_unscannable_attachments():
"""skip_unscannable_attachments configured in litellm_params reaches the guardrail instance."""
from litellm.proxy.guardrails.guardrail_hooks.model_armor import initialize_guardrail
from litellm.types.guardrails import Guardrail, LitellmParams
litellm_params = LitellmParams(
guardrail="model_armor",
mode="pre_call",
template_id="demo-template",
project_id="demo-project",
skip_unscannable_attachments=True,
)
guardrail = initialize_guardrail(
litellm_params=litellm_params,
guardrail=Guardrail(guardrail_name="model-armor-config-test"),
)
assert guardrail.optional_params.get("skip_unscannable_attachments") is True
def test_initialize_guardrail_skip_unscannable_defaults_false():
"""A config that omits skip_unscannable_attachments keeps the secure default (block)."""
from litellm.proxy.guardrails.guardrail_hooks.model_armor import initialize_guardrail
from litellm.types.guardrails import Guardrail, LitellmParams
litellm_params = LitellmParams(
guardrail="model_armor",
mode="pre_call",
template_id="demo-template",
project_id="demo-project",
)
guardrail = initialize_guardrail(
litellm_params=litellm_params,
guardrail=Guardrail(guardrail_name="model-armor-config-default"),
)
assert guardrail.optional_params.get("skip_unscannable_attachments") is False
@pytest.mark.asyncio
async def test_skip_unscannable_still_fails_closed_on_api_error():
"""skip_unscannable_attachments only affects references; a real API error still fails closed."""
guardrail = _make_guardrail(skip_unscannable_attachments=True, fail_on_error=True)
pdf_b64 = base64.b64encode(PDF_BYTES).decode("utf-8")
request_data = {
"model": "gpt-4",
"messages": [_file_message(pdf_b64)],
"metadata": {"guardrails": ["model-armor-test"]},
}
with patch.object(
guardrail.async_handler,
"post",
AsyncMock(side_effect=Exception("model armor upstream 500")),
):
with pytest.raises(HTTPException) as exc_info:
with pytest.raises(Exception) as exc_info:
await guardrail.async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(),
cache=MagicMock(spec=DualCache),
@ -2236,8 +2336,35 @@ async def test_pre_call_blocks_when_attachment_count_exceeds_cap():
call_type="completion",
)
assert exc_info.value.status_code == 400
assert "per-request scan limit" in str(exc_info.value.detail)
assert "model armor upstream 500" in str(exc_info.value)
@pytest.mark.asyncio
async def test_pre_call_scans_every_attachment_without_a_count_cap():
"""There is no per-request attachment cap: every scannable attachment is submitted to Model Armor."""
guardrail = _make_guardrail()
pdf_b64 = base64.b64encode(PDF_BYTES).decode("utf-8")
block = {
"type": "file",
"file": {"file_data": f"data:application/pdf;base64,{pdf_b64}"},
}
count = 25
request_data = {
"model": "gpt-4",
"messages": [{"role": "user", "content": [block] * count}],
"metadata": {"guardrails": ["model-armor-test"]},
}
mock_post = AsyncMock(return_value=_armor_response(blocked=False))
with patch.object(guardrail.async_handler, "post", mock_post):
await guardrail.async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(),
cache=MagicMock(spec=DualCache),
data=request_data,
call_type="completion",
)
assert len(_byte_items_sent(mock_post)) == count
@pytest.mark.asyncio