better resolve binary (#19856)

Author: logan-markewich

llama-index-core/llama_index/core/utils.py

@@ -34,6 +34,7 @@
 
 import platformdirs
 import requests
+from urllib.parse import urlparse
 
 if TYPE_CHECKING:
     from nltk.tokenize import PunktSentenceTokenizer
@@ -657,6 +658,37 @@ def resolve_binary(
         return BytesIO(data)
 
     elif url is not None:
+        parsed_url = urlparse(url)
+        if parsed_url.scheme == "data":
+            # Parse data URL: data:[<mediatype>][;base64],<data>
+            # The path contains everything after "data:"
+            data_part = parsed_url.path
+
+            # Split on the first comma to separate metadata from data
+            if "," not in data_part:
+                raise ValueError("Invalid data URL format: missing comma separator")
+
+            metadata, url_data = data_part.split(",", 1)
+            is_base64_encoded = metadata.endswith(";base64")
+
+            if is_base64_encoded:
+                # Data is base64 encoded in the URL
+                decoded_data = base64.b64decode(url_data)
+                if as_base64:
+                    # Return as base64 bytes
+                    return BytesIO(base64.b64encode(decoded_data))
+                else:
+                    # Return decoded binary data
+                    return BytesIO(decoded_data)
+            else:
+                # Data is not base64 encoded in the URL (URL-encoded text)
+                if as_base64:
+                    # Encode the text data as base64
+                    return BytesIO(base64.b64encode(url_data.encode("utf-8")))
+                else:
+                    # Return as text bytes
+                    return BytesIO(url_data.encode("utf-8"))
+
         headers = {
             "User-Agent": "LlamaIndex/0.0 (https://llamaindex.ai; [email protected]) llama-index-core/0.0"
         }

llama-index-core/tests/base/llms/test_types.py

@@ -201,6 +201,46 @@ def test_image_block_resolve_image_url(png_1px_b64: bytes, png_1px: bytes):
         assert img.read() == png_1px_b64
 
 
+def test_image_block_resolve_image_data_url_base64(png_1px_b64: bytes, png_1px: bytes):
+    # Test data URL with base64 encoding
+    data_url = f"data:image/png;base64,{png_1px_b64.decode('utf-8')}"
+    b = ImageBlock(url=AnyUrl(url=data_url))
+
+    img = b.resolve_image()
+    assert isinstance(img, BytesIO)
+    assert img.read() == png_1px
+
+    img = b.resolve_image(as_base64=True)
+    assert isinstance(img, BytesIO)
+    assert img.read() == png_1px_b64
+
+
+def test_image_block_resolve_image_data_url_plain_text():
+    # Test data URL with plain text (no base64)
+    test_text = "Hello, World!"
+    data_url = f"data:text/plain,{test_text}"
+    b = ImageBlock(url=AnyUrl(url=data_url))
+
+    img = b.resolve_image()
+    assert isinstance(img, BytesIO)
+    assert img.read() == test_text.encode("utf-8")
+
+    img = b.resolve_image(as_base64=True)
+    assert isinstance(img, BytesIO)
+    assert img.read() == base64.b64encode(test_text.encode("utf-8"))
+
+
+def test_image_block_resolve_image_data_url_invalid():
+    # Test invalid data URL format (missing comma)
+    invalid_data_url = "data:image/png;base64"
+    b = ImageBlock(url=AnyUrl(url=invalid_data_url))
+
+    with pytest.raises(
+        ValueError, match="Invalid data URL format: missing comma separator"
+    ):
+        b.resolve_image()
+
+
 def test_image_block_resolve_error():
     with pytest.raises(
         ValueError, match="No valid source provided to resolve binary data!"
@@ -383,6 +423,20 @@ def test_video_block_resolve_video_url(mp4_bytes: bytes, mp4_base64: bytes):
         assert vid.read() == mp4_base64
 
 
+def test_video_block_resolve_video_data_url_base64(mp4_bytes: bytes, mp4_base64: bytes):
+    # Test data URL with base64 encoding
+    data_url = f"data:video/mp4;base64,{mp4_base64.decode('utf-8')}"
+    b = VideoBlock(url=AnyUrl(url=data_url))
+
+    vid = b.resolve_video()
+    assert isinstance(vid, BytesIO)
+    assert vid.read() == mp4_bytes
+
+    vid = b.resolve_video(as_base64=True)
+    assert isinstance(vid, BytesIO)
+    assert vid.read() == mp4_base64
+
+
 def test_video_block_resolve_error():
     b = VideoBlock()
     with pytest.raises(ValueError, match="No valid source provided"):