[Bugfix] Merge audio final output into the matching text choice (#7376) (#7879)

Signed-off-by: mershi <mershi@tencent.com>
Co-authored-by: mershi <mershi@tencent.com>

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"

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: