From 10849c880b4cd7ed279f8c71eb4a4abdcce74cff Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Mon, 29 Jun 2026 19:19:29 -0700 Subject: [PATCH] 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 --- .../model_armor/file_scanning.py | 234 ++++ .../model_armor/model_armor.py | 181 ++- .../guardrail_hooks/test_model_armor.py | 1097 +++++++++++++++++ 3 files changed, 1494 insertions(+), 18 deletions(-) create mode 100644 litellm/proxy/guardrails/guardrail_hooks/model_armor/file_scanning.py diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/file_scanning.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/file_scanning.py new file mode 100644 index 00000000000..0bc6e67eb35 --- /dev/null +++ b/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) diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py index 19b6fa77911..bebd9b28745 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/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={ diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py index 54a8042d4cd..64df9ee7ab5 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py @@ -1,4 +1,5 @@ import asyncio +import base64 import io import json import os @@ -13,6 +14,7 @@ from fastapi import HTTPException import litellm import litellm.types.utils +from litellm._logging import verbose_proxy_logger from litellm.caching import DualCache from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.model_armor import ModelArmorGuardrail @@ -1843,3 +1845,1098 @@ async def test_async_moderation_hook_api_error_fail_on_error_false(): ) assert "API Error" in str(exc_info.value) + + +# ===== FILE / DOCUMENT ATTACHMENT SCANNING TESTS (LIT-4084) ===== + +PDF_BYTES = b"%PDF-1.4\nfake pdf payload with policy-violating content\n%%EOF" +DOCX_BYTES = b"PK\x03\x04 fake docx zip payload" + + +def _make_guardrail(**overrides): + params = dict( + template_id="test-template", + project_id="test-project", + location="us-central1", + guardrail_name="model-armor-test", + ) + params.update(overrides) + guardrail = ModelArmorGuardrail(**params) + guardrail._ensure_access_token_async = AsyncMock(return_value=("test-token", "test-project")) + return guardrail + + +def _armor_response(blocked: bool): + mock_response = AsyncMock() + mock_response.status_code = 200 + if blocked: + body = { + "sanitizationResult": { + "filterMatchState": "MATCH_FOUND", + "filterResults": {"rai": {"raiFilterResult": {"matchState": "MATCH_FOUND"}}}, + } + } + else: + body = {"sanitizationResult": {"filterMatchState": "NO_MATCH_FOUND"}} + mock_response.json = AsyncMock(return_value=body) + return mock_response + + +def _byte_items_sent(mock_post): + """Return every byteItem payload (file scans) submitted to Model Armor.""" + items = [] + for call in mock_post.call_args_list: + body = call.kwargs.get("json", {}) + byte_item = body.get("userPromptData", {}).get("byteItem") + if byte_item is not None: + items.append(byte_item) + return items + + +def _text_payloads_sent(mock_post): + """Return every text payload (text scans) submitted to Model Armor.""" + texts = [] + for call in mock_post.call_args_list: + body = call.kwargs.get("json", {}) + user_prompt = body.get("userPromptData", {}) + if "text" in user_prompt: + texts.append(user_prompt["text"]) + return texts + + +def _file_message(file_data_b64: str, mime: str = "application/pdf", filename: str = "doc.pdf"): + return { + "role": "user", + "content": [ + { + "type": "file", + "file": { + "file_data": f"data:{mime};base64,{file_data_b64}", + "filename": filename, + }, + } + ], + } + + +@pytest.mark.asyncio +async def test_pre_call_blocks_harmful_pdf_attachment(): + """Pre-call hook must scan an inline PDF attachment and block on a Model Armor match. + + Regression for LIT-4084: before the fix, a file-only message has no extractable + text, so the hook returned early and the document was never sent to Model Armor. + """ + guardrail = _make_guardrail() + 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(return_value=_armor_response(blocked=True)), + ) as mock_post: + with pytest.raises(HTTPException) as exc_info: + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=MagicMock(spec=DualCache), + data=request_data, + call_type="completion", + ) + + assert exc_info.value.status_code == 400 + assert "Content blocked by Model Armor" in str(exc_info.value.detail) + + byte_items = _byte_items_sent(mock_post) + assert len(byte_items) == 1 + assert byte_items[0]["byteDataType"] == "PDF" + assert base64.b64decode(byte_items[0]["byteData"]) == PDF_BYTES + + assert "model-armor-test" in request_data["metadata"]["applied_guardrails"] + assert request_data["metadata"]["_model_armor_status"] == "blocked" + + +@pytest.mark.asyncio +async def test_pre_call_allows_safe_pdf_attachment_but_still_scans_it(): + """A safe PDF attachment passes through, but the bytes are still submitted to Model Armor.""" + guardrail = _make_guardrail() + 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(return_value=_armor_response(blocked=False)), + ) as mock_post: + result = await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=MagicMock(spec=DualCache), + data=request_data, + call_type="completion", + ) + + assert result == request_data + byte_items = _byte_items_sent(mock_post) + assert len(byte_items) == 1 + assert byte_items[0]["byteDataType"] == "PDF" + assert base64.b64decode(byte_items[0]["byteData"]) == PDF_BYTES + + +@pytest.mark.asyncio +async def test_moderation_hook_blocks_harmful_file_attachment(): + """The during-call moderation hook must scan file attachments the same way.""" + guardrail = _make_guardrail() + 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(return_value=_armor_response(blocked=True)), + ) as mock_post: + with pytest.raises(HTTPException) as exc_info: + await guardrail.async_moderation_hook( + data=request_data, + user_api_key_dict=UserAPIKeyAuth(), + call_type="completion", + ) + + assert exc_info.value.status_code == 400 + assert "Content blocked by Model Armor" in str(exc_info.value.detail) + assert len(_byte_items_sent(mock_post)) == 1 + assert "model-armor-test" in request_data["metadata"]["applied_guardrails"] + + +@pytest.mark.asyncio +async def test_pre_call_scans_both_text_and_file(): + """When a message has both text and a file, both are submitted to Model Armor.""" + guardrail = _make_guardrail() + pdf_b64 = base64.b64encode(PDF_BYTES).decode("utf-8") + request_data = { + "model": "gpt-4", + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "summarize this"}, + { + "type": "file", + "file": {"file_data": f"data:application/pdf;base64,{pdf_b64}"}, + }, + ], + } + ], + "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 len(_byte_items_sent(mock_post)) == 1 + assert "summarize this" in _text_payloads_sent(mock_post) + + +@pytest.mark.asyncio +async def test_pre_call_scans_anthropic_document_block(): + """Anthropic-style `type: document` blocks with inline base64 are scanned and typed.""" + guardrail = _make_guardrail() + docx_b64 = base64.b64encode(DOCX_BYTES).decode("utf-8") + request_data = { + "model": "gpt-4", + "messages": [ + { + "role": "user", + "content": [ + { + "type": "document", + "source": { + "type": "base64", + "media_type": "application/vnd.openxmlformats-officedocument.wordprocessingml.document", + "data": docx_b64, + }, + } + ], + } + ], + "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", + ) + + byte_items = _byte_items_sent(mock_post) + assert len(byte_items) == 1 + assert byte_items[0]["byteDataType"] == "WORD_DOCUMENT" + assert base64.b64decode(byte_items[0]["byteData"]) == DOCX_BYTES + + +@pytest.mark.asyncio +async def test_pre_call_blocks_unresolvable_file_id_reference(): + """A bare file_id has no inline bytes to scan, so by default the guardrail fails closed.""" + guardrail = _make_guardrail() + request_data = { + "model": "gpt-4", + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "summarize this"}, + {"type": "file", "file": {"file_id": "file-abc123"}}, + ], + } + ], + "metadata": {"guardrails": ["model-armor-test"]}, + } + + with patch.object( + guardrail.async_handler, + "post", + AsyncMock(return_value=_armor_response(blocked=False)), + ): + with pytest.raises(HTTPException) as exc_info: + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=MagicMock(spec=DualCache), + data=request_data, + call_type="completion", + ) + + assert exc_info.value.status_code == 400 + assert "could not scan" in str(exc_info.value.detail) + + +@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.""" + guardrail = _make_guardrail() + request_data = { + "model": "gpt-4", + "messages": [ + { + "role": "user", + "content": [ + { + "type": "file", + "file": {"file_data": "https://example.com/secret.pdf", "filename": "secret.pdf"}, + } + ], + } + ], + "metadata": {"guardrails": ["model-armor-test"]}, + } + + with patch.object( + guardrail.async_handler, + "post", + AsyncMock(return_value=_armor_response(blocked=False)), + ): + with pytest.raises(HTTPException) as exc_info: + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=MagicMock(spec=DualCache), + data=request_data, + call_type="completion", + ) + + assert exc_info.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_pre_call_file_id_reference_skipped_when_fail_open(): + """With fail_on_error=False an unresolvable reference is skipped and the text is still scanned.""" + guardrail = _make_guardrail(fail_on_error=False) + request_data = { + "model": "gpt-4", + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "summarize this"}, + {"type": "file", "file": {"file_id": "file-abc123"}}, + ], + } + ], + "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) == [] + assert _text_payloads_sent(mock_post) == ["summarize this"] + + +@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}"}, + } + request_data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": [block] * (MAX_FILE_ATTACHMENTS_PER_REQUEST + 1)}], + "metadata": {"guardrails": ["model-armor-test"]}, + } + + with patch.object( + guardrail.async_handler, + "post", + AsyncMock(return_value=_armor_response(blocked=False)), + ): + with pytest.raises(HTTPException) as exc_info: + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=MagicMock(spec=DualCache), + data=request_data, + call_type="completion", + ) + + assert exc_info.value.status_code == 400 + assert "per-request scan limit" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +async def test_file_scan_error_isolated_when_fail_open(): + """A transient error on one attachment does not skip the remaining attachments (fail_on_error=False).""" + guardrail = _make_guardrail(fail_on_error=False) + pdf_b64 = base64.b64encode(PDF_BYTES).decode("utf-8") + block = { + "type": "file", + "file": {"file_data": f"data:application/pdf;base64,{pdf_b64}"}, + } + request_data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": [block, block]}], + "metadata": {"guardrails": ["model-armor-test"]}, + } + + # First attachment raises a transient error, second returns a normal response + post = AsyncMock(side_effect=[Exception("transient"), _armor_response(blocked=False)]) + with patch.object(guardrail.async_handler, "post", post): + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=MagicMock(spec=DualCache), + data=request_data, + call_type="completion", + ) + + # Both attachments are attempted: the first errors and is isolated, the second still scans + assert post.call_count == 2 + + +@pytest.mark.asyncio +async def test_pre_call_skips_unsupported_file_type(): + """An image attachment (no Model Armor byteDataType) is not submitted as a document.""" + guardrail = _make_guardrail() + png_b64 = base64.b64encode(b"\x89PNG fake").decode("utf-8") + request_data = { + "model": "gpt-4", + "messages": [_file_message(png_b64, mime="image/png", filename="x.png")], + "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", + ) + + mock_post.assert_not_called() + + +@pytest.mark.asyncio +async def test_pre_call_blocks_file_over_size_limit(): + """A recognized document over Model Armor's 4 MB limit cannot be scanned, so it is blocked.""" + from litellm.proxy.guardrails.guardrail_hooks.model_armor.file_scanning import ( + MODEL_ARMOR_MAX_FILE_SIZE_BYTES, + ) + + guardrail = _make_guardrail() + oversize_b64 = base64.b64encode(b"x" * (MODEL_ARMOR_MAX_FILE_SIZE_BYTES + 1)).decode("utf-8") + request_data = { + "model": "gpt-4", + "messages": [_file_message(oversize_b64)], + "metadata": {"guardrails": ["model-armor-test"]}, + } + + with patch.object( + guardrail.async_handler, + "post", + AsyncMock(return_value=_armor_response(blocked=False)), + ) as mock_post: + with pytest.raises(HTTPException) as exc_info: + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=MagicMock(spec=DualCache), + data=request_data, + call_type="completion", + ) + + assert exc_info.value.status_code == 400 + assert "scan limit" in str(exc_info.value.detail) + # The oversized document is never forwarded to the Model Armor API + mock_post.assert_not_called() + + +@pytest.mark.asyncio +async def test_pre_call_oversize_file_skipped_when_fail_open(): + """With fail_on_error=False the operator opts into fail-open, so an oversized file proceeds.""" + from litellm.proxy.guardrails.guardrail_hooks.model_armor.file_scanning import ( + MODEL_ARMOR_MAX_FILE_SIZE_BYTES, + ) + + guardrail = _make_guardrail(fail_on_error=False) + oversize_b64 = base64.b64encode(b"x" * (MODEL_ARMOR_MAX_FILE_SIZE_BYTES + 1)).decode("utf-8") + request_data = { + "model": "gpt-4", + "messages": [_file_message(oversize_b64)], + "metadata": {"guardrails": ["model-armor-test"]}, + } + + with patch.object( + guardrail.async_handler, + "post", + AsyncMock(return_value=_armor_response(blocked=False)), + ) as mock_post: + result = await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=MagicMock(spec=DualCache), + data=request_data, + call_type="completion", + ) + + assert result == request_data + mock_post.assert_not_called() + + +def _armor_sdp_deidentify_response(): + """A response that only trips the SDP deidentify (PII masking) filter.""" + mock_response = AsyncMock() + mock_response.status_code = 200 + mock_response.json = AsyncMock( + return_value={ + "sanitizationResult": { + "filterMatchState": "MATCH_FOUND", + "filterResults": { + "sdp": { + "sdpFilterResult": { + "deidentifyResult": { + "matchState": "MATCH_FOUND", + "data": {"text": "[REDACTED]"}, + } + } + } + }, + } + } + ) + return mock_response + + +@pytest.mark.asyncio +async def test_pre_call_blocks_pii_document_even_when_masking_enabled(): + """A PII document must block, not pass, even when mask_request_content=True. + + Documents have no masking fallback (Model Armor returns findings, not a sanitized + file), so a deidentify-only match has to block. Without this the original bytes + would reach the provider with PII intact. + """ + guardrail = _make_guardrail(mask_request_content=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(return_value=_armor_sdp_deidentify_response()), + ): + with pytest.raises(HTTPException) as exc_info: + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=MagicMock(spec=DualCache), + data=request_data, + call_type="completion", + ) + + assert exc_info.value.status_code == 400 + assert request_data["metadata"]["_model_armor_status"] == "blocked" + + +@pytest.mark.asyncio +async def test_pre_call_scans_raw_base64_file_without_data_uri(): + """A `type: file` with raw base64 (no data: URI) resolves its MIME from the filename.""" + guardrail = _make_guardrail() + pdf_b64 = base64.b64encode(PDF_BYTES).decode("utf-8") + request_data = { + "model": "gpt-4", + "messages": [ + { + "role": "user", + "content": [ + { + "type": "file", + "file": {"file_data": pdf_b64, "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", + ) + + byte_items = _byte_items_sent(mock_post) + assert len(byte_items) == 1 + assert byte_items[0]["byteDataType"] == "PDF" + assert base64.b64decode(byte_items[0]["byteData"]) == PDF_BYTES + + +@pytest.mark.asyncio +async def test_first_of_multiple_attachments_blocks(): + """Scanning stops and blocks at the first flagged attachment.""" + guardrail = _make_guardrail() + pdf_b64 = base64.b64encode(PDF_BYTES).decode("utf-8") + file_block = { + "type": "file", + "file": {"file_data": f"data:application/pdf;base64,{pdf_b64}"}, + } + request_data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": [file_block, file_block]}], + "metadata": {"guardrails": ["model-armor-test"]}, + } + + with patch.object( + guardrail.async_handler, + "post", + AsyncMock(return_value=_armor_response(blocked=True)), + ) as mock_post: + with pytest.raises(HTTPException): + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=MagicMock(spec=DualCache), + data=request_data, + call_type="completion", + ) + + assert mock_post.call_count == 1 + + +@pytest.mark.asyncio +async def test_file_scan_fail_on_error_false_proceeds(): + """When the Model Armor call errors and fail_on_error=False, the request proceeds.""" + guardrail = _make_guardrail(fail_on_error=False) + 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("Connection error")), + ): + result = await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=MagicMock(spec=DualCache), + data=request_data, + call_type="completion", + ) + + assert result == request_data + + +SUPPORTED_MIME_TYPE_MATRIX = [ + ("application/pdf", "PDF"), + ("application/msword", "WORD_DOCUMENT"), + ( + "application/vnd.openxmlformats-officedocument.wordprocessingml.document", + "WORD_DOCUMENT", + ), + ("application/vnd.ms-excel", "EXCEL_DOCUMENT"), + ( + "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", + "EXCEL_DOCUMENT", + ), + ("application/vnd.ms-powerpoint", "POWERPOINT_DOCUMENT"), + ( + "application/vnd.openxmlformats-officedocument.presentationml.presentation", + "POWERPOINT_DOCUMENT", + ), + ("text/csv", "CSV"), + ("text/plain", "TXT"), +] + + +@pytest.mark.parametrize("mime,expected_byte_data_type", SUPPORTED_MIME_TYPE_MATRIX) +@pytest.mark.asyncio +async def test_pre_call_submits_correct_byte_data_type_for_every_supported_mime(mime, expected_byte_data_type): + """Every supported MIME type maps to the right Model Armor byteDataType and is submitted.""" + guardrail = _make_guardrail() + payload = b"file content for %s" % mime.encode() + payload_b64 = base64.b64encode(payload).decode("utf-8") + request_data = { + "model": "gpt-4", + "messages": [_file_message(payload_b64, mime=mime, filename="attachment")], + "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", + ) + + byte_items = _byte_items_sent(mock_post) + assert len(byte_items) == 1 + assert byte_items[0]["byteDataType"] == expected_byte_data_type + assert base64.b64decode(byte_items[0]["byteData"]) == payload + + +@pytest.mark.asyncio +async def test_pre_call_resolves_mime_from_filename_when_data_uri_is_generic(): + """A data URI with a generic MIME still scans when the filename identifies a document. + + Regression for the case where attachments were skipped because only the data URI + header MIME was consulted, ignoring file.format and the filename. + """ + guardrail = _make_guardrail() + pdf_b64 = base64.b64encode(PDF_BYTES).decode("utf-8") + request_data = { + "model": "gpt-4", + "messages": [ + { + "role": "user", + "content": [ + { + "type": "file", + "file": { + "file_data": f"data:application/octet-stream;base64,{pdf_b64}", + "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", + ) + + byte_items = _byte_items_sent(mock_post) + assert len(byte_items) == 1 + assert byte_items[0]["byteDataType"] == "PDF" + assert base64.b64decode(byte_items[0]["byteData"]) == PDF_BYTES + + +@pytest.mark.asyncio +async def test_pre_call_normalizes_mime_with_charset_suffix(): + """A MIME with a charset parameter (text/plain; charset=utf-8) still maps to TXT.""" + guardrail = _make_guardrail() + payload = b"plain text body" + payload_b64 = base64.b64encode(payload).decode("utf-8") + request_data = { + "model": "gpt-4", + "messages": [ + { + "role": "user", + "content": [ + { + "type": "file", + "file": { + "file_data": f"data:text/plain;charset=utf-8;base64,{payload_b64}", + "filename": "notes.txt", + }, + } + ], + } + ], + "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", + ) + + byte_items = _byte_items_sent(mock_post) + assert len(byte_items) == 1 + assert byte_items[0]["byteDataType"] == "TXT" + + +@pytest.mark.asyncio +async def test_pre_call_scans_macro_enabled_office_document(): + """Macro-enabled and template Office MIME types map to their document family, not skipped.""" + guardrail = _make_guardrail() + payload = b"macro enabled word document bytes" + payload_b64 = base64.b64encode(payload).decode("utf-8") + request_data = { + "model": "gpt-4", + "messages": [ + { + "role": "user", + "content": [ + { + "type": "file", + "file": { + "file_data": f"data:application/vnd.ms-word.document.macroEnabled.12;base64,{payload_b64}", + "filename": "report.docm", + }, + } + ], + } + ], + "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", + ) + + byte_items = _byte_items_sent(mock_post) + assert len(byte_items) == 1 + assert byte_items[0]["byteDataType"] == "WORD_DOCUMENT" + + +@pytest.mark.asyncio +async def test_file_scan_does_not_log_document_bytes(): + """Debug logging must never emit the scanned document's base64 bytes.""" + guardrail = _make_guardrail() + pdf_b64 = base64.b64encode(PDF_BYTES).decode("utf-8") + request_data = { + "model": "gpt-4", + "messages": [_file_message(pdf_b64)], + "metadata": {"guardrails": ["model-armor-test"]}, + } + + logged_args = [] + + def _capture(*args, **kwargs): + logged_args.append(args) + + with patch.object(verbose_proxy_logger, "debug", side_effect=_capture): + with patch.object( + guardrail.async_handler, + "post", + AsyncMock(return_value=_armor_response(blocked=False)), + ): + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=MagicMock(spec=DualCache), + data=request_data, + call_type="completion", + ) + + flattened = " ".join(str(arg) for call in logged_args for arg in call) + assert pdf_b64 not in flattened + # the file request is still logged, just with type and size instead of the bytes + assert "byteDataType" in flattened + + +@pytest.mark.asyncio +async def test_pre_call_prefers_filename_over_conflicting_data_uri_mime(): + """A data URI mislabeled text/plain must not downgrade a .pdf attachment to TXT scanning.""" + guardrail = _make_guardrail() + pdf_b64 = base64.b64encode(PDF_BYTES).decode("utf-8") + request_data = { + "model": "gpt-4", + "messages": [ + { + "role": "user", + "content": [ + { + "type": "file", + "file": { + "file_data": f"data:text/plain;base64,{pdf_b64}", + "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", + ) + + byte_items = _byte_items_sent(mock_post) + assert len(byte_items) == 1 + assert byte_items[0]["byteDataType"] == "PDF" + + +@pytest.mark.asyncio +async def test_file_and_text_responses_are_both_recorded(): + """A request with both a file and text records both Model Armor responses, not just the last.""" + guardrail = _make_guardrail() + pdf_b64 = base64.b64encode(PDF_BYTES).decode("utf-8") + request_data = { + "model": "gpt-4", + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "summarize this"}, + { + "type": "file", + "file": {"file_data": f"data:application/pdf;base64,{pdf_b64}"}, + }, + ], + } + ], + "metadata": {"guardrails": ["model-armor-test"]}, + } + + with patch.object( + guardrail.async_handler, + "post", + AsyncMock(return_value=_armor_response(blocked=False)), + ): + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=MagicMock(spec=DualCache), + data=request_data, + call_type="completion", + ) + + recorded = request_data["metadata"]["_model_armor_response"] + # A list (not a tuple) so the guardrail logging redaction/serialization can walk it + assert isinstance(recorded, list) + assert len(recorded) == 2 + + +@pytest.mark.asyncio +async def test_single_scan_response_stays_a_dict(): + """A single scan keeps the backward-compatible single-dict response shape.""" + guardrail = _make_guardrail() + 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(return_value=_armor_response(blocked=False)), + ): + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=MagicMock(spec=DualCache), + data=request_data, + call_type="completion", + ) + + assert isinstance(request_data["metadata"]["_model_armor_response"], dict) + + +@pytest.mark.asyncio +async def test_pre_call_blocks_supported_document_with_undecodable_base64(): + """A supported document whose inline base64 will not decode cannot be scanned, so it fails closed.""" + guardrail = _make_guardrail() + request_data = { + "model": "gpt-4", + "messages": [ + { + "role": "user", + "content": [ + { + "type": "file", + "file": { + "file_data": "data:application/pdf;base64,@@@not-valid-base64@@@", + "filename": "broken.pdf", + }, + } + ], + } + ], + "metadata": {"guardrails": ["model-armor-test"]}, + } + + with patch.object( + guardrail.async_handler, + "post", + AsyncMock(return_value=_armor_response(blocked=False)), + ) as mock_post: + with pytest.raises(HTTPException) as exc_info: + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=MagicMock(spec=DualCache), + data=request_data, + call_type="completion", + ) + + assert exc_info.value.status_code == 400 + assert "could not scan" in str(exc_info.value.detail) + # The malformed document is never submitted to Model Armor + mock_post.assert_not_called() + + +@pytest.mark.asyncio +async def test_pre_call_undecodable_document_skipped_when_fail_open(): + """With fail_on_error=False a malformed supported document is skipped rather than blocking.""" + guardrail = _make_guardrail(fail_on_error=False) + request_data = { + "model": "gpt-4", + "messages": [ + { + "role": "user", + "content": [ + { + "type": "file", + "file": { + "file_data": "data:application/pdf;base64,@@@not-valid-base64@@@", + "filename": "broken.pdf", + }, + } + ], + } + ], + "metadata": {"guardrails": ["model-armor-test"]}, + } + + with patch.object( + guardrail.async_handler, + "post", + AsyncMock(return_value=_armor_response(blocked=False)), + ) as mock_post: + result = await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=MagicMock(spec=DualCache), + data=request_data, + call_type="completion", + ) + + assert result == request_data + mock_post.assert_not_called() + + +def test_accumulated_responses_are_redactable_as_a_list(): + """Accumulated file+text responses must be a list so guardrail logging can redact nested keys. + + Regression: a tuple is skipped by redact_nested_match_and_regex_keys (it only recurses into + dicts and lists), which would leave sensitive match/regex findings un-redacted in logs. + """ + from litellm.litellm_core_utils.core_helpers import ( + redact_nested_match_and_regex_keys, + ) + + first = {"sanitizationResult": {"filterResults": {"f": {"match": "secret-one"}}}} + second = {"sanitizationResult": {"filterResults": {"f": {"match": "secret-two"}}}} + + accumulated = ModelArmorGuardrail._append_armor_response(first, second) + assert isinstance(accumulated, list) + + redacted = redact_nested_match_and_regex_keys(accumulated) + blob = json.dumps(redacted) + assert "secret-one" not in blob + assert "secret-two" not in blob + assert blob.count("[REDACTED]") == 2