feat: added support for images types (stanfordnlp#7872)

Author: glesperance

dspy/adapters/types/image.py

@@ -4,6 +4,7 @@
 from typing import Any, Dict, List, Union
 from urllib.parse import urlparse
 import re
+import mimetypes
 
 import pydantic
 import requests
@@ -79,23 +80,23 @@ def is_url(string: str) -> bool:
 
 def encode_image(image: Union[str, bytes, "PILImage.Image", dict], download_images: bool = False) -> str:
     """
-    Encode an image to a base64 data URI.
+    Encode an image or file to a base64 data URI.
 
     Args:
-        image: The image to encode. Can be a PIL Image, file path, URL, or data URI.
+        image: The image or file to encode. Can be a PIL Image, file path, URL, or data URI.
         download_images: Whether to download images from URLs.
 
     Returns:
-        str: The data URI of the image or the URL if download_images is False.
+        str: The data URI of the file or the URL if download_images is False.
 
     Raises:
-        ValueError: If the image type is not supported.
+        ValueError: If the file type is not supported.
     """
     if isinstance(image, dict) and "url" in image:
         # NOTE: Not doing other validation for now
         return image["url"]
     elif isinstance(image, str):
-        if image.startswith("data:image/"):
+        if image.startswith("data:"):
             # Already a data URI
             return image
         elif os.path.isfile(image):
@@ -110,8 +111,8 @@ def encode_image(image: Union[str, bytes, "PILImage.Image", dict], download_imag
                 return image
         else:
             # Unsupported string format
-            print(f"Unsupported image string: {image}")
-            raise ValueError(f"Unsupported image string: {image}")
+            print(f"Unsupported file string: {image}")
+            raise ValueError(f"Unsupported file string: {image}")
     elif PIL_AVAILABLE and isinstance(image, PILImage.Image):
         # PIL Image
         return _encode_pil_image(image)
@@ -129,34 +130,52 @@ def encode_image(image: Union[str, bytes, "PILImage.Image", dict], download_imag
 
 
 def _encode_image_from_file(file_path: str) -> str:
-    """Encode an image from a file path to a base64 data URI."""
-    with open(file_path, "rb") as image_file:
-        image_data = image_file.read()
-    file_extension = _get_file_extension(file_path)
-    encoded_image = base64.b64encode(image_data).decode("utf-8")
-    return f"data:image/{file_extension};base64,{encoded_image}"
+    """Encode a file from a file path to a base64 data URI."""
+    with open(file_path, "rb") as file:
+        file_data = file.read()
+
+    # Use mimetypes to guess directly from the file path
+    mime_type, _ = mimetypes.guess_type(file_path)
+    if mime_type is None:
+        raise ValueError(f"Could not determine MIME type for file: {file_path}")
+
+    encoded_data = base64.b64encode(file_data).decode("utf-8")
+    return f"data:{mime_type};base64,{encoded_data}"
 
 
 def _encode_image_from_url(image_url: str) -> str:
-    """Encode an image from a URL to a base64 data URI."""
+    """Encode a file from a URL to a base64 data URI."""
     response = requests.get(image_url)
     response.raise_for_status()
     content_type = response.headers.get("Content-Type", "")
-    if content_type.startswith("image/"):
-        file_extension = content_type.split("/")[-1]
+
+    # Use the content type from the response headers if available
+    if content_type:
+        mime_type = content_type
     else:
-        # Fallback to file extension from URL or default to 'png'
-        file_extension = _get_file_extension(image_url) or "png"
-    encoded_image = base64.b64encode(response.content).decode("utf-8")
-    return f"data:image/{file_extension};base64,{encoded_image}"
+        # Try to guess MIME type from URL
+        mime_type, _ = mimetypes.guess_type(image_url)
+        if mime_type is None:
+            raise ValueError(f"Could not determine MIME type for URL: {image_url}")
+
+    encoded_data = base64.b64encode(response.content).decode("utf-8")
+    return f"data:{mime_type};base64,{encoded_data}"
+
 
 def _encode_pil_image(image: 'PILImage') -> str:
     """Encode a PIL Image object to a base64 data URI."""
     buffered = io.BytesIO()
-    file_extension = (image.format or "PNG").lower()
-    image.save(buffered, format=file_extension)
-    encoded_image = base64.b64encode(buffered.getvalue()).decode("utf-8")
-    return f"data:image/{file_extension};base64,{encoded_image}"
+    file_format = image.format or "PNG"
+    image.save(buffered, format=file_format)
+
+    # Get the correct MIME type using the image format
+    file_extension = file_format.lower()
+    mime_type, _ = mimetypes.guess_type(f"file.{file_extension}")
+    if mime_type is None:
+        raise ValueError(f"Could not determine MIME type for image format: {file_format}")
+
+    encoded_data = base64.b64encode(buffered.getvalue()).decode("utf-8")
+    return f"data:{mime_type};base64,{encoded_data}"
 
 
 def _get_file_extension(path_or_url: str) -> str:
@@ -166,11 +185,11 @@ def _get_file_extension(path_or_url: str) -> str:
 
 
 def is_image(obj) -> bool:
-    """Check if the object is an image or a valid image reference."""
+    """Check if the object is an image or a valid media file reference."""
     if PIL_AVAILABLE and isinstance(obj, PILImage.Image):
         return True
     if isinstance(obj, str):
-        if obj.startswith("data:image/"):
+        if obj.startswith("data:"):
             return True
         elif os.path.isfile(obj):
             return True

tests/signatures/test_adapter_image.py

@@ -11,6 +11,8 @@
 from dspy.adapters.types.image import encode_image
 import tempfile
 import pydantic
+import os
+
 
 @pytest.fixture
 def sample_pil_image():
@@ -318,6 +320,114 @@ class OptionalImageSignature(dspy.Signature):
     assert result.output == "Hello"
     assert count_messages_with_image_url_pattern(lm.history[-1]["messages"]) == 0
 
+
+def test_pdf_url_support():
+    """Test support for PDF files from URLs"""
+    pdf_url = "https://www.w3.org/WAI/ER/tests/xhtml/testfiles/resources/pdf/dummy.pdf"
+
+    # Create a dspy.Image object from the PDF URL with download=True
+    pdf_image = dspy.Image.from_url(pdf_url, download=True)
+
+    # The data URI should contain application/pdf in the MIME type
+    assert "data:application/pdf" in pdf_image.url
+    assert ";base64," in pdf_image.url
+
+    # Test using it in a predictor
+    class PDFSignature(dspy.Signature):
+        document: dspy.Image = dspy.InputField(desc="A PDF document")
+        summary: str = dspy.OutputField(desc="A summary of the PDF")
+
+    predictor, lm = setup_predictor(PDFSignature, {"summary": "This is a dummy PDF"})
+    result = predictor(document=pdf_image)
+
+    assert result.summary == "This is a dummy PDF"
+    assert count_messages_with_image_url_pattern(lm.history[-1]["messages"]) == 1
+
+    # Ensure the URL was properly expanded in messages
+    messages_str = str(lm.history[-1]["messages"])
+    assert "application/pdf" in messages_str
+
+
+def test_different_mime_types():
+    """Test support for different file types and MIME type detection"""
+    # Test with various file types
+    file_urls = {
+        "pdf": "https://www.w3.org/WAI/ER/tests/xhtml/testfiles/resources/pdf/dummy.pdf",
+        "image": "https://images.dog.ceo/breeds/dane-great/n02109047_8912.jpg",
+    }
+
+    expected_mime_types = {
+        "pdf": "application/pdf",
+        "image": "image/jpeg",
+    }
+
+    for file_type, url in file_urls.items():
+        # Download and encode
+        encoded = encode_image(url, download_images=True)
+
+        # Check for correct MIME type in the encoded data - using 'in' instead of startswith
+        # to account for possible parameters in the MIME type
+        assert f"data:{expected_mime_types[file_type]}" in encoded
+        assert ";base64," in encoded
+
+
+def test_mime_type_from_response_headers():
+    """Test that MIME types from response headers are correctly used"""
+    # This URL returns proper Content-Type header
+    pdf_url = "https://www.w3.org/WAI/ER/tests/xhtml/testfiles/resources/pdf/dummy.pdf"
+
+    # Make an actual request to get the content type from headers
+    response = requests.get(pdf_url)
+    expected_mime_type = response.headers.get("Content-Type", "")
+
+    # Should be application/pdf or similar
+    assert "pdf" in expected_mime_type.lower()
+
+    # Encode with download to test MIME type from headers
+    encoded = encode_image(pdf_url, download_images=True)
+
+    # The encoded data should contain the correct MIME type
+    assert "application/pdf" in encoded
+    assert ";base64," in encoded
+
+
+def test_pdf_from_file():
+    """Test handling a PDF file from disk"""
+    # Download a PDF to a temporary file
+    pdf_url = "https://www.w3.org/WAI/ER/tests/xhtml/testfiles/resources/pdf/dummy.pdf"
+    response = requests.get(pdf_url)
+    response.raise_for_status()
+
+    with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as tmp_file:
+        tmp_file.write(response.content)
+        tmp_file_path = tmp_file.name
+
+    try:
+        # Create a dspy.Image from the file
+        pdf_image = dspy.Image.from_file(tmp_file_path)
+
+        # Check that the MIME type is correct
+        assert "data:application/pdf" in pdf_image.url
+        assert ";base64," in pdf_image.url
+
+        # Test the image in a predictor
+        class FilePDFSignature(dspy.Signature):
+            document: dspy.Image = dspy.InputField(desc="A PDF document from file")
+            summary: str = dspy.OutputField(desc="A summary of the PDF")
+
+        predictor, lm = setup_predictor(FilePDFSignature, {"summary": "This is a PDF from file"})
+        result = predictor(document=pdf_image)
+
+        assert result.summary == "This is a PDF from file"
+        assert count_messages_with_image_url_pattern(lm.history[-1]["messages"]) == 1
+    finally:
+        # Clean up the temporary file
+        try:
+            os.unlink(tmp_file_path)
+        except:
+            pass
+
+
 def test_image_repr():
     """Test string representation of Image objects"""
     url_image = dspy.Image.from_url("https://example.com/dog.jpg", download=False)