Handle audio data URL MIME variants

Author: fus3r

src/mistral_common/protocol/instruct/chunk.py

@@ -22,7 +22,7 @@
 
 def _strip_audio_data_url_prefix(data: str) -> str:
     r"""Remove the optional base64 audio data URL prefix."""
-    if re.match(r"^data:audio/\w+;base64,", data):
+    if re.match(r"^data:audio/[^;,]+(?:;[^;,]*)*;base64,", data):
         return data.split(",", 1)[1]
     return data
 

src/mistral_common/tokens/tokenizers/audio.py

@@ -2,7 +2,6 @@
 import io
 import logging
 import math
-import re
 import warnings
 from dataclasses import dataclass
 from enum import Enum
@@ -19,7 +18,12 @@
     is_soundfile_installed,
     is_soxr_installed,
 )
-from mistral_common.protocol.instruct.chunk import AudioChunk, AudioURLChunk, AudioURLType
+from mistral_common.protocol.instruct.chunk import (
+    AudioChunk,
+    AudioURLChunk,
+    AudioURLType,
+    _strip_audio_data_url_prefix,
+)
 
 if is_soxr_installed():
     import soxr
@@ -114,8 +118,7 @@ def from_base64(audio_base64: str, strict: bool = True) -> "Audio":
         """
         assert_soundfile_installed()
 
-        if re.match(r"^data:audio/\w+;base64,", audio_base64):
-            audio_base64 = audio_base64.split(",")[1]
+        audio_base64 = _strip_audio_data_url_prefix(audio_base64)
 
         try:
             audio_bytes = base64.b64decode(audio_base64)

tests/test_audio.py

@@ -144,6 +144,15 @@ def rmse(a: np.ndarray, b: np.ndarray) -> float:
             raise ValueError(f"Unknown format {format}")
 
 
+@pytest.mark.parametrize("mime_type", ["audio/x-wav", "audio/vnd.wave", "audio/wav;codec=pcm"])
+def test_audio_from_audio_chunk_accepts_data_url_mime_variants(mime_type: str) -> None:
+    b64 = _make_dummy_base64()
+    audio = Audio.from_audio_chunk(AudioChunk(input_audio=f"data:{mime_type};base64,{b64}"))
+
+    assert audio.sampling_rate == 16000
+    assert audio.format == "wav"
+
+
 @pytest.mark.parametrize(
     "freq, expected_mel",
     [

tests/test_converters.py

@@ -1,6 +1,7 @@
 import base64
 import copy
 import io
+import time
 import warnings
 from pathlib import Path
 from typing import Any
@@ -44,6 +45,7 @@
     ImageURLChunk,
     TextChunk,
     ThinkChunk,
+    _strip_audio_data_url_prefix,
 )
 from mistral_common.protocol.instruct.messages import (
     AssistantMessage,
@@ -167,6 +169,27 @@ def test_convert_input_audio_chunk() -> None:
     assert AudioChunk.from_openai(typeddict_openai) == chunk
 
 
+@pytest.mark.parametrize("mime_type", ["audio/x-wav", "audio/vnd.wave", "audio/wav;codec=pcm"])
+def test_convert_input_audio_chunk_with_data_url_mime_variants(mime_type: str) -> None:
+    input_audio = DUMMY_AUDIO_CHUNK.input_audio
+    assert isinstance(input_audio, str)
+    chunk = AudioChunk(input_audio=f"data:{mime_type};base64,{input_audio}")
+
+    openai_dict = chunk.to_openai()
+
+    assert openai_dict["input_audio"]["data"] == input_audio
+    assert openai_dict["input_audio"]["format"] == "wav"
+
+
+def test_audio_data_url_prefix_malformed_parameters_stays_linear() -> None:
+    malformed = f"data:audio/wav{';' * 30}x"
+
+    start = time.monotonic()
+
+    assert _strip_audio_data_url_prefix(malformed) == malformed
+    assert time.monotonic() - start < 1
+
+
 @pytest.mark.parametrize(
     ["openai_image_url_chunk", "image_url_chunk"],
     [