minor adjustment to image testing to support png (#152)

.mailmap

@@ -3,3 +3,4 @@ James Braza <james@futurehouse.org> <jamesbraza@gmail.com>
 Michael Skarlinski <mskarlinski@futurehouse.org> mskarlin <12701035+mskarlin@users.noreply.github.com>
 Ryan-Rhys Griffiths <ryan@futurehouse.org> <ryangriff123@gmail.com>
 Siddharth Narayanan <sid@futurehouse.org> <sidnarayanan@users.noreply.github.com>
+Ludovico Mitchener <ludo@futurehouse.org> ludomitch

.pre-commit-config.yaml

@@ -47,6 +47,7 @@ repos:
       - id: codespell
         additional_dependencies: [".[toml]"]
         exclude_types: [jupyter]
+        exclude: '.*\.b64$'
   - repo: https://github.com/pappasam/toml-sort
     rev: v0.24.2
     hooks:

src/aviary/message.py

@@ -6,7 +6,7 @@
 
 from pydantic import BaseModel, Field, field_validator, model_validator
 
-from aviary.utils import check_if_valid_base64, encode_image_to_base64
+from aviary.utils import encode_image_to_base64, validate_base64_image
 
 if TYPE_CHECKING:
     from logging import LogRecord
@@ -147,7 +147,7 @@ def create_message(
                 {
                     "type": "image_url",
                     "image_url": {
-                        "url": check_if_valid_base64(image)
+                        "url": validate_base64_image(image)
                         # If image is a string, assume it's already a base64 encoded image
                         if isinstance(image, str)
                         else encode_image_to_base64(image)

src/aviary/utils.py

@@ -70,10 +70,12 @@ def encode_image_to_base64(img: "np.ndarray") -> str:
     )
 
 
-def check_if_valid_base64(image: str) -> str:
-    """Check if the input string is a valid base64 encoded image."""
+def validate_base64_image(image: str) -> str:
+    """Validate if the input string is a valid base64 encoded image and if it is, return the image."""
     try:
-        base64.b64decode(image)
+        # Support for inclusion of the data:image/ url prefix
+        test_image = image.split(",")[1] if image.startswith("data:image/") else image
+        base64.b64decode(test_image)
     except Exception as err:
         raise ValueError("Invalid base64 encoded image") from err
     return image

tests/test_messages.py

@@ -136,7 +136,7 @@ def test_dump(self, message: Message, expected: dict) -> None:
             (
                 [
                     np.zeros((32, 32, 3), dtype=np.uint8),  # red square
-                    load_base64_image("sample_image1.b64"),
+                    load_base64_image("sample_jpeg_image.b64"),
                 ],
                 "What color are these squares? List each color.",
                 None,
@@ -151,7 +151,14 @@ def test_dump(self, message: Message, expected: dict) -> None:
             ),
             # Case 4: A string should be converted to a base64 encoded image
             (
-                load_base64_image("sample_image1.b64"),
+                load_base64_image("sample_jpeg_image.b64"),
+                "What color is this square?",
+                None,
+                2,  # 1 image + 1 text
+            ),
+            # Case 5: A PNG image should be converted to a base64 encoded image
+            (
+                load_base64_image("sample_png_image.b64"),
                 "What color is this square?",
                 None,
                 2,  # 1 image + 1 text