Fix ImageUrl, VideoUrl, AudioUrl and DocumentUrl not being serializable (#2422)

Author: DouweM

pydantic_ai_slim/pydantic_ai/messages.py

@@ -106,7 +106,7 @@ class FileUrl(ABC):
     - `GoogleModel`: `VideoUrl.vendor_metadata` is used as `video_metadata`: https://ai.google.dev/gemini-api/docs/video-understanding#customize-video-processing
     """
 
-    _media_type: str | None = field(init=False, repr=False)
+    _media_type: str | None = field(init=False, repr=False, compare=False)
 
     def __init__(
         self,
@@ -120,19 +120,21 @@ def __init__(
         self.force_download = force_download
         self._media_type = media_type
 
-    @abstractmethod
-    def _infer_media_type(self) -> str:
-        """Return the media type of the file, based on the url."""
-
     @property
     def media_type(self) -> str:
-        """Return the media type of the file, based on the url or the provided `_media_type`."""
+        """Return the media type of the file, based on the URL or the provided `media_type`."""
         return self._media_type or self._infer_media_type()
 
+    @abstractmethod
+    def _infer_media_type(self) -> str:
+        """Infer the media type of the file based on the URL."""
+        raise NotImplementedError
+
     @property
     @abstractmethod
     def format(self) -> str:
         """The file format."""
+        raise NotImplementedError
 
     __repr__ = _utils.dataclasses_no_defaults_repr
 
@@ -182,7 +184,9 @@ def _infer_media_type(self) -> VideoMediaType:
         elif self.is_youtube:
             return 'video/mp4'
         else:
-            raise ValueError(f'Unknown video file extension: {self.url}')
+            raise ValueError(
+                f'Could not infer media type from video URL: {self.url}. Explicitly provide a `media_type` instead.'
+            )
 
     @property
     def is_youtube(self) -> bool:
@@ -238,7 +242,9 @@ def _infer_media_type(self) -> AudioMediaType:
         if self.url.endswith('.aac'):
             return 'audio/aac'
 
-        raise ValueError(f'Unknown audio file extension: {self.url}')
+        raise ValueError(
+            f'Could not infer media type from audio URL: {self.url}. Explicitly provide a `media_type` instead.'
+        )
 
     @property
     def format(self) -> AudioFormat:
@@ -278,7 +284,9 @@ def _infer_media_type(self) -> ImageMediaType:
         elif self.url.endswith('.webp'):
             return 'image/webp'
         else:
-            raise ValueError(f'Unknown image file extension: {self.url}')
+            raise ValueError(
+                f'Could not infer media type from image URL: {self.url}. Explicitly provide a `media_type` instead.'
+            )
 
     @property
     def format(self) -> ImageFormat:
@@ -324,10 +332,16 @@ def _infer_media_type(self) -> str:
             return 'application/pdf'
         elif self.url.endswith('.rtf'):
             return 'application/rtf'
+        elif self.url.endswith('.docx'):
+            return 'application/vnd.openxmlformats-officedocument.wordprocessingml.document'
+        elif self.url.endswith('.xlsx'):
+            return 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet'
 
         type_, _ = guess_type(self.url)
         if type_ is None:
-            raise ValueError(f'Unknown document file extension: {self.url}')
+            raise ValueError(
+                f'Could not infer media type from document URL: {self.url}. Explicitly provide a `media_type` instead.'
+            )
         return type_
 
     @property

tests/test_agent.py

@@ -2922,7 +2922,7 @@ def validate_output(ctx: RunContext[None], o: Any) -> Any:  # pragma: no cover
         agent.run_sync('Hello', output_type=int)
 
 
-def test_binary_content_all_messages_json():
+def test_binary_content_serializable():
     agent = Agent('test')
 
     content = BinaryContent(data=b'Hello', media_type='text/plain')
@@ -2974,6 +2974,57 @@ def test_binary_content_all_messages_json():
     assert messages == result.all_messages()
 
 
+def test_image_url_serializable():
+    agent = Agent('test')
+
+    content = ImageUrl('https://example.com/chart', media_type='image/jpeg')
+    result = agent.run_sync(['Hello', content])
+
+    serialized = result.all_messages_json()
+    assert json.loads(serialized) == snapshot(
+        [
+            {
+                'parts': [
+                    {
+                        'content': [
+                            'Hello',
+                            {
+                                'url': 'https://example.com/chart',
+                                'force_download': False,
+                                'vendor_metadata': None,
+                                'kind': 'image-url',
+                            },
+                        ],
+                        'timestamp': IsStr(),
+                        'part_kind': 'user-prompt',
+                    }
+                ],
+                'instructions': None,
+                'kind': 'request',
+            },
+            {
+                'parts': [{'content': 'success (no tool calls)', 'part_kind': 'text'}],
+                'usage': {
+                    'requests': 1,
+                    'request_tokens': 51,
+                    'response_tokens': 4,
+                    'total_tokens': 55,
+                    'details': None,
+                },
+                'model_name': 'test',
+                'timestamp': IsStr(),
+                'kind': 'response',
+                'vendor_details': None,
+                'vendor_id': None,
+            },
+        ]
+    )
+
+    # We also need to be able to round trip the serialized messages.
+    messages = ModelMessagesTypeAdapter.validate_json(serialized)
+    assert messages == result.all_messages()
+
+
 def test_tool_return_part_binary_content_serialization():
     """Test that ToolReturnPart can properly serialize BinaryContent."""
     png_data = b'\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01\x08\x02\x00\x00\x00\x90wS\xde\x00\x00\x00\x0cIDATx\x9cc```\x00\x00\x00\x04\x00\x01\xf6\x178\x00\x00\x00\x00IEND\xaeB`\x82'

tests/test_messages.py

@@ -153,7 +153,7 @@ def test_audio_url(audio_url: AudioUrl, media_type: str, format: str):
 
 
 def test_audio_url_invalid():
-    with pytest.raises(ValueError, match='Unknown audio file extension: foobar.potato'):
+    with pytest.raises(ValueError, match='Could not infer media type from audio URL: foobar.potato'):
         AudioUrl('foobar.potato').media_type
 
 
@@ -173,10 +173,10 @@ def test_image_url_formats(image_url: ImageUrl, media_type: str, format: str):
 
 
 def test_image_url_invalid():
-    with pytest.raises(ValueError, match='Unknown image file extension: foobar.potato'):
+    with pytest.raises(ValueError, match='Could not infer media type from image URL: foobar.potato'):
         ImageUrl('foobar.potato').media_type
 
-    with pytest.raises(ValueError, match='Unknown image file extension: foobar.potato'):
+    with pytest.raises(ValueError, match='Could not infer media type from image URL: foobar.potato'):
         ImageUrl('foobar.potato').format
 
 
@@ -213,7 +213,7 @@ def test_document_url_formats(document_url: DocumentUrl, media_type: str, format
 
 
 def test_document_url_invalid():
-    with pytest.raises(ValueError, match='Unknown document file extension: foobar.potato'):
+    with pytest.raises(ValueError, match='Could not infer media type from document URL: foobar.potato'):
         DocumentUrl('foobar.potato').media_type
 
     with pytest.raises(ValueError, match='Unknown document media type: text/x-python'):
@@ -301,7 +301,7 @@ def test_video_url_formats(video_url: VideoUrl, media_type: str, format: str):
 
 
 def test_video_url_invalid():
-    with pytest.raises(ValueError, match='Unknown video file extension: foobar.potato'):
+    with pytest.raises(ValueError, match='Could not infer media type from video URL: foobar.potato'):
         VideoUrl('foobar.potato').media_type