fix(guardrails): scan file and document attachments with Model Armor (#31655)

The Model Armor guardrail only sent text extracted from user messages to
sanitizeUserPrompt, so harmful content inside attached PDFs, Office docs,
and CSVs reached the LLM unscanned. A file-only message had no extractable
text, so the pre-call and moderation hooks returned early and the document
was never submitted to Model Armor at all.

Wire inline document/file scanning into async_pre_call_hook and
async_moderation_hook. extract_file_attachments walks message content blocks
(OpenAI type:file file_data and Anthropic type:document source), decodes the
base64 bytes, maps the MIME type to a Model Armor byteDataType, and skips
remote URLs, bare file_id references, oversize files past the 4 MB limit, and
unsupported types. Each attachment is sent through the byte API and a
MATCH_FOUND blocks the request before it reaches the LLM.

Resolves LIT-4084

litellm/proxy/guardrails/guardrail_hooks/model_armor/file_scanning.py

@@ -0,0 +1,234 @@
+"""Resolve inline file/document attachments in chat messages to Model Armor byte payloads.
+
+Model Armor scans documents through its ``byteItem`` API (PDF, Office docs, CSV, plaintext).
+This module walks message content blocks (``type: file`` with inline ``file_data`` and
+``type: document`` with an inline base64 ``source``), validates each block into a typed model,
+maps its MIME type to a Model Armor ``byteDataType``, and returns the decoded bytes so the
+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
+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.
+"""
+
+import base64
+import binascii
+import mimetypes
+from dataclasses import dataclass
+from typing import Annotated, Literal, Sequence
+
+from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError
+
+from litellm._logging import verbose_proxy_logger
+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"]
+
+_MIME_TO_BYTE_DATA_TYPE: tuple[tuple[str, ModelArmorByteDataType], ...] = (
+    ("application/pdf", "PDF"),
+    # Word family: legacy, OOXML, macro-enabled, and templates all map to WORD_DOCUMENT
+    ("application/msword", "WORD_DOCUMENT"),
+    ("application/vnd.openxmlformats-officedocument.wordprocessingml.document", "WORD_DOCUMENT"),
+    ("application/vnd.openxmlformats-officedocument.wordprocessingml.template", "WORD_DOCUMENT"),
+    ("application/vnd.ms-word.document.macroenabled.12", "WORD_DOCUMENT"),
+    ("application/vnd.ms-word.template.macroenabled.12", "WORD_DOCUMENT"),
+    # Excel family
+    ("application/vnd.ms-excel", "EXCEL_DOCUMENT"),
+    ("application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", "EXCEL_DOCUMENT"),
+    ("application/vnd.openxmlformats-officedocument.spreadsheetml.template", "EXCEL_DOCUMENT"),
+    ("application/vnd.ms-excel.sheet.macroenabled.12", "EXCEL_DOCUMENT"),
+    ("application/vnd.ms-excel.template.macroenabled.12", "EXCEL_DOCUMENT"),
+    # PowerPoint family
+    ("application/vnd.ms-powerpoint", "POWERPOINT_DOCUMENT"),
+    ("application/vnd.openxmlformats-officedocument.presentationml.presentation", "POWERPOINT_DOCUMENT"),
+    ("application/vnd.openxmlformats-officedocument.presentationml.template", "POWERPOINT_DOCUMENT"),
+    ("application/vnd.openxmlformats-officedocument.presentationml.slideshow", "POWERPOINT_DOCUMENT"),
+    ("application/vnd.ms-powerpoint.presentation.macroenabled.12", "POWERPOINT_DOCUMENT"),
+    ("application/vnd.ms-powerpoint.template.macroenabled.12", "POWERPOINT_DOCUMENT"),
+    ("application/vnd.ms-powerpoint.slideshow.macroenabled.12", "POWERPOINT_DOCUMENT"),
+    ("text/csv", "CSV"),
+    ("text/plain", "TXT"),
+)
+
+
+@dataclass(frozen=True, slots=True)
+class ModelArmorFileAttachment:
+    file_bytes: bytes
+    byte_data_type: ModelArmorByteDataType
+
+
+@dataclass(frozen=True, slots=True)
+class FileScanPlan:
+    # Decoded attachments ready to submit to Model Armor.
+    attachments: tuple[ModelArmorFileAttachment, ...]
+    # Document/file blocks the guardrail recognized but could not turn into scannable bytes
+    # (file_id/remote references, or a supported type whose inline base64 failed to decode).
+    unscannable_count: int
+
+
+class _FileData(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+    file_data: str | None = None
+    format: str | None = None
+    filename: str | None = None
+
+
+class _FileBlock(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+    type: Literal["file"]
+    file: _FileData
+
+
+class _DocumentSource(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+    data: str | None = None
+    media_type: str | None = None
+
+
+class _DocumentBlock(BaseModel):
+    model_config = ConfigDict(extra="ignore")
+    type: Literal["document"]
+    source: _DocumentSource
+
+
+_AttachmentBlock = Annotated[_FileBlock | _DocumentBlock, Field(discriminator="type")]
+_BLOCK_ADAPTER: TypeAdapter[_FileBlock | _DocumentBlock] = TypeAdapter(_AttachmentBlock)
+
+
+def plan_file_scans(messages: Sequence[AllMessageValues]) -> FileScanPlan:
+    """Classify every document/file block into scannable attachments vs unscannable ones.
+
+    Unscannable covers references with no inline bytes and supported documents whose inline
+    base64 fails to decode; the hook fails closed on these. Inline content of an unsupported
+    type (for example an image) is neither scanned nor counted, it is simply left alone.
+    """
+    classified = tuple(_classify_block(block) for message in messages for block in _content_blocks(message))
+    attachments = tuple(attachment for attachment, _ in classified if attachment is not None)
+    unscannable_count = sum(1 for attachment, is_unscannable in classified if attachment is None and is_unscannable)
+    return FileScanPlan(attachments=attachments, unscannable_count=unscannable_count)
+
+
+def _content_blocks(message: AllMessageValues) -> tuple[object, ...]:
+    content = message.get("content")
+    return tuple(content) if isinstance(content, list) else ()
+
+
+def _classify_block(block: object) -> tuple[ModelArmorFileAttachment | None, bool]:
+    """Return (attachment, is_unscannable). At most one is meaningful; (None, False) means skip."""
+    parsed = _parse_block(block)
+    if parsed is None:
+        return None, False
+    if _is_reference(parsed):
+        return None, True
+
+    byte_data_type, data = _block_byte_data_type_and_data(parsed)
+    if data is None:
+        return None, True
+    if byte_data_type is None:
+        # Recognized inline content of a type Model Armor's byte API does not scan (e.g. an image).
+        return None, False
+
+    decoded = _safe_b64decode(data)
+    if decoded is None:
+        # A supported document whose base64 will not decode cannot be scanned, so fail closed.
+        return None, True
+
+    return ModelArmorFileAttachment(file_bytes=decoded, byte_data_type=byte_data_type), False
+
+
+def _is_reference(block: _FileBlock | _DocumentBlock) -> bool:
+    if isinstance(block, _DocumentBlock):
+        return not block.source.data
+    raw = block.file.file_data
+    return not raw or _is_remote_uri(raw)
+
+
+def _parse_block(block: object) -> _FileBlock | _DocumentBlock | None:
+    try:
+        return _BLOCK_ADAPTER.validate_python(block)
+    except ValidationError:
+        return None
+
+
+def _block_byte_data_type_and_data(
+    block: _FileBlock | _DocumentBlock,
+) -> tuple[ModelArmorByteDataType | None, str | None]:
+    if isinstance(block, _DocumentBlock):
+        return _mime_to_byte_data_type(block.source.media_type), block.source.data
+
+    raw = block.file.file_data
+    if not raw:
+        return None, None
+    uri_mime, data = _parse_data_uri(raw)
+    if data is None:
+        data = raw
+    # The data URI header is the least reliable signal: it can be generic (application/octet-stream)
+    # or mislabeled (text/plain for a PDF). Prefer the explicit format and filename, falling back to
+    # the header only when neither resolves, and warn rather than let a conflicting header downgrade a
+    # recognized document to the wrong filter.
+    declared = _first_supported_byte_data_type((block.file.format, _mime_from_filename(block.file.filename)))
+    header = _mime_to_byte_data_type(uri_mime)
+    if declared is None:
+        return header, data
+    if header is not None and header != declared:
+        verbose_proxy_logger.warning(
+            "Model Armor: data URI MIME %s maps to %s but the attachment declares %s; scanning as %s",
+            uri_mime,
+            header,
+            declared,
+            declared,
+        )
+    return declared, data
+
+
+def _first_supported_byte_data_type(
+    mimes: tuple[str | None, ...],
+) -> ModelArmorByteDataType | None:
+    return next(
+        (byte_data_type for mime in mimes for byte_data_type in (_mime_to_byte_data_type(mime),) if byte_data_type),
+        None,
+    )
+
+
+def _parse_data_uri(raw: str) -> tuple[str | None, str | None]:
+    if not raw.startswith("data:") or ";base64," not in raw:
+        return None, None
+    header, data = raw.split(";base64,", 1)
+    return header[len("data:") :] or None, data
+
+
+def _mime_to_byte_data_type(mime: str | None) -> ModelArmorByteDataType | None:
+    if mime is None:
+        return None
+    normalized = mime.split(";")[0].strip().lower()
+    return next(
+        (byte_data_type for candidate, byte_data_type in _MIME_TO_BYTE_DATA_TYPE if candidate == normalized), None
+    )
+
+
+def _mime_from_filename(filename: str | None) -> str | None:
+    if filename is None:
+        return None
+    guessed, _ = mimetypes.guess_type(filename)
+    return guessed
+
+
+def _safe_b64decode(data: str) -> bytes | None:
+    try:
+        return base64.b64decode(data, validate=True)
+    except (binascii.Error, ValueError):
+        verbose_proxy_logger.warning("Model Armor: skipping attachment with undecodable base64 content")
+        return None
+
+
+def _is_remote_uri(raw: str) -> bool:
+    return raw.strip().lower().startswith(_REMOTE_URI_SCHEMES)

litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py

@@ -4,7 +4,9 @@ from typing import (
     AsyncGenerator,
     List,
     Literal,
+    Mapping,
     Optional,
+    Sequence,
     Type,
     Union,
 )
@@ -29,7 +31,13 @@ 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,
+)
 from litellm.types.guardrails import GuardrailEventHooks
+from litellm.types.llms.openai import AllMessageValues
 from litellm.types.utils import (
     CallTypesLiteral,
     Choices,
@@ -166,11 +174,21 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
             "Authorization": f"Bearer {access_token}",
         }
 
-        verbose_proxy_logger.debug(
-            "Model Armor request - URL: %s, Body: %s",
-            url,
-            body,
-        )
+        # Never log byteData: it is the full base64 of the scanned document. Log only its
+        # type and size so debug deployments cannot leak the contents the guardrail inspects.
+        if file_bytes is not None and file_type is not None:
+            verbose_proxy_logger.debug(
+                "Model Armor file request - URL: %s, byteDataType: %s, bytes: %d",
+                url,
+                file_type,
+                len(file_bytes),
+            )
+        else:
+            verbose_proxy_logger.debug(
+                "Model Armor request - URL: %s, Body: %s",
+                url,
+                body,
+            )
 
         # Make request
         if self.async_handler is None:
@@ -293,6 +311,21 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
         # Fallback: if Model Armor put sanitized text at the root, use it
         return armor_response.get("sanitizedText") or armor_response.get("text")
 
+    @staticmethod
+    def _append_armor_response(existing: object, armor_response: Mapping[str, object]) -> object:
+        """Accumulate scan responses so a later text scan does not drop an earlier file scan.
+
+        Returns the single response on its own (backward compatible) and a list once a request
+        carries more than one scan. A list (not a tuple) is required because the guardrail logging
+        pipeline (redact_nested_match_and_regex_keys and the StandardLoggingGuardrailInformation
+        dict | list[dict] contract) only recurses into dicts and lists when redacting and serializing.
+        """
+        if existing is None:
+            return armor_response
+        if isinstance(existing, list):
+            return [*existing, armor_response]  # mutable-ok: logging pipeline requires list[dict], not tuple
+        return [existing, armor_response]  # mutable-ok: logging pipeline requires list[dict], not tuple
+
     def _process_response(
         self,
         response: Optional[dict],
@@ -326,6 +359,108 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
         )
         return response
 
+    @staticmethod
+    def _unscannable_block_error(reason: str) -> HTTPException:
+        return HTTPException(
+            status_code=400,
+            detail={"error": f"Model Armor could not scan an attachment and blocked the request: {reason}"},
+        )
+
+    async def _scan_request_files(self, messages: Sequence[AllMessageValues], data: dict) -> None:
+        """Submit inline document/file attachments to Model Armor and block on any findings.
+
+        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.
+        """
+        from litellm.proxy.common_utils.callback_utils import (
+            _get_or_create_proxy_metadata_bucket,
+            add_guardrail_to_applied_guardrails_header,
+        )
+
+        plan = plan_file_scans(messages)
+        attachments = plan.attachments
+        unscannable_references = plan.unscannable_count
+        if not attachments and unscannable_references == 0:
+            return
+
+        add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
+        # Use the same metadata bucket the header helper writes to, so the logged Model Armor
+        # payload and status land where _process_response reads them on every route.
+        _, metadata = _get_or_create_proxy_metadata_bucket(data)
+        fail_on_error = bool(self.optional_params.get("fail_on_error", True))
+
+        if unscannable_references > 0:
+            reason = (
+                f"{unscannable_references} attachment(s) reference a document with no inline bytes "
+                "(file_id or remote URL) that Model Armor cannot scan"
+            )
+            verbose_proxy_logger.warning("Model Armor: %s", reason)
+            if fail_on_error:
+                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 = (
+                    f"attachment of {len(attachment.file_bytes)} bytes exceeds Model Armor's "
+                    f"{MODEL_ARMOR_MAX_FILE_SIZE_BYTES} byte scan limit"
+                )
+                verbose_proxy_logger.warning("Model Armor: %s", reason)
+                if not fail_on_error:
+                    continue
+                metadata["_model_armor_status"] = "blocked"
+                raise self._unscannable_block_error(reason)
+
+            try:
+                armor_response = await self.make_model_armor_request(
+                    source="user_prompt",
+                    request_data=data,
+                    file_bytes=attachment.file_bytes,
+                    file_type=attachment.byte_data_type,
+                )
+            except HTTPException:
+                raise
+            except Exception as e:
+                # Isolate transient errors per attachment so one failure does not leave the
+                # remaining attachments in the same request unscanned.
+                verbose_proxy_logger.error("Model Armor file scan error: %s", str(e), exc_info=True)
+                if fail_on_error:
+                    raise
+                continue
+
+            # Model Armor returns findings for documents, not a sanitized file, so there is no
+            # masking fallback. Any finding must block, even when mask_request_content is enabled,
+            # otherwise a PII-only (SDP deidentify) document would pass through unscrubbed.
+            blocked = self._should_block_content(armor_response, allow_sanitization=False)
+            metadata["_model_armor_response"] = self._append_armor_response(
+                metadata.get("_model_armor_response"), armor_response
+            )
+            if blocked or metadata.get("_model_armor_status") == "blocked":
+                metadata["_model_armor_status"] = "blocked"
+            else:
+                metadata["_model_armor_status"] = "success"
+
+            if blocked:
+                raise HTTPException(
+                    status_code=400,
+                    detail={
+                        "error": "Content blocked by Model Armor",
+                        "model_armor_response": armor_response,
+                    },
+                )
+
     @log_guardrail_information
     async def async_pre_call_hook(
         self,
@@ -355,6 +490,8 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
             get_last_user_message,
         )
 
+        await self._scan_request_files(messages=messages, data=data)
+
         content = get_last_user_message(messages)
         if not content:
             return data
@@ -372,24 +509,27 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
             #   race-conditions between concurrent requests which share the same guardrail instance.
             #   This ensures each request logs its own Model Armor response instead of a potentially stale value
             #   overwritten by another coroutine.
+            blocked = self._should_block_content(armor_response, allow_sanitization=self.mask_request_content)
             if isinstance(data, dict):
                 metadata = data.setdefault("metadata", {})  # ensures metadata exists and is unique per request
-                metadata["_model_armor_response"] = armor_response
+                # Accumulate so a prior file scan on the same request is not overwritten by this text scan.
+                metadata["_model_armor_response"] = self._append_armor_response(
+                    metadata.get("_model_armor_response"), armor_response
+                )
                 # Pre-compute guardrail status for downstream logging. A blocked response will eventually raise
                 #   an HTTPException, however in scenarios where the caller decides to ignore the exception (e.g.
                 #   fail_on_error=False) we still want the correct status reflected.
-                metadata["_model_armor_status"] = (
-                    "blocked"
-                    if self._should_block_content(armor_response, allow_sanitization=self.mask_request_content)
-                    else "success"
-                )
+                if blocked or metadata.get("_model_armor_status") == "blocked":
+                    metadata["_model_armor_status"] = "blocked"
+                else:
+                    metadata["_model_armor_status"] = "success"
 
             # Add guardrail to applied_guardrails BEFORE potential blocking
             # This ensures guardrail is recorded even when it blocks the request
             add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
 
             # Check if content should be blocked
-            if self._should_block_content(armor_response, allow_sanitization=self.mask_request_content):
+            if blocked:
                 raise HTTPException(
                     status_code=400,
                     detail={
@@ -447,6 +587,8 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
             get_last_user_message,
         )
 
+        await self._scan_request_files(messages=messages, data=data)
+
         content = get_last_user_message(messages)
         if not content:
             return data
@@ -459,22 +601,25 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
                 request_data=data,
             )
 
+            blocked = self._should_block_content(armor_response, allow_sanitization=self.mask_request_content)
             # Store the armor response for logging
             if isinstance(data, dict):
                 metadata = data.setdefault("metadata", {})
-                metadata["_model_armor_response"] = armor_response
-                metadata["_model_armor_status"] = (
-                    "blocked"
-                    if self._should_block_content(armor_response, allow_sanitization=self.mask_request_content)
-                    else "success"
+                # Accumulate so a prior file scan on the same request is not overwritten by this text scan.
+                metadata["_model_armor_response"] = self._append_armor_response(
+                    metadata.get("_model_armor_response"), armor_response
                 )
+                if blocked or metadata.get("_model_armor_status") == "blocked":
+                    metadata["_model_armor_status"] = "blocked"
+                else:
+                    metadata["_model_armor_status"] = "success"
 
             # Add guardrail to applied_guardrails BEFORE potential blocking
             # This ensures guardrail is recorded even when it blocks the request
             add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
 
             # Check if content should be blocked
-            if self._should_block_content(armor_response, allow_sanitization=self.mask_request_content):
+            if blocked:
                 raise HTTPException(
                     status_code=400,
                     detail={