fix(vertex_ai): label TTS audio bytes with their real content-type

litellm/litellm_core_utils/audio_utils/utils.py

@@ -11,6 +11,7 @@ from litellm.types.files import (
     AUDIO_FILE_TYPES,
     FILE_EXTENSIONS,
     FILE_MIME_TYPES,
+    FileType,
     get_file_mime_type_from_extension,
 )
 from litellm.types.utils import FileTypes
@@ -351,3 +352,24 @@ def resolve_speech_media_type(upstream_content_type: str | None, response_format
         None if response_format is None else _speech_media_type_for_response_format(response_format)
     )
     return requested_media_type or DEFAULT_SPEECH_MEDIA_TYPE
+
+
+_OGG_OPUS_HEAD_WINDOW: Final = 64
+_MPEG_FRAME_SYNC_MASK: Final = 0xE0
+_MPEG_FRAME_LAYER_MASK: Final = 0x06
+
+
+def speech_media_type_from_audio_bytes(audio: bytes) -> str | None:
+    if audio[:4] == b"RIFF" and audio[8:12] == b"WAVE":
+        return FILE_MIME_TYPES[FileType.WAV]
+    if audio[:4] == b"fLaC":
+        return FILE_MIME_TYPES[FileType.FLAC]
+    if audio[:4] == b"OggS":
+        is_opus: Final = b"OpusHead" in audio[:_OGG_OPUS_HEAD_WINDOW]
+        return FILE_MIME_TYPES[FileType.OPUS if is_opus else FileType.OGG]
+    if audio[:3] == b"ID3":
+        return FILE_MIME_TYPES[FileType.MP3]
+    if len(audio) < 2 or audio[0] != 0xFF or (audio[1] & _MPEG_FRAME_SYNC_MASK) != _MPEG_FRAME_SYNC_MASK:
+        return None
+    is_adts_aac: Final = (audio[1] & _MPEG_FRAME_LAYER_MASK) == 0
+    return FILE_MIME_TYPES[FileType.AAC if is_adts_aac else FileType.MP3]

litellm/llms/vertex_ai/text_to_speech/transformation.py

@@ -11,6 +11,9 @@ from typing import TYPE_CHECKING, Any, Final, Union
 
 import httpx
 
+from litellm.litellm_core_utils.audio_utils.utils import (
+    speech_media_type_from_audio_bytes,
+)
 from litellm.llms.base_llm.text_to_speech.transformation import (
     BaseTextToSpeechConfig,
     TextToSpeechRequestData,
@@ -457,12 +460,11 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase):
         if not response_content:
             raise ValueError("No audioContent in Vertex AI TTS response")
 
-        # Decode base64 to get binary content
         binary_data: Final = base64.b64decode(response_content)
-
-        # Create an httpx.Response object with the binary data
+        media_type: Final = speech_media_type_from_audio_bytes(binary_data)
         response: Final = httpx.Response(
             status_code=200,
+            headers={} if media_type is None else {"content-type": media_type},
             content=binary_data,
         )
 

tests/test_litellm/litellm_core_utils/audio_utils/test_utils.py

@@ -1,30 +0,0 @@
-import pytest
-
-from litellm.litellm_core_utils.audio_utils.utils import resolve_speech_media_type
-
-
-@pytest.mark.parametrize(
-    ("upstream_content_type", "response_format", "expected"),
-    [
-        ("audio/wav", None, "audio/wav"),
-        ("AUDIO/WAV", None, "audio/wav"),
-        ("audio/flac; charset=binary", "mp3", "audio/flac"),
-        ("application/json", "flac", "audio/flac"),
-        ("application/octet-stream", "pcm", "audio/pcm"),
-        (None, "wav", "audio/wav"),
-        (None, "WAV", "audio/wav"),
-        (None, "opus", "audio/opus"),
-        (None, "aac", "audio/aac"),
-        (None, "mp3", "audio/mpeg"),
-        (None, "mp4", "audio/mpeg"),
-        (None, "bogus", "audio/mpeg"),
-        (None, None, "audio/mpeg"),
-        ("", None, "audio/mpeg"),
-    ],
-)
-def test_resolve_speech_media_type(upstream_content_type, response_format, expected):
-    resolved = resolve_speech_media_type(
-        upstream_content_type=upstream_content_type,
-        response_format=response_format,
-    )
-    assert resolved == expected

tests/test_litellm/litellm_core_utils/test_audio_utils.py

@@ -347,3 +347,59 @@ class TestNormalizeTranscriptionLanguageToBcp47:
         )
 
         assert normalize_transcription_language_to_bcp47(language) == expected
