From 64cfe75bfd398aa856c117aedfb6b72c70f48d16 Mon Sep 17 00:00:00 2001 From: lif <1835304752@qq.com> Date: Sun, 4 Jan 2026 03:17:38 +0800 Subject: [PATCH] fix: extract pure base64 data from data URLs for Ollama (#18465) Fix Ollama_chatException "illegal base64 data at input byte 4" error when using images with ollama_chat provider. Ollama expects pure base64 data, not the full data URL format (data:image/png;base64,...). Fixes #18338 Signed-off-by: majiayu000 <1835304752@qq.com> --- .../prompt_templates/common_utils.py | 32 +++- .../test_extract_base64_image.py | 156 ++++++++++++++++++ 2 files changed, 185 insertions(+), 3 deletions(-) create mode 100644 tests/test_litellm/litellm_core_utils/test_extract_base64_image.py diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index ca2a092dbc..b100b9b516 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -1087,9 +1087,35 @@ def _parse_content_for_reasoning( return None, message_text +def _extract_base64_data(image_url: str) -> str: + """ + Extract pure base64 data from an image URL. + + If the URL is a data URL (e.g., "data:image/png;base64,iVBOR..."), + extract and return only the base64 data portion. + Otherwise, return the original URL unchanged. + + This is needed for providers like Ollama that expect pure base64 data + rather than full data URLs. + + Args: + image_url: The image URL or data URL to process + + Returns: + The base64 data if it's a data URL, otherwise the original URL + """ + if image_url.startswith("data:") and ";base64," in image_url: + return image_url.split(";base64,", 1)[1] + return image_url + + def extract_images_from_message(message: AllMessageValues) -> List[str]: """ - Extract images from a message + Extract images from a message. + + For data URLs (e.g., "data:image/png;base64,iVBOR..."), only the base64 + data portion is extracted. This is required for providers like Ollama + that expect pure base64 data rather than full data URLs. """ images = [] message_content = message.get("content") @@ -1098,7 +1124,7 @@ def extract_images_from_message(message: AllMessageValues) -> List[str]: image_url = m.get("image_url") if image_url: if isinstance(image_url, str): - images.append(image_url) + images.append(_extract_base64_data(image_url)) elif isinstance(image_url, dict) and "url" in image_url: - images.append(image_url["url"]) + images.append(_extract_base64_data(image_url["url"])) return images diff --git a/tests/test_litellm/litellm_core_utils/test_extract_base64_image.py b/tests/test_litellm/litellm_core_utils/test_extract_base64_image.py new file mode 100644 index 0000000000..b17c02d700 --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_extract_base64_image.py @@ -0,0 +1,156 @@ +""" +Unit tests for _extract_base64_data and extract_images_from_message functions. + +These tests verify that base64 image data is correctly extracted from data URLs, +which fixes the Ollama error "illegal base64 data at input byte 4". + +Related issue: https://github.com/BerriAI/litellm/issues/18338 +""" +import pytest + +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + _extract_base64_data, + extract_images_from_message, +) + + +class TestExtractBase64Data: + """Tests for _extract_base64_data function""" + + def test_extract_base64_from_png_data_url(self): + """Test extracting base64 data from a PNG data URL""" + data_url = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk" + expected = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk" + assert _extract_base64_data(data_url) == expected + + def test_extract_base64_from_jpeg_data_url(self): + """Test extracting base64 data from a JPEG data URL""" + data_url = "data:image/jpeg;base64,/9j/4AAQSkZJRgABAQAAAQABAAD" + expected = "/9j/4AAQSkZJRgABAQAAAQABAAD" + assert _extract_base64_data(data_url) == expected + + def test_extract_base64_from_gif_data_url(self): + """Test extracting base64 data from a GIF data URL""" + data_url = "data:image/gif;base64,R0lGODlhAQABAIAAAAAAAP" + expected = "R0lGODlhAQABAIAAAAAAAP" + assert _extract_base64_data(data_url) == expected + + def test_regular_url_unchanged(self): + """Test that regular HTTP URLs are returned unchanged""" + url = "https://example.com/image.png" + assert _extract_base64_data(url) == url + + def test_file_path_unchanged(self): + """Test that file paths are returned unchanged""" + path = "/path/to/image.png" + assert _extract_base64_data(path) == path + + def test_data_url_without_base64_unchanged(self): + """Test that data URLs without base64 encoding are returned unchanged""" + # This is a data URL with URL encoding, not base64 + url = "data:text/plain,Hello%20World" + assert _extract_base64_data(url) == url + + def test_base64_data_with_special_chars(self): + """Test extracting base64 data that contains valid special characters""" + # Base64 can contain +, /, and = characters + data_url = "data:image/png;base64,abc+def/ghi===" + expected = "abc+def/ghi===" + assert _extract_base64_data(data_url) == expected + + +class TestExtractImagesFromMessage: + """Tests for extract_images_from_message function""" + + def test_extract_from_message_with_data_url_string(self): + """Test extracting images when image_url is a string data URL""" + message = { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": "data:image/png;base64,iVBORw0KGgo", + } + ], + } + result = extract_images_from_message(message) + assert result == ["iVBORw0KGgo"] + + def test_extract_from_message_with_data_url_dict(self): + """Test extracting images when image_url is a dict with url key""" + message = { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": {"url": "data:image/png;base64,iVBORw0KGgo"}, + } + ], + } + result = extract_images_from_message(message) + assert result == ["iVBORw0KGgo"] + + def test_extract_from_message_with_regular_url(self): + """Test that regular URLs are preserved""" + message = { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": {"url": "https://example.com/image.png"}, + } + ], + } + result = extract_images_from_message(message) + assert result == ["https://example.com/image.png"] + + def test_extract_multiple_images(self): + """Test extracting multiple images from a single message""" + message = { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": "data:image/png;base64,image1base64", + }, + { + "type": "image_url", + "image_url": {"url": "data:image/jpeg;base64,image2base64"}, + }, + { + "type": "image_url", + "image_url": "https://example.com/image3.png", + }, + ], + } + result = extract_images_from_message(message) + assert result == [ + "image1base64", + "image2base64", + "https://example.com/image3.png", + ] + + def test_empty_content(self): + """Test message with empty content""" + message = {"role": "user", "content": []} + result = extract_images_from_message(message) + assert result == [] + + def test_no_images_in_content(self): + """Test message with content but no images""" + message = { + "role": "user", + "content": [{"type": "text", "text": "Hello world"}], + } + result = extract_images_from_message(message) + assert result == [] + + def test_string_content(self): + """Test message with string content (no images possible)""" + message = {"role": "user", "content": "Hello world"} + result = extract_images_from_message(message) + assert result == [] + + +if __name__ == "__main__": + pytest.main([__file__, "-v"])