Feature: 🚀 Add Audio Content Support to OpenAISpec Request (#439)

* feat: add AudioContent and InputAudio models to support audio data in chat messages

* chore: add fixtures for OpenAI request data with audio support in WAV and FLAC formats

* test: add end-to-end tests for OpenAI Spec audio input handling with different audio formats

* test: rename test function for audio input parity with OpenAI Spec

* fix: update comment style for data attribute in InputAudio model

* refactor: remove unnecessary line breaks

Author: bhimrazy

src/litserve/specs/openai.py

@@ -77,6 +77,16 @@ class ImageContent(BaseModel):
     image_url: Union[str, ImageContentURL]
 
 
+class InputAudio(BaseModel):
+    data: str  # base64 encoded audio data.
+    format: Literal["wav", "mp3"]
+
+
+class AudioContent(BaseModel):
+    type: Literal["input_audio"]
+    input_audio: InputAudio
+
+
 class Function(BaseModel):
     name: str
     description: str
@@ -133,7 +143,7 @@ class ResponseFormatJSONSchema(BaseModel):
 
 class ChatMessage(BaseModel):
     role: str
-    content: Optional[Union[str, List[Union[TextContent, ImageContent]]]] = None
+    content: Optional[Union[str, List[Union[TextContent, ImageContent, AudioContent]]]] = None
     name: Optional[str] = None
     tool_calls: Optional[List[ToolCall]] = None
     tool_call_id: Optional[str] = None

tests/conftest.py

@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import base64
 import time
 from typing import Generator
 
@@ -193,6 +194,44 @@ def openai_request_data_with_image():
     }
 
 
+@pytest.fixture
+def openai_request_data_with_audio_wav(openai_request_data):
+    # Create a base64 encoded string from a list of audio data
+    audio_data = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
+    encoded_string = base64.b64encode(bytearray(audio_data)).decode("utf-8")
+
+    request_data = openai_request_data.copy()
+    request_data["messages"] = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "What is in this recording?"},
+                {"type": "input_audio", "input_audio": {"data": encoded_string, "format": "wav"}},
+            ],
+        },
+    ]
+    return request_data
+
+
+@pytest.fixture
+def openai_request_data_with_audio_flac(openai_request_data):
+    # Create a base64 encoded string from a list of audio data
+    audio_data = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
+    encoded_string = base64.b64encode(bytearray(audio_data)).decode("utf-8")
+
+    request_data = openai_request_data.copy()
+    request_data["messages"] = [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "What is in this recording?"},
+                {"type": "input_audio", "input_audio": {"data": encoded_string, "format": "flac"}},
+            ],
+        },
+    ]
+    return request_data
+
+
 @pytest.fixture
 def openai_request_data_with_tools():
     return {

tests/e2e/test_e2e.py

@@ -183,6 +183,34 @@ def test_openai_parity_with_image_input():
         )
 
 
+@e2e_from_file("tests/e2e/default_openaispec.py")
+def test_openai_parity_with_audio_input(openai_request_data_with_audio_wav):
+    client = OpenAI(
+        base_url="http://127.0.0.1:8000/v1",
+        api_key="lit",  # required, but unused
+    )
+    messages = openai_request_data_with_audio_wav["messages"]
+    response = client.chat.completions.create(
+        model="lit",
+        messages=messages,
+    )
+    assert response.choices[0].message.content == "This is a generated output", (
+        f"Server didn't return expected output\nOpenAI client output: {response}"
+    )
+
+    response = client.chat.completions.create(
+        model="lit",
+        messages=messages,
+        stream=True,
+    )
+
+    expected_outputs = ["This is a generated output", None]
+    for r, expected_out in zip(response, expected_outputs):
+        assert r.choices[0].delta.content == expected_out, (
+            f"Server didn't return expected output.\nOpenAI client output: {r}"
+        )
+
+
 @e2e_from_file("tests/e2e/default_openaispec_tools.py")
 def test_openai_parity_with_tools():
     client = OpenAI(

tests/test_specs.py

@@ -147,6 +147,30 @@ async def test_openai_spec_with_image(openai_request_data_with_image):
             )
 
 
+@pytest.mark.asyncio
+async def test_openai_spec_with_audio(openai_request_data_with_audio_wav, openai_request_data_with_audio_flac):
+    server = ls.LitServer(TestAPI(), spec=OpenAISpec())
+
+    with wrap_litserve_start(server) as server:
+        async with LifespanManager(server.app) as manager, AsyncClient(
+            transport=ASGITransport(app=manager.app), base_url="http://test"
+        ) as ac:
+            resp = await ac.post("/v1/chat/completions", json=openai_request_data_with_audio_wav, timeout=10)
+            assert resp.status_code == 200, "Status code should be 200"
+
+            assert resp.json()["choices"][0]["message"]["content"] == "This is a generated output", (
+                "LitAPI predict response should match with the generated output"
+            )
+
+            # test for unsupported audio format
+            resp = await ac.post("/v1/chat/completions", json=openai_request_data_with_audio_flac, timeout=10)
+            assert resp.status_code == 422, "Status code should be 422"
+            errors = resp.json()["detail"]
+            assert any(error["msg"] == "Input should be 'wav' or 'mp3'" for error in errors), (
+                "Error message for unsupported audio format should be present"
+            )
+
+
 @pytest.mark.asyncio
 async def test_override_encode(openai_request_data):
     server = ls.LitServer(TestAPIWithCustomEncode(), spec=OpenAISpec())