From bb3c7cbedfcd865f9c13e1e9b1d996fe2863ade0 Mon Sep 17 00:00:00 2001
From: Johnny Chadda <johnny@chadda.se>
Date: Mon, 19 May 2025 11:35:32 +0200
Subject: [PATCH] Get proper mime type for images

---
 src/opperai/types/__init__.py | 47 ++++++++++++++++++++++++-----------
 tests/test_types.py           | 39 ++++++++++++++++++++++++++++-
 2 files changed, 70 insertions(+), 16 deletions(-)

diff --git a/src/opperai/types/__init__.py b/src/opperai/types/__init__.py
index 0198204..d515647 100644
--- a/src/opperai/types/__init__.py
+++ b/src/opperai/types/__init__.py
@@ -2,6 +2,7 @@
 from __future__ import annotations
 
 import base64
+import mimetypes
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Union
@@ -167,25 +168,41 @@ class ImageInput(BaseModel):
     @computed_field
     @property
     def _opper_image_input(self) -> str:
-        if self.path:
-            suffix = self.path.suffix
-            if suffix == ".png":
-                with open(self.path, "rb") as image_file:
-                    base64_image = base64.b64encode(image_file.read()).decode("utf-8")
-                data = f"data:image/png;base64,{base64_image}"
-            elif suffix in (".jpg", ".jpeg"):
-                with open(self.path, "rb") as image_file:
-                    base64_image = base64.b64encode(image_file.read()).decode("utf-8")
-                data = f"data:image/jpeg;base64,{base64_image}"
+        if not self.path:
+            raise ValueError("no path provided")
+
+        # Initialize MIME types database if not already done
+        if not mimetypes.inited:
+            mimetypes.init()
+
+        # Read the image file
+        with open(self.path, "rb") as image_file:
+            image_data = image_file.read()
+
+        # Get MIME type based on file extension
+        mime_type, _ = mimetypes.guess_type(str(self.path))
+
+        # Verify it's a supported image type
+        if mime_type == "image/png":
+            data_uri_mime = "image/png"
+        elif mime_type in ("image/jpeg", "image/jpg"):
+            data_uri_mime = "image/jpeg"
+        else:
+            # Fallback to checking file headers if mime_type is None or not an image
+            if image_data.startswith(b"\x89PNG\r\n\x1a\n"):
+                # PNG signature
+                data_uri_mime = "image/png"
+            elif image_data.startswith(b"\xff\xd8\xff"):
+                # JPEG signature
+                data_uri_mime = "image/jpeg"
             else:
                 raise ValueError(
-                    "File type not supported. Supported types: .png, .jpg, .jpeg"
+                    "File type not supported. Supported types: png, jpg, jpeg"
                 )
 
-        if not data:
-            raise ValueError("no path or url provided")
-
-        return data
+        # Encode image data
+        base64_image = base64.b64encode(image_data).decode("utf-8")
+        return f"data:{data_uri_mime};base64,{base64_image}"
 
     @classmethod
     def from_path(cls, path: FilePath):
diff --git a/tests/test_types.py b/tests/test_types.py
index e50d3ae..5393463 100644
--- a/tests/test_types.py
+++ b/tests/test_types.py
@@ -1,5 +1,9 @@
+import os
+from pathlib import Path
+
 import pytest
-from opperai.types import validate_uuid_xor_path
+
+from opperai.types import ImageInput, validate_uuid_xor_path
 
 
 def test_id_xor_path():
@@ -19,3 +23,36 @@ def test_function(uuid=None, path=None):
 
     with pytest.raises(ValueError):
         test_function(uuid=1, path="test", other="other")
+
+
+def test_image_input_mime_detection():
+    # Test PNG detection
+    png_path = Path("tests/fixtures/images/minimal.png")
+    png_input = ImageInput(path=png_path)
+    png_result = png_input._opper_image_input
+    assert "data:image/png;base64," in png_result
+
+    # Test JPEG detection
+    jpg_path = Path("tests/fixtures/images/fossil.jpg")
+    jpg_input = ImageInput(path=jpg_path)
+    jpg_result = jpg_input._opper_image_input
+    assert "data:image/jpeg;base64," in jpg_result
+
+    # Test invalid path
+    with pytest.raises(ValueError, match="no path provided"):
+        invalid_input = ImageInput()
+        invalid_input._opper_image_input
+
+    # Test unsupported file type (create a temporary text file)
+    import tempfile
+
+    with tempfile.NamedTemporaryFile(suffix=".txt", delete=False) as tmp:
+        tmp.write(b"This is a text file, not an image")
+        tmp_path = tmp.name
+
+    try:
+        with pytest.raises(ValueError, match="File type not supported"):
+            text_input = ImageInput(path=Path(tmp_path))
+            text_input._opper_image_input
+    finally:
+        os.unlink(tmp_path)  # Clean up the temporary file