+
+
+class TestResolveSpeechMediaType:
+    @pytest.mark.parametrize(
+        ("upstream_content_type", "response_format", "expected"),
+        [
+            ("audio/wav", None, "audio/wav"),
+            ("AUDIO/WAV", None, "audio/wav"),
+            ("audio/flac; charset=binary", "mp3", "audio/flac"),
+            ("application/json", "flac", "audio/flac"),
+            ("application/octet-stream", "pcm", "audio/pcm"),
+            (None, "wav", "audio/wav"),
+            (None, "WAV", "audio/wav"),
+            (None, "opus", "audio/opus"),
+            (None, "aac", "audio/aac"),
+            (None, "mp3", "audio/mpeg"),
+            (None, "mp4", "audio/mpeg"),
+            (None, "bogus", "audio/mpeg"),
+            (None, None, "audio/mpeg"),
+            ("", None, "audio/mpeg"),
+        ],
+    )
+    def test_resolution(self, upstream_content_type, response_format, expected):
+        from litellm.litellm_core_utils.audio_utils.utils import resolve_speech_media_type
+
+        resolved = resolve_speech_media_type(
+            upstream_content_type=upstream_content_type,
+            response_format=response_format,
+        )
+        assert resolved == expected
+
+
+class TestSpeechMediaTypeFromAudioBytes:
+    @pytest.mark.parametrize(
+        ("audio", "expected"),
+        [
+            (b"RIFF\x24\x00\x00\x00WAVEfmt ", "audio/wav"),
+            (b"fLaC\x00\x00\x00\x22", "audio/flac"),
+            (b"OggS" + b"\x00" * 24 + b"OpusHead", "audio/opus"),
+            (b"OggS" + b"\x00" * 24 + b"\x01vorbis", "audio/ogg"),
+            (b"ID3\x04\x00\x00\x00\x00\x00\x00", "audio/mpeg"),
+            (b"\xff\xfb\x90\x64", "audio/mpeg"),
+            (b"\xff\xf3\x80\x00", "audio/mpeg"),
+            (b"\xff\xf1\x50\x80", "audio/aac"),
+            (b"\xff\xf9\x50\x80", "audio/aac"),
+            (b"RIFF\x24\x00\x00\x00AVI LIST", None),
+            (b"\xff\x00\x00\x00", None),
+            (b"\x00\x01\x02\x03\x04\x05", None),
+            (b"\xff", None),
+            (b"", None),
+        ],
+    )
+    def test_sniffing(self, audio, expected):
+        from litellm.litellm_core_utils.audio_utils.utils import speech_media_type_from_audio_bytes
+
+        assert speech_media_type_from_audio_bytes(audio) == expected

tests/test_litellm/llms/vertex_ai/text_to_speech/test_transformation.py

@@ -1,3 +1,4 @@
+import base64
 from unittest.mock import MagicMock, Mock, patch
 
 import httpx
@@ -126,6 +127,48 @@ class TestVertexAITextToSpeechConfig:
         assert voice_dict == voice_input
 
 
+@pytest.mark.parametrize(
+    ("audio", "expected_content_type"),
+    [
+        (b"RIFF\x24\x00\x00\x00WAVEfmt \x10\x00\x00\x00", "audio/wav"),
+        (b"\xff\xfb\x90\x64\x00\x00\x00\x00", "audio/mpeg"),
+        (b"OggS" + b"\x00" * 24 + b"OpusHead", "audio/opus"),
+        (b"fLaC\x00\x00\x00\x22", "audio/flac"),
+    ],
+)
+def test_transform_text_to_speech_response_labels_content_type(audio, expected_content_type):
+    raw_response = httpx.Response(
+        status_code=200,
+        json={"audioContent": base64.b64encode(audio).decode()},
+    )
+
+    result = VertexAITextToSpeechConfig().transform_text_to_speech_response(
+        model="vertex_ai/chirp",
+        raw_response=raw_response,
+        logging_obj=MagicMock(),
+    )
+
+    assert result.response.headers["content-type"] == expected_content_type
+    assert result.response.content == audio
+
+
+def test_transform_text_to_speech_response_leaves_unknown_bytes_unlabeled():
+    raw_pcm = b"\x00\x01\x02\x03\x04\x05\x06\x07"
+    raw_response = httpx.Response(
+        status_code=200,
+        json={"audioContent": base64.b64encode(raw_pcm).decode()},
+    )
+
+    result = VertexAITextToSpeechConfig().transform_text_to_speech_response(
+        model="vertex_ai/chirp",
+        raw_response=raw_response,
+        logging_obj=MagicMock(),
+    )
+
+    assert "content-type" not in result.response.headers
+    assert result.response.content == raw_pcm
+
+
 @patch("litellm.llms.custom_httpx.llm_http_handler.HTTPHandler.post")
 @patch.object(VertexAITextToSpeechConfig, "_ensure_access_token")
 @patch.object(VertexAITextToSpeechConfig, "_get_token_and_url")