From 0ccc493455463fa3994b33250f1b5bfd0a9bedb2 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 5 Aug 2025 17:07:10 -0700 Subject: [PATCH] [Bug]: Fix Mimetype Resolution Error in Bedrock Document Understanding (#13309) * fix _validate_format for BedrockImageProcessor * add test * fix _validate_format for bedrock * _get_document_format * test_bedrock_get_document_format_fallback_mimes * fix: add fallback method for mime type detection --- .../prompt_templates/factory.py | 79 +++++++++++++++---- ...llm_core_utils_prompt_templates_factory.py | 66 ++++++++++++++++ 2 files changed, 130 insertions(+), 15 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index b4ace1545d2..26388dc2362 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -1,5 +1,6 @@ import copy import json +import mimetypes import re import uuid import xml.etree.ElementTree as ET @@ -13,6 +14,7 @@ import litellm.types import litellm.types.llms from litellm import verbose_logger from litellm.llms.custom_httpx.http_handler import HTTPHandler, get_async_httpx_client +from litellm.types.files import get_file_extension_from_mime_type from litellm.types.llms.anthropic import * from litellm.types.llms.bedrock import MessageBlock as BedrockMessageBlock from litellm.types.llms.custom_http import httpxSpecialProvider @@ -2351,7 +2353,6 @@ def stringify_json_tool_call_content(messages: List) -> List: ###### AMAZON BEDROCK ####### import base64 -import mimetypes from email.message import Message import httpx @@ -2479,20 +2480,11 @@ class BedrockImageProcessor: ) if is_document: - potential_extensions = mimetypes.guess_all_extensions(mime_type) - valid_extensions = [ - ext[1:] - for ext in potential_extensions - if ext[1:] in supported_doc_formats - ] + return BedrockImageProcessor._get_document_format( + mime_type=mime_type, + supported_doc_formats=supported_doc_formats + ) - if not valid_extensions: - raise ValueError( - f"No supported extensions for MIME type: {mime_type}. Supported formats: {supported_doc_formats}" - ) - - # Use first valid extension instead of provided image_format - return valid_extensions[0] else: ######################################################### # Check if image_format is an image or video @@ -2502,6 +2494,60 @@ class BedrockImageProcessor: f"Unsupported image format: {image_format}. Supported formats: {supported_image_and_video_formats}" ) return image_format + + @staticmethod + def _get_document_format( + mime_type: str, + supported_doc_formats: List[str] + ) -> str: + """ + Get the document format from the mime type + + - Primary method - uses `mimetypes.guess_all_extensions` + - Fallback method - uses `get_file_extension_from_mime_type` + + Relevant Issue: https://github.com/BerriAI/litellm/issues/12260 + + `mimetypes` is not available in docker containers, so we fallback to `get_file_extension_from_mime_type` + + Args: + mime_type: The mime type of the document + supported_doc_formats: The supported document formats for the current model + + Returns: + The document format + """ + valid_extensions: Optional[List[str]] = None + potential_extensions = mimetypes.guess_all_extensions( + mime_type, strict=False + ) + valid_extensions = [ + ext[1:] + for ext in potential_extensions + if ext[1:] in supported_doc_formats + ] + + # Fallback to types/files.py if mimetypes doesn't return valid extensions + ################# + # litellm runs on docker containers and `mimetypes` depends on the installed mimetypes of the OS + # we fallback to well known mime types in types/files.py if mimetypes doesn't return valid extensions + if not valid_extensions: + try: + fallback_extension = get_file_extension_from_mime_type(mime_type) + if fallback_extension in supported_doc_formats: + valid_extensions = [fallback_extension] + except ValueError: + # Neither mimetypes nor files.py could handle this MIME type + # get_file_extension_from_mime_type raises ValueError if the mime type is not supported + pass + + if not valid_extensions: + raise ValueError( + f"No supported extensions for MIME type: {mime_type}. Supported formats: {supported_doc_formats}" + ) + + # Use first valid extension instead of provided image_format + return valid_extensions[0] @staticmethod def _create_bedrock_block( @@ -2950,7 +2996,10 @@ def process_empty_text_blocks( ] modified_message = message.copy() - modified_message["content"] = modified_content_block + modified_message["content"] = cast( + Union[List[ChatCompletionTextObject], List[ChatCompletionThinkingBlock]], + modified_content_block, + ) return modified_message diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index 0cfb9f87d14..e3ca7101c6c 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -139,6 +139,70 @@ def test_bedrock_validate_format_image_or_video(): result = BedrockImageProcessor._validate_format(f"video/{format}", format) assert result == format, f"Expected {format}, got {result}" + # Test valid document formats + valid_document_formats = { + "application/pdf": "pdf", + "text/csv": "csv", + "application/vnd.openxmlformats-officedocument.wordprocessingml.document": "docx", + "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet": "xlsx", + } + for mime, expected in valid_document_formats.items(): + print("testing mime", mime, "expected", expected) + result = BedrockImageProcessor._validate_format( + mime, mime.split("/")[1] + ) + assert result == expected, f"Expected {expected}, got {result}" + + +def test_bedrock_get_document_format_fallback_mimes(): + """ + Test the _get_document_format method with fallback MIME types for DOCX and XLSX. + + This tests the fallback mechanism when mimetypes.guess_all_extensions returns empty results, + which can happen in Docker containers where mimetypes depends on OS-installed MIME types. + """ + from unittest.mock import patch + + # Test DOCX fallback + docx_mime = "application/vnd.openxmlformats-officedocument.wordprocessingml.document" + supported_formats = ["pdf", "docx", "xlsx", "csv"] + + # Mock mimetypes.guess_all_extensions to return empty list (simulating Docker container scenario) + with patch('mimetypes.guess_all_extensions', return_value=[]): + result = BedrockImageProcessor._get_document_format( + mime_type=docx_mime, + supported_doc_formats=supported_formats + ) + assert result == "docx", f"Expected 'docx', got '{result}'" + + # Test XLSX fallback + xlsx_mime = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet" + + with patch('mimetypes.guess_all_extensions', return_value=[]): + result = BedrockImageProcessor._get_document_format( + mime_type=xlsx_mime, + supported_doc_formats=supported_formats + ) + assert result == "xlsx", f"Expected 'xlsx', got '{result}'" + + +def test_bedrock_get_document_format_mimetypes_success(): + """ + Test the _get_document_format method when mimetypes.guess_all_extensions works normally. + """ + docx_mime = "application/vnd.openxmlformats-officedocument.wordprocessingml.document" + supported_formats = ["pdf", "docx", "xlsx", "csv"] + + # Test normal mimetypes behavior (should not hit fallback) + result = BedrockImageProcessor._get_document_format( + mime_type=docx_mime, + supported_doc_formats=supported_formats + ) + assert result == "docx", f"Expected 'docx', got '{result}'" + + + + # def test_ollama_pt_consecutive_system_messages(): # """Test handling consecutive system messages""" @@ -505,3 +569,5 @@ def test_bedrock_tools_unpack_defs(): ] _bedrock_tools_pt(tools=tools) + +