diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 795fb36961e..b4df3cf2518 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -3412,9 +3412,8 @@ class BedrockImageProcessor: # Extract MIME type using regular expression mime_type_match: Final = re.match(r"data:(.*?);base64", image_metadata) - if mime_type_match: - mime_type = mime_type_match.group(1) - mime_type = mime_type.split(";")[0] + mime_type = mime_type_match.group(1).split(";")[0] if mime_type_match else "" + if "/" in mime_type: image_format = mime_type.split("/")[1] else: mime_type = "image/jpeg" 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 72d26f31c60..1fe864df6b8 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 @@ -327,6 +327,21 @@ def test_bedrock_validate_format_image_or_video(): assert result == expected, f"Expected {expected}, got {result}" +def test_bedrock_parse_base64_image_missing_mime_subtype(): + """A base64 data URI without a mime subtype must fall back to jpeg instead of raising IndexError.""" + + data = base64.b64encode(b"fake").decode() + + for prefix in ["data:", "data:application"]: + _, mime_type, image_format = BedrockImageProcessor._parse_base64_image(f"{prefix};base64,{data}") + assert mime_type == "image/jpeg", f"Expected image/jpeg fallback, got {mime_type}" + assert image_format == "jpeg", f"Expected jpeg fallback, got {image_format}" + + _, mime_type, image_format = BedrockImageProcessor._parse_base64_image(f"data:image/png;base64,{data}") + assert mime_type == "image/png" + assert image_format == "png" + + def test_bedrock_get_document_format_fallback_mimes(): """ Test the _get_document_format method with fallback MIME types for DOCX and XLSX.