diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index ae36484b294..926441395c5 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -77,7 +77,7 @@ from .common_utils import ( parse_tool_call_arguments, salvage_concatenated_tool_arguments, ) -from .image_handling import convert_url_to_base64, is_pdf_media_type +from .image_handling import convert_url_to_base64, is_pdf_media_type, parse_data_uri def default_pt(messages): @@ -890,18 +890,14 @@ def convert_to_anthropic_image_obj(openai_image_url: str, format: str | None) -> try: if openai_image_url.startswith("http"): openai_image_url = convert_url_to_base64(url=openai_image_url) - # Extract the media type and base64 data - media_type, base64_data = openai_image_url.split("data:")[1].split(";base64,") - - if format: - media_type = format - else: - media_type = media_type.replace("\\/", "/") + parsed: Final = parse_data_uri(openai_image_url) + if parsed is None or not parsed.is_base64: + raise ValueError("expected a base64 data URI") return GenericImageParsingChunk( type="base64", - media_type=media_type, - data=base64_data, + media_type=format or parsed.media_type, + data=parsed.data, ) except litellm.ImageFetchError: raise @@ -1683,16 +1679,10 @@ def convert_to_gemini_tool_call_result( # Detect data-URL images (e.g. from Anthropic tool_result with a single image block # that was serialised as a plain string by translate_anthropic_messages_to_openai) # and promote them to inline_data so Gemini receives actual image bytes. - if content_str[:5].lower() == "data:" and ";base64," in content_str: - try: - mime_rest: Final = content_str[5:].split(";base64,", 1) - if len(mime_rest) == 2 and mime_rest[0].startswith("image/"): - # Strip any extra parameters (e.g. ";charset=UTF-8") from the MIME segment - clean_mime: Final = mime_rest[0].split(";")[0].strip() - media_parts.append({"inline_data": BlobType(data=mime_rest[1], mime_type=clean_mime)}) - content_str = "" - except Exception as e: - verbose_logger.warning("Failed to parse data URL in tool response: %s", e) + data_uri: Final = parse_data_uri(content_str) + if data_uri is not None and data_uri.is_base64 and data_uri.media_type.startswith("image/"): + media_parts.append({"inline_data": BlobType(data=data_uri.data, mime_type=data_uri.media_type)}) + content_str = "" elif isinstance(message["content"], list): content_list: Final = message["content"] for content in content_list: @@ -1813,10 +1803,8 @@ def _is_anthropic_document_data_uri(url: str) -> bool: # text/plain (see select_anthropic_content_block_type_for_file). Routing # other mimes here would produce a document block the API rejects, so we # leave them on the image code path. - match: Final = re.match(r"data:([^;,]+)", url) - if not match: - return False - return match.group(1) in _ANTHROPIC_DOCUMENT_BASE64_MEDIA_TYPES + parsed: Final = parse_data_uri(url) + return parsed is not None and parsed.media_type in _ANTHROPIC_DOCUMENT_BASE64_MEDIA_TYPES def convert_to_anthropic_tool_result( @@ -4104,20 +4092,10 @@ class BedrockImageProcessor: @staticmethod def _parse_base64_image(image_url: str) -> tuple[str, str, str]: """Parse base64 encoded image data.""" - image_metadata, img_without_base_64 = image_url.split(",") - - # 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] - image_format = mime_type.split("/")[1] - else: - mime_type = "image/jpeg" - image_format = "jpeg" - - return img_without_base_64, mime_type, image_format + parsed: Final = parse_data_uri(image_url) + if parsed is None or not parsed.media_type: + return image_url.rpartition(",")[2], "image/jpeg", "jpeg" + return parsed.data, parsed.media_type, parsed.media_type.rpartition("/")[2] @staticmethod def validate_format(mime_type: str, image_format: str) -> str: diff --git a/litellm/litellm_core_utils/prompt_templates/image_handling.py b/litellm/litellm_core_utils/prompt_templates/image_handling.py index fa6f1c1ed0a..dd3958fa764 100644 --- a/litellm/litellm_core_utils/prompt_templates/image_handling.py +++ b/litellm/litellm_core_utils/prompt_templates/image_handling.py @@ -233,6 +233,28 @@ def _bare_media_type(media_type: str) -> str: return media_type.split(";", 1)[0].strip().lower() +@dataclass(frozen=True, slots=True) +class DataUri: + media_type: str + data: str + is_base64: bool + + +def parse_data_uri(url: str) -> DataUri | None: + """Split a `data:` URI into its bare media type (parameters such as `;charset=` dropped) and payload.""" + if url[:5].lower() != "data:": + return None + header, comma, data = url[5:].partition(",") + if not comma: + return None + unescaped_header: Final = header.replace("\\/", "/") + return DataUri( + media_type=_bare_media_type(unescaped_header), + data=data, + is_base64=unescaped_header.rpartition(";")[2].strip().lower() == "base64", + ) + + def is_pdf_media_type(media_type: str) -> bool: return _bare_media_type(media_type) in _PDF_MEDIA_TYPES @@ -438,9 +460,11 @@ def _inlined_file(file: Mapping[str, object], url: str, data_url: str) -> Mappin def _base64_source(url: str, data_url: str) -> Mapping[str, str]: - fetched_media_type, data = data_url.removeprefix("data:").split(";base64,", 1) - media_type: Final = "application/pdf" if url.lower().endswith(".pdf") else fetched_media_type - return {"type": "base64", "media_type": media_type, "data": data} + fetched: Final = parse_data_uri(data_url) + if fetched is None: + raise litellm.ImageFetchError(f"Error: Fetched media is not a base64 data URI. url={url}") + media_type: Final = "application/pdf" if url.lower().endswith(".pdf") else fetched.media_type + return {"type": "base64", "media_type": media_type, "data": fetched.data} def _inline(remote: _RemoteImage | _RemoteFile | _RemoteSource, data_url: str) -> Mapping[str, object]: diff --git a/litellm/llms/bedrock/embed/amazon_nova_transformation.py b/litellm/llms/bedrock/embed/amazon_nova_transformation.py index cc8c8d37c42..59b7fdfb917 100644 --- a/litellm/llms/bedrock/embed/amazon_nova_transformation.py +++ b/litellm/llms/bedrock/embed/amazon_nova_transformation.py @@ -14,6 +14,7 @@ Docs - https://docs.aws.amazon.com/bedrock/latest/userguide/nova-embed.html from typing import Final +from litellm.litellm_core_utils.prompt_templates.image_handling import parse_data_uri from litellm.types.utils import ( Embedding, EmbeddingResponse, @@ -62,26 +63,10 @@ class AmazonNovaEmbeddingConfig: media_type: e.g., "image/jpeg", "video/mp4", "audio/mpeg" base64_data: The base64-encoded data without the prefix """ - if not data_url.startswith("data:"): + parsed: Final = parse_data_uri(data_url) + if parsed is None: raise ValueError(f"Invalid data URL format: {data_url[:50]}...") - - # Split by comma to separate metadata from data - # Format: data:image/jpeg;base64, - if "," not in data_url: - raise ValueError(f"Invalid data URL format (missing comma): {data_url[:50]}...") - - metadata, base64_data = data_url.split(",", 1) - - # Extract media type from metadata - # Remove 'data:' prefix and ';base64' suffix - metadata = metadata[5:] # Remove 'data:' - - if ";" in metadata: - media_type = metadata.split(";")[0] - else: - media_type = metadata - - return media_type, base64_data + return parsed.media_type, parsed.data def _transform_request( self, diff --git a/litellm/llms/gigachat/file_handler.py b/litellm/llms/gigachat/file_handler.py index d53ae3af334..1b16b508613 100644 --- a/litellm/llms/gigachat/file_handler.py +++ b/litellm/llms/gigachat/file_handler.py @@ -7,7 +7,6 @@ GigaChat requires files to be uploaded first, then referenced by file_id. import base64 import hashlib -import re import uuid from collections.abc import Mapping from typing import Final @@ -15,6 +14,7 @@ from typing import Final from pydantic import TypeAdapter from litellm._logging import verbose_logger +from litellm.litellm_core_utils.prompt_templates.image_handling import parse_data_uri from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, @@ -43,16 +43,14 @@ def _parse_data_url(data_url: str) -> tuple[bytes, str, str] | None: Returns: Tuple of (content_bytes, content_type, extension) or None """ - match: Final = re.match(r"data:([^;]+);base64,(.+)", data_url) - if not match: + parsed: Final = parse_data_uri(data_url) + if parsed is None or not parsed.is_base64 or not parsed.media_type or not parsed.data: return None - content_type: Final = match.group(1) - base64_data: Final = match.group(2) - content_bytes: Final = base64.b64decode(base64_data) - ext: Final = content_type.split("/")[-1].split(";")[0] or "jpg" + content_bytes: Final = base64.b64decode(parsed.data) + ext: Final = parsed.media_type.rpartition("/")[2] or "jpg" - return content_bytes, content_type, ext + return content_bytes, parsed.media_type, ext def _content_type_or_default(headers: Mapping[str, str]) -> str: diff --git a/litellm/llms/snowflake/chat/transformation.py b/litellm/llms/snowflake/chat/transformation.py index 928dcc3bc52..a7e071107f1 100644 --- a/litellm/llms/snowflake/chat/transformation.py +++ b/litellm/llms/snowflake/chat/transformation.py @@ -10,7 +10,6 @@ Ref: https://docs.snowflake.com/en/user-guide/snowflake-cortex/cortex-rest-api import copy import json -import re from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, Protocol, TypedDict @@ -24,7 +23,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( create_anthropic_image_param, select_anthropic_content_block_type_for_file, ) -from litellm.litellm_core_utils.prompt_templates.image_handling import async_inline_remote_media +from litellm.litellm_core_utils.prompt_templates.image_handling import async_inline_remote_media, parse_data_uri from litellm.litellm_core_utils.url_utils import join_api_base from litellm.llms.anthropic.chat.handler import ModelResponseIterator as AnthropicStreamParser from litellm.llms.anthropic.chat.transformation import AnthropicConfig @@ -150,8 +149,8 @@ def _str_field(container: object, key: str) -> str | None: def _data_uri_media_type(url: str) -> str: - match: Final = re.match(r"data:([^;,]+)", url) - return match.group(1) if match else "" + parsed: Final = parse_data_uri(url) + return parsed.media_type if parsed else "" def _convert_media_block_to_anthropic(block: object) -> object: diff --git a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py index 96a785c06c9..2b203c48f6b 100644 --- a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py +++ b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py @@ -10,6 +10,7 @@ from typing import Annotated, Final, Literal, cast from pydantic import ConfigDict, Field, TypeAdapter, ValidationError from litellm.exceptions import BadRequestError +from litellm.litellm_core_utils.prompt_templates.image_handling import parse_data_uri from litellm.llms.vertex_ai.common_utils import GEMINI_FILES_API_URI_PREFIX, gemini_video_metadata_from_openai from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.vertex_ai import ( @@ -116,25 +117,20 @@ def _parse_data_url(data_url: str, mime_type_override: str | None = None) -> tup Raises: ValueError: If data URL format is invalid or MIME type is unsupported """ - if not data_url.startswith("data:"): + parsed: Final = parse_data_uri(data_url) + if parsed is None: raise ValueError(f"Invalid data URL format: {data_url[:50]}...") - if "," not in data_url: - raise ValueError(f"Invalid data URL format (missing comma): {data_url[:50]}...") - - metadata, base64_data = data_url.split(",", 1) - declared_media_type: Final = metadata[5:].split(";")[0] - if mime_type_override is not None: - return mime_type_override, base64_data + return mime_type_override, parsed.data - if declared_media_type not in SUPPORTED_EMBEDDING_MIME_TYPES: + if parsed.media_type not in SUPPORTED_EMBEDDING_MIME_TYPES: raise ValueError( - f"Unsupported MIME type for embedding: {declared_media_type}. " + f"Unsupported MIME type for embedding: {parsed.media_type}. " f"Supported types: {', '.join(sorted(SUPPORTED_EMBEDDING_MIME_TYPES))}" ) - return declared_media_type, base64_data + return parsed.media_type, parsed.data def _is_data_url(s: str) -> bool: diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index fafc098ed5d..6c155ad2c76 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -6604,6 +6604,8 @@ def test_bedrock_tool_calling_pt(): ("data:image/jpeg;base64,1234", "image/jpeg"), ("data:application/pdf;base64,1234", "application/pdf"), (r"data:image\/jpeg;base64,1234", "image/jpeg"), + ("data:image/png;charset=utf-8;base64,1234", "image/png"), + ("data:image/PNG;base64,1234", "image/png"), ], ) def test_base64_image_input(url, expected_media_type): @@ -6612,6 +6614,21 @@ def test_base64_image_input(url, expected_media_type): assert response["media_type"] == expected_media_type +@pytest.mark.parametrize( + "url", + [ + "data:image/png;base64,iVBORw0KGgo=", + "data:image/png;charset=utf-8;base64,iVBORw0KGgo=", + "data:Image/WEBP;base64,UklGRg==", + ], +) +def test_anthropic_and_bedrock_converse_agree_on_data_uri_media_type(url): + anthropic_media_type = convert_to_anthropic_image_obj(openai_image_url=url, format=None)["media_type"] + bedrock_format = BedrockImageProcessor.process_image_sync(url)["image"]["format"] + + assert anthropic_media_type == f"image/{bedrock_format}" + + def test_create_anthropic_image_param_with_http_url(): """Test that HTTP/HTTPS URLs are passed as URL references, not base64.""" image_param = create_anthropic_image_param("https://example.com/image.jpg", format=None) diff --git a/tests/unit/litellm_core_utils/test_image_handling.py b/tests/unit/litellm_core_utils/test_image_handling.py index 2d4cd4b9a6d..03517f6f546 100644 --- a/tests/unit/litellm_core_utils/test_image_handling.py +++ b/tests/unit/litellm_core_utils/test_image_handling.py @@ -18,6 +18,7 @@ from litellm.litellm_core_utils import url_utils from litellm.litellm_core_utils.prompt_templates import image_handling from litellm.litellm_core_utils.prompt_templates.image_handling import ( MAX_CONCURRENT_REMOTE_MEDIA_FETCHES, + DataUri, RemoteMedia, async_convert_url_to_base64, async_format_pdf_links, @@ -27,6 +28,7 @@ from litellm.litellm_core_utils.prompt_templates.image_handling import ( format_pdf_links, inline_remote_media, is_unformatted_pdf_link_image, + parse_data_uri, serves_pdf, ) from litellm.litellm_core_utils.url_utils import SSRFError @@ -236,6 +238,23 @@ def test_streaming_download_handles_petabyte_file(monkeypatch): assert "exceeds maximum allowed size" in str(excinfo.value) +@pytest.mark.parametrize( + "url, expected", + [ + ("data:image/png;base64,iVBORw0KGgo=", DataUri("image/png", "iVBORw0KGgo=", True)), + ("data:image/png;charset=utf-8;base64,iVBORw0KGgo=", DataUri("image/png", "iVBORw0KGgo=", True)), + ("DATA:Image/PNG;BASE64,iVBORw0KGgo=", DataUri("image/png", "iVBORw0KGgo=", True)), + (r"data:image\/jpeg;base64,/9j/", DataUri("image/jpeg", "/9j/", True)), + ("data:text/plain,a,b", DataUri("text/plain", "a,b", False)), + ("data:;base64,AAAA", DataUri("", "AAAA", True)), + ("data:image/png;base64", None), + ("https://example.com/a.png", None), + ], +) +def test_parse_data_uri(url, expected): + assert parse_data_uri(url) == expected + + def test_data_url_is_returned_unchanged_without_fetch(monkeypatch): """ A data URL is already inline base64 image data, so convert_url_to_base64 diff --git a/tests/unit/llms/gigachat/test_file_handler.py b/tests/unit/llms/gigachat/test_file_handler.py index 0739842d7bf..e1c9b437a5c 100644 --- a/tests/unit/llms/gigachat/test_file_handler.py +++ b/tests/unit/llms/gigachat/test_file_handler.py @@ -90,16 +90,9 @@ class TestParseDataUrl: assert content_type == "image/jpeg" assert ext == "jpeg" - def test_valid_base64_with_semicolon_in_type(self): - """Data URLs with charset before base64 segment do not match the regex.""" - # The regex `data:([^;]+);base64,(.+)` requires the pattern to be - # `data:;base64,`. If `;charset=utf-8` appears before - # `;base64,`, the regex sees `data:image/png` as group 1 but then - # looks for `;base64,` immediately after — which isn't there because - # `;charset=utf-8;base64,` has extra text before `;base64,` + def test_media_type_parameters_are_dropped(self): data_url = "data:image/png;charset=utf-8;base64," + _RED_PNG_B64 - result = _parse_data_url(data_url) - assert result is None + assert _parse_data_url(data_url) == _parse_data_url(_RED_PNG_DATA_URL) def test_invalid_data_url_returns_none(self): assert _parse_data_url("not-a-data-url") is None diff --git a/tests/unit/llms/vertex_ai/test_gemini_batch_embeddings.py b/tests/unit/llms/vertex_ai/test_gemini_batch_embeddings.py index 7271fb48861..cdd234ec0de 100644 --- a/tests/unit/llms/vertex_ai/test_gemini_batch_embeddings.py +++ b/tests/unit/llms/vertex_ai/test_gemini_batch_embeddings.py @@ -194,7 +194,7 @@ def test_parse_data_url_invalid_format(): with pytest.raises(ValueError, match="Invalid data URL format"): _parse_data_url("not-a-data-url") - with pytest.raises(ValueError, match="missing comma"): + with pytest.raises(ValueError, match="Invalid data URL format"): _parse_data_url("data:image/png;base64")