Add all audio types supported by Gemini to AudioUrl (#2151)

Author: ChenghaoMou

pydantic_ai_slim/pydantic_ai/messages.py

@@ -25,7 +25,7 @@
     from .models.instrumented import InstrumentationSettings
 
 
-AudioMediaType: TypeAlias = Literal['audio/wav', 'audio/mpeg']
+AudioMediaType: TypeAlias = Literal['audio/wav', 'audio/mpeg', 'audio/ogg', 'audio/flac', 'audio/aiff', 'audio/aac']
 ImageMediaType: TypeAlias = Literal['image/jpeg', 'image/png', 'image/gif', 'image/webp']
 DocumentMediaType: TypeAlias = Literal[
     'application/pdf',
@@ -48,7 +48,7 @@
     'video/3gpp',
 ]
 
-AudioFormat: TypeAlias = Literal['wav', 'mp3']
+AudioFormat: TypeAlias = Literal['wav', 'mp3', 'oga', 'flac', 'aiff', 'aac']
 ImageFormat: TypeAlias = Literal['jpeg', 'png', 'gif', 'webp']
 DocumentFormat: TypeAlias = Literal['csv', 'doc', 'docx', 'html', 'md', 'pdf', 'txt', 'xls', 'xlsx']
 VideoFormat: TypeAlias = Literal['mkv', 'mov', 'mp4', 'webm', 'flv', 'mpeg', 'mpg', 'wmv', 'three_gp']
@@ -182,13 +182,25 @@ class AudioUrl(FileUrl):
 
     @property
     def media_type(self) -> AudioMediaType:
-        """Return the media type of the audio file, based on the url."""
+        """Return the media type of the audio file, based on the url.
+
+        References:
+        - Gemini: https://ai.google.dev/gemini-api/docs/audio#supported-formats
+        """
         if self.url.endswith('.mp3'):
             return 'audio/mpeg'
-        elif self.url.endswith('.wav'):
+        if self.url.endswith('.wav'):
             return 'audio/wav'
-        else:
-            raise ValueError(f'Unknown audio file extension: {self.url}')
+        if self.url.endswith('.flac'):
+            return 'audio/flac'
+        if self.url.endswith('.oga'):
+            return 'audio/ogg'
+        if self.url.endswith('.aiff'):
+            return 'audio/aiff'
+        if self.url.endswith('.aac'):
+            return 'audio/aac'
+
+        raise ValueError(f'Unknown audio file extension: {self.url}')
 
     @property
     def format(self) -> AudioFormat:
@@ -358,6 +370,10 @@ class ToolReturn:
 _audio_format_lookup: dict[str, AudioFormat] = {
     'audio/mpeg': 'mp3',
     'audio/wav': 'wav',
+    'audio/flac': 'flac',
+    'audio/ogg': 'oga',
+    'audio/aiff': 'aiff',
+    'audio/aac': 'aac',
 }
 _image_format_lookup: dict[str, ImageFormat] = {
     'image/jpeg': 'jpeg',

tests/evals/test_otel.py

@@ -65,9 +65,10 @@ async def task2():
     task1_root = tree1.roots[0]
     assert len(task1_root.children) == 2, 'task1 should have exactly two children'
     task1_child_names = {child.name for child in task1_root.children}
-    assert task1_child_names == {'task1_child1', 'task1_child2'}, (
-        "task1's children should be task1_child1 and task1_child2"
-    )
+    assert task1_child_names == {
+        'task1_child1',
+        'task1_child2',
+    }, "task1's children should be task1_child1 and task1_child2"
 
     # Verify that tree2 only contains spans from task2
     assert len(tree2.roots) == 1, 'tree2 should have exactly one root span'

tests/test_messages.py

@@ -125,6 +125,10 @@ def test_binary_content_document(media_type: str, format: str):
     [
         pytest.param(AudioUrl('foobar.mp3'), 'audio/mpeg', 'mp3', id='mp3'),
         pytest.param(AudioUrl('foobar.wav'), 'audio/wav', 'wav', id='wav'),
+        pytest.param(AudioUrl('foobar.oga'), 'audio/ogg', 'oga', id='oga'),
+        pytest.param(AudioUrl('foobar.flac'), 'audio/flac', 'flac', id='flac'),
+        pytest.param(AudioUrl('foobar.aiff'), 'audio/aiff', 'aiff', id='aiff'),
+        pytest.param(AudioUrl('foobar.aac'), 'audio/aac', 'aac', id='aac'),
     ],
 )
 def test_audio_url(audio_url: AudioUrl, media_type: str, format: str):