diff --git a/tests/entrypoints/openai_api/test_chat_audio_choice_merge.py b/tests/entrypoints/openai_api/test_chat_audio_choice_merge.py new file mode 100644 index 000000000..f24a2582c --- /dev/null +++ b/tests/entrypoints/openai_api/test_chat_audio_choice_merge.py @@ -0,0 +1,135 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Tests for merging audio final outputs into text choices. + +Covers #7376: a non-streaming chat request with +``modalities=["text", "audio"]`` used to return two choices with duplicate +``index=0`` — one carrying the text and one carrying the audio — so clients +reading ``choices[0]`` silently lost the audio. +""" + +from __future__ import annotations + +import json + +import pytest +from openai.types.chat.chat_completion_audio import ChatCompletionAudio as OpenAIChatCompletionAudio +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionResponseChoice, + ChatMessage, +) +from vllm.entrypoints.serve.engine.protocol import UsageInfo + +from vllm_omni.entrypoints.openai.protocol.audio import AudioChunkMetadata +from vllm_omni.entrypoints.openai.protocol.chat_completion import ( + OmniChatCompletionResponse, + OmniChatCompletionResponseChoice, +) +from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _text_choice(index: int, content: str) -> ChatCompletionResponseChoice: + return ChatCompletionResponseChoice( + index=index, + message=ChatMessage(role="assistant", content=content), + logprobs=None, + finish_reason="stop", + stop_reason=None, + ) + + +def _audio_choice(index: int, data: str = "Zm9v") -> OmniChatCompletionResponseChoice: + return OmniChatCompletionResponseChoice( + index=index, + message=ChatMessage( + role="assistant", + audio=OpenAIChatCompletionAudio(id=f"audio-{index}", data=data, expires_at=0, transcript=""), + ), + audio_metadata=AudioChunkMetadata(format="wav", sample_rate_hz=24000, frame_count=17, channels=1), + logprobs=None, + finish_reason="stop", + stop_reason=None, + ) + + +@pytest.fixture +def serving_chat(): + return object.__new__(OmniOpenAIServingChat) + + +class TestMergeAudioChoices: + def test_text_and_audio_merge_into_single_choice(self, serving_chat): + merged = serving_chat._merge_audio_choices([_text_choice(0, "hello")], [_audio_choice(0)]) + + assert len(merged) == 1 + assert merged[0].index == 0 + assert merged[0].message.content == "hello" + assert merged[0].message.audio is not None + assert merged[0].message.audio.data == "Zm9v" + assert isinstance(merged[0], OmniChatCompletionResponseChoice) + assert merged[0].audio_metadata is not None + assert merged[0].audio_metadata.format == "wav" + + def test_no_duplicate_indexes_in_serialized_response(self, serving_chat): + """The exact #7376 symptom: two choices with index=0 in one response.""" + merged = serving_chat._merge_audio_choices([_text_choice(0, "hello")], [_audio_choice(0)]) + + response = OmniChatCompletionResponse( + id="chatcmpl-7376", + created=0, + model="test-model", + choices=merged, + usage=UsageInfo(prompt_tokens=0, completion_tokens=0, total_tokens=0), + ) + serialized = json.loads(response.model_dump_json()) + indexes = [c["index"] for c in serialized["choices"]] + assert indexes == [0] + assert serialized["choices"][0]["message"]["content"] == "hello" + assert serialized["choices"][0]["message"]["audio"]["data"] == "Zm9v" + assert serialized["choices"][0]["audio_metadata"]["format"] == "wav" + + def test_audio_only_request_keeps_standalone_choice(self, serving_chat): + merged = serving_chat._merge_audio_choices([], [_audio_choice(0)]) + + assert len(merged) == 1 + assert merged[0].message.audio is not None + assert merged[0].message.content is None + + def test_n_gt_1_merges_by_index(self, serving_chat): + text_choices = [_text_choice(0, "first"), _text_choice(1, "second")] + audio_choices = [_audio_choice(0, "AAAA"), _audio_choice(1, "BBBB")] + + merged = serving_chat._merge_audio_choices(text_choices, audio_choices) + + assert len(merged) == 2 + assert [c.index for c in merged] == [0, 1] + assert merged[0].message.content == "first" + assert merged[0].message.audio.data == "AAAA" + assert merged[1].message.content == "second" + assert merged[1].message.audio.data == "BBBB" + + def test_audio_without_matching_text_choice_is_appended(self, serving_chat): + # n=1 text but the audio stage reported a second index (defensive): + # the unmatched audio choice must not be dropped. + merged = serving_chat._merge_audio_choices([_text_choice(0, "only")], [_audio_choice(1)]) + + assert len(merged) == 2 + assert merged[0].message.audio is None + assert merged[1].message.audio is not None + + def test_text_choice_fields_survive_merge(self, serving_chat): + text = ChatCompletionResponseChoice( + index=0, + message=ChatMessage(role="assistant", content="kept"), + logprobs=None, + finish_reason="length", + stop_reason="max_tokens", + ) + + merged = serving_chat._merge_audio_choices([text], [_audio_choice(0)]) + + assert merged[0].finish_reason == "length" + assert merged[0].stop_reason == "max_tokens" + assert merged[0].message.content == "kept" diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index 6926253c5..d8aaa4e1d 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -2437,6 +2437,12 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): kv_transfer_params = None response_metrics: dict[str, Any] | None = None + # For text+audio requests the audio final output shares output + # indexes with the text final output. Collect audio choices and + # merge them into the matching text choice after the loop instead + # of appending duplicate-index choices (#7376). + pending_audio_choices: list[OmniChatCompletionResponseChoice] = [] + # Build requested modalities set for filtering requested_modalities = ( set(request.modalities) if hasattr(request, "modalities") and request.modalities else None @@ -2486,9 +2492,11 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): ) ] elif omni_outputs.final_output_type == "audio": - choices_data = self._create_audio_choice(omni_outputs, role, request, stream=False) - if isinstance(choices_data, ErrorResponse): - return choices_data + audio_data = self._create_audio_choice(omni_outputs, role, request, stream=False) + if isinstance(audio_data, ErrorResponse): + return audio_data + pending_audio_choices.extend(audio_data) + choices_data = [] elif omni_outputs.final_output_type == "image": choices_data = self._create_image_choice(omni_outputs, role, request, stream=False) else: @@ -2507,6 +2515,9 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): response_metrics.update(extra) choices.extend(choices_data) + if pending_audio_choices: + choices = self._merge_audio_choices(choices, pending_audio_choices) + response_metrics = self._filter_stage_metrics_detail(response_metrics, request) # Compute prompt_text for non-streaming response (upstream #42052) @@ -2819,6 +2830,37 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): return choices, usage, prompt_logprobs, prompt_token_ids, kv_transfer_params + def _merge_audio_choices( + self, + choices: list[ChatCompletionResponseChoice], + audio_choices: list[OmniChatCompletionResponseChoice], + ) -> list[ChatCompletionResponseChoice]: + """Fold audio outputs into the matching text choice (#7376). + + A non-streaming chat request with ``modalities=["text", "audio"]`` + produces two final outputs (text and audio) whose entries share the + same output index. Appending both as separate choices yields + duplicate ``index`` values, and clients that read ``choices[0]`` + never see the audio. Merge the audio object and its metadata into + the choice with the matching index; audio-only requests (no text + choice) keep the standalone audio choice. + """ + for audio_choice in audio_choices: + for i, existing in enumerate(choices): + if existing.index != audio_choice.index: + continue + merged = OmniChatCompletionResponseChoice( + **existing.model_dump(exclude={"message"}), + message=existing.message, + audio_metadata=audio_choice.audio_metadata, + ) + merged.message.audio = audio_choice.message.audio + choices[i] = merged + break + else: + choices.append(audio_choice) + return choices + def _create_audio_choice( self, omni_outputs: OmniRequestOutput, role: str, request: ChatCompletionRequest, stream: bool = False ) -> list[ChatCompletionResponseChoice] | list[ChatCompletionResponseStreamChoice] | ErrorResponse: