From 349120a9f3553afddf4a3a9c313243ebfc95cf56 Mon Sep 17 00:00:00 2001 From: mateo <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 12 Nov 2026 14:09:37 +0000 Subject: [PATCH] fix(model-armor): fail closed on remote s3 documents Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> (cherry picked from commit f9ad9f8aad6e9df611d8b4f470aeacee7ac0ff85) --- .../model_armor/file_scanning.py | 4 +- .../guardrail_hooks/test_model_armor.py | 51 +++++++++++++++++-- 2 files changed, 48 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/file_scanning.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/file_scanning.py index 366ea7326ae..931aacbe8be 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/file_scanning.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/file_scanning.py @@ -7,7 +7,7 @@ maps its MIME type to a Model Armor ``byteDataType``, and returns the decoded by guardrail hooks can submit them. ``plan_file_scans`` classifies each block: blocks with no inline bytes (``file_id`` or remote -``gs://`` / ``http(s)`` references) and supported documents whose base64 will not decode are +``gs://`` / ``s3://`` / ``http(s)`` references) and supported documents whose base64 will not decode are reported as unscannable so the guardrail hook can fail closed (blocking unless ``fail_on_error`` is false) rather than letting an unscanned document reach the model. """ @@ -27,7 +27,7 @@ from litellm.types.llms.openai import AllMessageValues MODEL_ARMOR_MAX_FILE_SIZE_BYTES: Final = 4 * 1024 * 1024 -_REMOTE_URI_SCHEMES: Final = ("gs://", "http://", "https://") +_REMOTE_URI_SCHEMES: Final = ("gs://", "s3://", "http://", "https://") ModelArmorByteDataType = Literal["PDF", "WORD_DOCUMENT", "EXCEL_DOCUMENT", "POWERPOINT_DOCUMENT", "CSV", "TXT"] diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py index 5aa3c7dfc17..c0901f83853 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py @@ -3,13 +3,11 @@ import base64 import io import json from collections.abc import Iterator, Sequence -from typing import cast +from typing import Final, cast from unittest.mock import AsyncMock, MagicMock, Mock, patch -import pytest - - import httpx +import pytest from fastapi import HTTPException import litellm @@ -17,13 +15,18 @@ import litellm.types.utils from litellm._logging import verbose_proxy_logger from litellm.caching import DualCache from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, MaskedHTTPStatusError -from litellm.proxy.guardrails.anthropic_sse import anthropic_sse_error_frames from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.guardrails.anthropic_sse import anthropic_sse_error_frames from litellm.proxy.guardrails.guardrail_hooks.model_armor import ModelArmorGuardrail +from litellm.proxy.guardrails.guardrail_hooks.model_armor.file_scanning import ( + FileScanPlan, + plan_file_scans, +) from litellm.proxy.guardrails.guardrail_hooks.model_armor.model_armor import ( ModelArmorAPIError, ) from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.llms.openai import AllMessageValues @pytest.mark.asyncio @@ -2670,6 +2673,44 @@ async def test_pre_call_blocks_unresolvable_file_id_reference(): assert "could not scan" in str(exc_info.value.detail) +@pytest.mark.asyncio +@pytest.mark.parametrize( + "uri", + [ + "s3://listed-bucket/report.pdf", + " S3://listed-bucket/report.pdf ", + "gs://listed-bucket/report.pdf", + "https://example.com/report.pdf", + "http://example.com/report.pdf", + ], +) +async def test_pre_call_blocks_remote_document_without_format_or_filename(uri: str) -> None: + messages: Final[list[AllMessageValues]] = [ + {"role": "user", "content": [{"type": "file", "file": {"file_data": uri}}]} + ] + assert plan_file_scans(messages) == FileScanPlan(attachments=(), unscannable_count=1) + + guardrail: Final = ModelArmorGuardrail( + template_id="test-template", + project_id="test-project", + guardrail_name="model-armor-test", + fail_on_error=True, + ) + data: Final = {"messages": messages, "metadata": {"guardrails": ["model-armor-test"]}} + + with pytest.raises(HTTPException) as exc_info: + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(), + data=data, + call_type="completion", + ) + + assert exc_info.value.status_code == 400 + assert "could not scan" in str(exc_info.value.detail) + assert data["metadata"]["_model_armor_status"] == "blocked" + + @pytest.mark.asyncio async def test_pre_call_blocks_remote_url_document_reference(): """A remote (https) document reference cannot be fetched here, so it fails closed by default."""