[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"