From e0c4f97f326f895c7830bdbdcc8c6861b8fcb10a Mon Sep 17 00:00:00 2001 From: bcsdhjew Date: Tue, 15 Sep 2026 00:05:55 +0800 Subject: [PATCH] [Bugfix][Frontend] Preserve duplex chat audio format and duration (#7316) Signed-off-by: Nolen Liang --- docs/serving/realtime_duplex_api.md | 43 +++ .../openai/test_duplex_protocol.py | 48 ++++ .../openai/test_duplex_session_attachment.py | 84 ++++++ .../openai_api/test_audio_format.py | 96 ++++++- .../openai_api/test_duplex_handler.py | 260 +++++++++++++++++- vllm_omni/entrypoints/duplex/chat_fallback.py | 57 +++- vllm_omni/entrypoints/duplex/protocol.py | 40 ++- .../entrypoints/duplex/session_attachment.py | 10 + .../entrypoints/duplex/session_runner.py | 19 +- .../entrypoints/openai/audio_utils_mixin.py | 18 +- .../entrypoints/openai/protocol/audio.py | 14 + .../openai/protocol/chat_completion.py | 23 +- vllm_omni/entrypoints/openai/serving_chat.py | 12 +- 13 files changed, 684 insertions(+), 40 deletions(-) diff --git a/docs/serving/realtime_duplex_api.md b/docs/serving/realtime_duplex_api.md index c855c4491..288c934e8 100644 --- a/docs/serving/realtime_duplex_api.md +++ b/docs/serving/realtime_duplex_api.md @@ -387,6 +387,49 @@ Semantic divergences hidden behind shared names: the Tier 3 surface on top of the same Tier 1 events, so both kinds of client can share one deployment and one event log. +### Turn-based Chat audio metadata + +For turn-based models served through Chat fallback, non-empty audio choices +produced by the Omni Chat audio encoder include `audio_metadata`, in both +streaming and complete Chat responses. Text-only choices do not require it. + +| Field | Meaning | +| --- | --- | +| `format` | Encoding of the returned audio, such as `pcm` or `wav`; `pcm16` requests normalize to `pcm`. | +| `sample_rate_hz` | Sample rate used by the encoder, in samples per second. | +| `frame_count` | Samples per channel in this chunk, or in the complete waveform for a non-streaming choice. Not a byte count. | +| `channels` | Number of audio channels. | + +These values describe the waveform after audio transformations. Compressed +formats may add codec padding that is not included in `frame_count`. +The encoder uses a model-reported sample rate when available; Qwen3-Omni +currently omits it and uses the existing 24 kHz default. This metadata does +not independently determine a model's sample rate. + +Chat fallback requests the session's output format explicitly. Its internal +`response.output_audio.delta` carries `format`, `sample_rate_hz`, `channels`, +and `audio_duration_ms`. The Realtime endpoint exposes output `format` and +`sample_rate_hz` on `response.audio.delta`, with duration in +`metadata.audio_duration_ms`; it does not currently forward `channels`. +The duration is cumulative for the response: +integer frames are summed before converting to milliseconds and rounding down. +Non-empty audio without valid metadata, or a sample-rate change within one +response, fails that fallback response with `response_error`. + +In this Chat fallback path, `response.done.response.metadata.playback` reports: + +- `generated_ms`: cumulative source-waveform duration. +- `sent_ms`: cumulative audio accepted into the server output path, including + queued or journaled audio. A journaled event remains accepted if the + subsequent send fails; this is not proof of socket delivery. +- `played_ms`: cumulative playback reported by the client. +- `committed_ms`: playback position used for history reconciliation. + +The native model path still records audio before sending; it does not yet +share Chat fallback's output-acceptance timing. Only client playback reports +establish what was played. The existing duration-proportional text truncation +remains an approximation, not word-level audio alignment. + ### Event catalogue Every event type the Realtime route accepts or emits, with a worked example in the next section: 21 diff --git a/tests/entrypoints/openai/test_duplex_protocol.py b/tests/entrypoints/openai/test_duplex_protocol.py index e5f3becbc..77c8f2cb0 100644 --- a/tests/entrypoints/openai/test_duplex_protocol.py +++ b/tests/entrypoints/openai/test_duplex_protocol.py @@ -462,6 +462,54 @@ def test_duplex_playback_ack_tracks_committed_cursor_separately(): assert session.playback.committed_ms == 2_000 +@pytest.mark.parametrize("sample_rate_hz,frame_count", [(22_050, 23), (24_000, 24)]) +def test_duplex_generated_audio_rounds_only_cumulative_frames(sample_rate_hz: int, frame_count: int): + session = DuplexSessionRegistry().create() + response_id = session.begin_response() + + durations = [ + session.record_generated_audio(response_id, frame_count=1, sample_rate_hz=sample_rate_hz) + for _ in range(frame_count) + ] + + assert durations[:-1] == [0] * (frame_count - 1) + assert durations[-1] == 1 + assert session.playback.generated_ms == 1 + assert session.playback.sent_ms == 0 + assert session.playback.played_ms == 0 + + +def test_duplex_generated_audio_rejects_rate_change_without_advancing(): + session = DuplexSessionRegistry().create() + response_id = session.begin_response() + session.record_generated_audio(response_id, frame_count=2205, sample_rate_hz=22_050) + + with pytest.raises(ValueError, match="sample rate changed"): + session.record_generated_audio(response_id, frame_count=2400, sample_rate_hz=24_000) + + assert session.playback.generated_ms == 100 + assert session.playback.sent_ms == 0 + + +def test_duplex_late_audio_accounting_cannot_change_new_response(): + session = DuplexSessionRegistry().create() + old_response = session.begin_response() + session.record_generated_audio(old_response, frame_count=2205, sample_rate_hz=22_050) + session.barge_in() + new_response = session.begin_response() + session.record_generated_audio(new_response, frame_count=2400, sample_rate_hz=48_000) + current_state = session.turn_state + + assert session.record_generated_audio(old_response, frame_count=2205, sample_rate_hz=22_050) is None + session.mark_audio_sent(100, response_id=old_response, text_chars=99) + + assert session.playback_for_response(old_response).sent_ms == 100 + assert session.playback.generated_ms == 50 + assert session.playback.sent_ms == 0 + assert session.assistant_audio_text_marks == () + assert session.turn_state == current_state + + def test_duplex_history_commit_uses_audio_text_alignment_marks(): registry = DuplexSessionRegistry() session = registry.create() diff --git a/tests/entrypoints/openai/test_duplex_session_attachment.py b/tests/entrypoints/openai/test_duplex_session_attachment.py index a73159064..13bd32761 100644 --- a/tests/entrypoints/openai/test_duplex_session_attachment.py +++ b/tests/entrypoints/openai/test_duplex_session_attachment.py @@ -122,6 +122,90 @@ def test_event_journal_overflow_is_explicit_and_does_not_evict_silently() -> Non journal.replay_after(0) +@pytest.mark.asyncio +@pytest.mark.parametrize("journal", [False, True]) +async def test_registry_failed_send_is_accepted_only_when_journaled(mocker, journal: bool) -> None: + registry = DuplexSessionAttachmentRegistry(replay_ttl_s=60.0, replay_max_bytes_per_session=4096) + accepted = mocker.Mock() + expected_calls = [mocker.call()] if journal else [] + + async def failing_send(payload): + assert accepted.call_args_list == expected_calls + raise RuntimeError("transport lost") + + await registry.create("sid-send-failure", incarnation=0, send=failing_send, close=mocker.AsyncMock()) + + with pytest.raises(RuntimeError, match="transport lost"): + await registry.send_event( + "sid-send-failure", + {"type": "response.audio.delta", "delta": "AAAA"}, + journal=journal, + on_accepted=accepted, + ) + + assert accepted.call_args_list == expected_calls + + +@pytest.mark.asyncio +@pytest.mark.parametrize("journal", [False, True]) +async def test_registry_detached_event_is_accepted_only_when_journaled(mocker, journal: bool) -> None: + registry = DuplexSessionAttachmentRegistry(replay_ttl_s=60.0, replay_max_bytes_per_session=4096) + send = mocker.AsyncMock() + accepted = mocker.Mock() + created = await registry.create("sid-detached-acceptance", incarnation=0, send=send, close=mocker.AsyncMock()) + await registry.detach("sid-detached-acceptance", attachment_generation=created.attachment_generation) + + entry = await registry.send_event( + "sid-detached-acceptance", + {"type": "response.audio.delta", "delta": "AAAA"}, + journal=journal, + on_accepted=accepted, + ) + + send.assert_not_awaited() + assert (entry is not None) is journal + assert accepted.call_count == int(journal) + + +@pytest.mark.asyncio +async def test_registry_unjournaled_event_is_accepted_after_successful_send(mocker) -> None: + registry = DuplexSessionAttachmentRegistry(replay_ttl_s=60.0, replay_max_bytes_per_session=4096) + accepted = mocker.Mock() + + async def send(payload): + accepted.assert_not_called() + + await registry.create("sid-live-acceptance", incarnation=0, send=send, close=mocker.AsyncMock()) + + entry = await registry.send_event( + "sid-live-acceptance", + {"type": "response.audio.delta", "delta": "AAAA"}, + journal=False, + on_accepted=accepted, + ) + + assert entry is None + accepted.assert_called_once_with() + + +@pytest.mark.asyncio +async def test_registry_overflow_does_not_accept_or_send_event(mocker) -> None: + registry = DuplexSessionAttachmentRegistry(replay_ttl_s=60.0, replay_max_bytes_per_session=100) + accepted = mocker.Mock() + send = mocker.AsyncMock() + await registry.create("sid-overflow-acceptance", incarnation=0, send=send, close=mocker.AsyncMock()) + + with pytest.raises(DuplexJournalOverflowError, match="byte limit"): + await registry.send_event( + "sid-overflow-acceptance", + {"type": "response.audio.delta", "delta": "A" * 200}, + on_accepted=accepted, + ) + + accepted.assert_not_called() + send.assert_not_awaited() + + @pytest.mark.asyncio async def test_registry_resume_rotates_token_replays_and_atomically_replaces_attachment() -> None: clock = _Clock() diff --git a/tests/entrypoints/openai_api/test_audio_format.py b/tests/entrypoints/openai_api/test_audio_format.py index 9c7eb9d89..c7fc4fbef 100644 --- a/tests/entrypoints/openai_api/test_audio_format.py +++ b/tests/entrypoints/openai_api/test_audio_format.py @@ -12,6 +12,8 @@ Covers: from __future__ import annotations +import base64 +import json from io import BytesIO import numpy as np @@ -27,6 +29,8 @@ from vllm_omni.entrypoints.openai.protocol.audio import ( SUPPORTED_CHAT_AUDIO_FORMATS, CreateAudio, ) +from vllm_omni.outputs import OmniRequestOutput +from vllm_omni.outputs.mm_outputs import MultimodalCompletionOutput, MultimodalPayload pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -65,6 +69,13 @@ class TestCreateAudio: ) response = mixin.create_audio(audio_obj) assert len(response.audio_data) > 0 + assert response.audio_metadata is not None + assert response.audio_metadata.model_dump() == { + "format": fmt, + "sample_rate_hz": 24000, + "frame_count": 24000, + "channels": 1, + } def test_wav_magic_bytes(self, mixin, audio_tensor): audio_obj = CreateAudio( @@ -113,10 +124,9 @@ class TestCreateAudio: response = mixin.create_audio(audio_obj) assert response.audio_data[:4] == b"RIFF" assert response.media_type == "audio/wav" + assert response.audio_metadata.format == "wav" def test_base64_encoding(self, mixin, audio_tensor): - import base64 - audio_obj = CreateAudio( audio_tensor=audio_tensor, sample_rate=24000, @@ -128,21 +138,33 @@ class TestCreateAudio: decoded = base64.b64decode(response.audio_data) assert decoded[:4] == b"RIFF" - def test_resamples_wav_to_requested_output_rate(self, mixin, audio_tensor): + @pytest.mark.parametrize( + "channels,speed,frame_count", + [(1, 1.0, 8000), (2, 2.0, 4000)], + ids=["mono-resampled", "stereo-resampled-faster"], + ) + def test_metadata_matches_transformed_waveform(self, mixin, audio_tensor, channels, speed, frame_count): response = mixin.create_audio( CreateAudio( - audio_tensor=audio_tensor, + audio_tensor=audio_tensor if channels == 1 else np.stack((audio_tensor, audio_tensor)), sample_rate=24000, output_sample_rate=8000, response_format="wav", - speed=1.0, + speed=speed, base64_encode=False, ) ) with soundfile.SoundFile(BytesIO(response.audio_data)) as audio_file: assert audio_file.samplerate == 8000 - assert audio_file.frames == 8000 + assert audio_file.frames == frame_count + assert audio_file.channels == channels + assert response.audio_metadata.model_dump() == { + "format": "wav", + "sample_rate_hz": audio_file.samplerate, + "frame_count": audio_file.frames, + "channels": audio_file.channels, + } class TestStreamingAudioResampler: @@ -262,3 +284,65 @@ class TestResolveAudioFormat: request = self._make_request({"format": fmt, "voice": "alloy"}) result = serving_chat._resolve_audio_format(request) assert not isinstance(result, ErrorResponse), f"Format {fmt} should be accepted" + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("fmt", ["pcm", "wav"]) +def test_chat_audio_metadata_survives_response_serialization(stream, fmt): + from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest + from vllm.entrypoints.openai.engine.protocol import ErrorResponse, UsageInfo + + from vllm_omni.entrypoints.openai.protocol.chat_completion import ( + OmniChatCompletionResponse, + OmniChatCompletionStreamResponse, + ) + from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat + + serving_chat = object.__new__(OmniOpenAIServingChat) + output = MultimodalCompletionOutput( + index=0, + text="", + token_ids=[], + cumulative_logprob=None, + logprobs=None, + multimodal_output=MultimodalPayload.from_dict( + {"audio": [torch.zeros(17), torch.zeros(23)], "sr": torch.tensor(22050)} + ), + finish_reason="stop", + stop_reason=None, + ) + request = ChatCompletionRequest( + model="test-model", + messages=[{"role": "user", "content": "hello"}], + audio={"format": fmt, "voice": "alloy"}, + ) + choices = serving_chat._create_audio_choice( + OmniRequestOutput(outputs=[output], final_output_type="audio"), "assistant", request, stream=stream + ) + assert not isinstance(choices, ErrorResponse) + response_type = OmniChatCompletionStreamResponse if stream else OmniChatCompletionResponse + response = response_type( + id="chatcmpl-audio-metadata", + created=0, + model="test-model", + choices=choices, + usage=UsageInfo(prompt_tokens=0, completion_tokens=0, total_tokens=0), + ) + + serialized = json.loads(response.model_dump_json(exclude_unset=True)) + choice = serialized["choices"][0] + frame_count = 23 if stream else 40 + assert choice["audio_metadata"] == { + "format": fmt, + "sample_rate_hz": 22050, + "frame_count": frame_count, + "channels": 1, + } + audio = choice["delta"]["content"] if stream else choice["message"]["audio"]["data"] + raw = base64.b64decode(audio) + if fmt == "wav": + with soundfile.SoundFile(BytesIO(raw)) as audio_file: + assert audio_file.frames == frame_count + assert audio_file.samplerate == 22050 + else: + assert len(raw) == frame_count * 2 diff --git a/tests/entrypoints/openai_api/test_duplex_handler.py b/tests/entrypoints/openai_api/test_duplex_handler.py index 206a38917..b4a943873 100644 --- a/tests/entrypoints/openai_api/test_duplex_handler.py +++ b/tests/entrypoints/openai_api/test_duplex_handler.py @@ -9,6 +9,8 @@ import io import json import struct import wave +from collections.abc import Callable +from contextlib import nullcontext from types import SimpleNamespace from typing import Any @@ -2182,7 +2184,8 @@ async def test_minicpmo_native_session_update_rejects_native_duplex_mode_flip(): assert "session.updated" not in ws.sent_types() -def test_native_realtime_protocol_audio_delta_preserves_sample_rate_hz(): +@pytest.mark.parametrize("sample_rate_hz", [22_050, 24_000]) +def test_native_realtime_protocol_audio_delta_preserves_sample_rate_hz(sample_rate_hz: int): ws = TimedWebSocket() protocol = NativeRealtimeSessionProtocol(ws) # type: ignore[arg-type] @@ -2192,7 +2195,8 @@ def test_native_realtime_protocol_audio_delta_preserves_sample_rate_hz(): "response_id": "resp-a", "audio": "AAAA", "format": "pcm", - "sample_rate_hz": 24000, + "sample_rate_hz": sample_rate_hz, + "audio_duration_ms": 135, } ) @@ -2201,7 +2205,8 @@ def test_native_realtime_protocol_audio_delta_preserves_sample_rate_hz(): ] assert {payload["type"] for payload in audio_events} == {"response.output_audio.delta"} assert {payload["format"] for payload in audio_events} == {"pcm16"} - assert {payload["sample_rate_hz"] for payload in audio_events} == {24000} + assert {payload["sample_rate_hz"] for payload in audio_events} == {sample_rate_hz} + assert audio_events[0]["metadata"]["audio_duration_ms"] == 135 def test_native_realtime_protocol_ignores_removed_legacy_event_switches(): @@ -5023,8 +5028,10 @@ async def test_duplex_chat_audio_stream_uses_output_audio_delta_event(): response_id = session.begin_response() sent: list[dict[str, Any]] = [] - async def send_json(data: dict[str, Any]) -> None: + async def send_json(data: dict[str, Any], *, on_accepted: Callable[[], None] | None = None) -> None: sent.append(data) + if on_accepted is not None: + on_accepted() await handler._emit_chat_payload( session, @@ -5034,7 +5041,13 @@ async def test_duplex_chat_audio_stream_uses_output_audio_delta_event(): { "delta": { "content": "AAAA", - } + }, + "audio_metadata": { + "format": "wav", + "sample_rate_hz": 22_050, + "frame_count": 2205, + "channels": 1, + }, } ], }, @@ -5051,8 +5064,222 @@ async def test_duplex_chat_audio_stream_uses_output_audio_delta_event(): "epoch": 0, "audio": "AAAA", "format": "wav", + "sample_rate_hz": 22_050, + "channels": 1, + "audio_duration_ms": 100, } ] + assert session.playback.generated_ms == 100 + assert session.playback.sent_ms == 100 + assert session.playback.played_ms == 0 + + +@pytest.mark.parametrize("response_format", ["pcm", "wav"]) +def test_duplex_chat_request_uses_response_audio_format(response_format: str): + handler = OmniDuplexSessionHandler( + chat_service=TurnBasedFakeChatService(FakeEngineClient()), + config_timeout_s=0.1, + idle_timeout_s=1, + ) + extra_audio = {"format": "flac", "voice": "alloy"} + session = DuplexSession( + session_id="sid-chat-audio-format", + config=DuplexSessionConfig(response_format=response_format, extra_body={"audio": extra_audio}), + ) + session.begin_response() + + request = handler._build_chat_request(session, "format-request") + + assert request.audio == {"format": response_format, "voice": "alloy"} + assert extra_audio == {"format": "flac", "voice": "alloy"} + + +def _chat_audio_completion_payload(*, frame_count: int = 2205, sample_rate_hz: int = 22_050) -> dict[str, object]: + return { + "modality": "audio", + "choices": [ + { + "delta": {"content": "AAAA"}, + "audio_metadata": { + "format": "pcm", + "sample_rate_hz": sample_rate_hz, + "frame_count": frame_count, + "channels": 1, + }, + } + ], + } + + +@pytest.mark.asyncio +async def test_duplex_chat_full_audio_message_preserves_actual_format_and_duration(): + handler = OmniDuplexSessionHandler( + chat_service=TurnBasedFakeChatService(FakeEngineClient()), + config_timeout_s=0.1, + idle_timeout_s=1, + ) + session = DuplexSession(session_id="sid-chat-full-audio", config=DuplexSessionConfig(response_format="wav")) + response_id = session.begin_response() + sent: list[dict[str, object]] = [] + + async def send_json(data: dict[str, object], *, on_accepted: Callable[[], None] | None = None) -> None: + sent.append(data) + if on_accepted is not None: + on_accepted() + + await handler._emit_chat_payload( + session, + { + "choices": [ + { + "message": {"role": "assistant", "content": "hello", "audio": {"data": "AAAA"}}, + "audio_metadata": { + "format": "pcm", + "sample_rate_hz": 22_050, + "frame_count": 2205, + "channels": 1, + }, + } + ], + }, + session.epoch, + response_id, + send_json, + ) + + assert [event["type"] for event in sent] == ["response.text.delta", "response.output_audio.delta"] + assert sent[1]["audio"] == "AAAA" + assert sent[1]["format"] == "pcm" + assert sent[1]["sample_rate_hz"] == 22_050 + assert sent[1]["audio_duration_ms"] == 100 + assert session.assistant_text_buffer == ("hello",) + assert session.playback.sent_ms == 100 + + +@pytest.mark.asyncio +async def test_duplex_chat_audio_does_not_round_each_fragment(): + handler = OmniDuplexSessionHandler( + chat_service=TurnBasedFakeChatService(FakeEngineClient()), + config_timeout_s=0.1, + idle_timeout_s=1, + ) + session = DuplexSession(session_id="sid-chat-sub-ms", config=DuplexSessionConfig()) + response_id = session.begin_response() + sent: list[dict[str, object]] = [] + + async def send_json(data: dict[str, object], *, on_accepted: Callable[[], None] | None = None) -> None: + sent.append(data) + if on_accepted is not None: + on_accepted() + + for _ in range(23): + await handler._emit_chat_payload( + session, + _chat_audio_completion_payload(frame_count=1), + session.epoch, + response_id, + send_json, + ) + + assert [event["audio_duration_ms"] for event in sent] == [0] * 22 + [1] + assert session.playback.generated_ms == 1 + assert session.playback.sent_ms == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("send_error", [None, RuntimeError("output unavailable")], ids=["rejected", "failed"]) +async def test_duplex_chat_unaccepted_audio_is_generated_but_not_sent(send_error: RuntimeError | None): + handler = OmniDuplexSessionHandler( + chat_service=TurnBasedFakeChatService(FakeEngineClient()), + config_timeout_s=0.1, + idle_timeout_s=1, + ) + session = DuplexSession(session_id="sid-chat-unaccepted-audio", config=DuplexSessionConfig()) + response_id = session.begin_response() + + async def send_json(data: dict[str, object], *, on_accepted: Callable[[], None] | None = None) -> None: + if send_error is not None: + raise send_error + + with pytest.raises(RuntimeError, match="output unavailable") if send_error is not None else nullcontext(): + await handler._emit_chat_payload( + session, _chat_audio_completion_payload(), session.epoch, response_id, send_json + ) + + assert session.playback.generated_ms == 100 + assert session.playback.sent_ms == 0 + assert session.playback.played_ms == 0 + + +@pytest.mark.asyncio +async def test_duplex_chat_late_send_completion_only_updates_old_response(): + handler = OmniDuplexSessionHandler( + chat_service=TurnBasedFakeChatService(FakeEngineClient()), + config_timeout_s=0.1, + idle_timeout_s=1, + ) + session = DuplexSession(session_id="sid-chat-late-audio", config=DuplexSessionConfig()) + old_response = session.begin_response() + send_started = asyncio.Event() + finish_send = asyncio.Event() + + async def delayed_send(data: dict[str, object], *, on_accepted: Callable[[], None] | None = None) -> None: + send_started.set() + await finish_send.wait() + assert on_accepted is not None + on_accepted() + + emit_task = asyncio.create_task( + handler._emit_chat_payload(session, _chat_audio_completion_payload(), session.epoch, old_response, delayed_send) + ) + try: + await asyncio.wait_for(send_started.wait(), timeout=1) + session.barge_in() + new_response = session.begin_response() + current_state = session.turn_state + finish_send.set() + await asyncio.wait_for(emit_task, timeout=1) + + assert session.active_response_id == new_response + assert session.playback_for_response(old_response).sent_ms == 100 + assert session.playback.generated_ms == 0 + assert session.playback.sent_ms == 0 + assert session.turn_state == current_state + finally: + emit_task.cancel() + await asyncio.gather(emit_task, return_exceptions=True) + + +@pytest.mark.asyncio +async def test_duplex_chat_missing_audio_metadata_reports_error(): + class MissingMetadataChatService(TurnBasedFakeChatService): + async def create_chat_completion(self, request, raw_request=None): + async def generate(): + payload = {"modality": "audio", "choices": [{"delta": {"content": "AAAA"}}]} + yield f"data: {json.dumps(payload)}\n\n" + + return generate() + + handler = OmniDuplexSessionHandler( + chat_service=MissingMetadataChatService(FakeEngineClient()), + config_timeout_s=0.1, + idle_timeout_s=1, + ) + session = DuplexSession(session_id="sid-chat-missing-metadata", config=DuplexSessionConfig()) + sent: list[dict[str, object]] = [] + + async def send_json(data: dict[str, object], *, on_accepted: Callable[[], None] | None = None) -> None: + sent.append(data) + if on_accepted is not None: + on_accepted() + + await handler._run_response(session, send_json) + + assert [event["type"] for event in sent] == ["response.created", "error"] + assert sent[-1]["code"] == "response_error" + assert "AudioChunkMetadata" in sent[-1]["error"] + assert session.playback.generated_ms == 0 + assert session.playback.sent_ms == 0 @pytest.mark.asyncio @@ -5071,7 +5298,17 @@ async def test_duplex_chat_stage_metrics_use_latest_streaming_snapshot(): for index, snapshot in enumerate(snapshots): payload = { "modality": "audio", - "choices": [{"delta": {"content": f"audio-{index}"}}], + "choices": [ + { + "delta": {"content": f"audio-{index}"}, + "audio_metadata": { + "format": "pcm", + "sample_rate_hz": 24_000, + "frame_count": 2400, + "channels": 1, + }, + } + ], "metrics": {"stage_metrics": {"0": snapshot}}, } yield f"data: {json.dumps(payload)}\n\n" @@ -5089,13 +5326,22 @@ async def test_duplex_chat_stage_metrics_use_latest_streaming_snapshot(): ) sent: list[dict[str, Any]] = [] - async def send_json(data: dict[str, Any]) -> None: + async def send_json(data: dict[str, Any], *, on_accepted: Callable[[], None] | None = None) -> None: sent.append(data) + if on_accepted is not None: + on_accepted() await handler._run_response(session, send_json) audio_deltas = [event for event in sent if event.get("type") == "response.output_audio.delta"] done = next(event for event in sent if event.get("type") == "response.done") + assert [event["audio_duration_ms"] for event in audio_deltas] == [100, 200, 300] + assert done["playback"] == { + "generated_ms": 300, + "sent_ms": 300, + "played_ms": 0, + "committed_ms": 0, + } assert all("vllm_omni" not in event for event in audio_deltas) assert done["vllm_omni"]["stage_metrics"]["0"] == { "num_tokens_out": 3, diff --git a/vllm_omni/entrypoints/duplex/chat_fallback.py b/vllm_omni/entrypoints/duplex/chat_fallback.py index f41d7205b..bd682c4f3 100644 --- a/vllm_omni/entrypoints/duplex/chat_fallback.py +++ b/vllm_omni/entrypoints/duplex/chat_fallback.py @@ -6,6 +6,7 @@ from __future__ import annotations import asyncio import json from collections.abc import AsyncGenerator +from functools import partial from typing import Any from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest @@ -115,6 +116,13 @@ class ChatFallbackProjectorMixin: ): model_extra.pop(protocol_key, None) kwargs.update(model_extra) + if "audio" in response_config.modalities: + audio_options = kwargs.get("audio") + audio_options = dict(audio_options) if isinstance(audio_options, dict) else {} + # Chat defaults to WAV. Request the session's internal output format + # explicitly rather than relabeling the returned bytes as PCM. + audio_options["format"] = response_config.response_format + kwargs["audio"] = audio_options if isinstance(tools, list): kwargs["tools"] = tools if isinstance(tool_choice, str | dict): @@ -190,6 +198,8 @@ class ChatFallbackProjectorMixin: response_id: str, send_json, ) -> None: + if session.epoch != epoch or session.active_response_id != response_id: + return metrics = payload.get("metrics") stage_metrics = metrics.get("stage_metrics") if isinstance(metrics, dict) else None if isinstance(stage_metrics, dict): @@ -233,19 +243,7 @@ class ChatFallbackProjectorMixin: content = message.get("content") if isinstance(content, str) and content: - if modality == "audio": - session.mark_audio_sent() - await send_json( - { - "type": "response.output_audio.delta", - "session_id": session.session_id, - "response_id": response_id, - "epoch": epoch, - "audio": content, - "format": session.response_config.response_format, - } - ) - else: + if modality != "audio": session.append_assistant_text(content) await send_json( { @@ -257,6 +255,39 @@ class ChatFallbackProjectorMixin: } ) + audio_content = content if modality == "audio" else None + if isinstance(message, dict): + message_audio = message.get("audio") + if isinstance(message_audio, dict): + audio_content = message_audio.get("data") + if isinstance(audio_content, str) and audio_content: + # The OpenAI package imports duplex serving during initialization. + # Defer this import until chat fallback has finished loading. + from vllm_omni.entrypoints.openai.protocol.audio import AudioChunkMetadata + + audio_metadata = AudioChunkMetadata.model_validate(choice.get("audio_metadata")) + duration_ms = session.record_generated_audio( + response_id, + frame_count=audio_metadata.frame_count, + sample_rate_hz=audio_metadata.sample_rate_hz, + ) + if duration_ms is None: + return + await send_json( + { + "type": "response.output_audio.delta", + "session_id": session.session_id, + "response_id": response_id, + "epoch": epoch, + "audio": audio_content, + "format": audio_metadata.format, + "sample_rate_hz": audio_metadata.sample_rate_hz, + "channels": audio_metadata.channels, + "audio_duration_ms": duration_ms, + }, + on_accepted=partial(session.mark_audio_sent, duration_ms, response_id=response_id), + ) + finish_reason = choice.get("finish_reason") if finish_reason is not None and modality != "audio": await send_json( diff --git a/vllm_omni/entrypoints/duplex/protocol.py b/vllm_omni/entrypoints/duplex/protocol.py index ff91e0db5..d28a7b009 100644 --- a/vllm_omni/entrypoints/duplex/protocol.py +++ b/vllm_omni/entrypoints/duplex/protocol.py @@ -178,6 +178,8 @@ class DuplexPlaybackCursor: sent_ms: int = 0 played_ms: int = 0 committed_ms: int = 0 + generated_frames: int = 0 + sample_rate_hz: int | None = None def acknowledge(self, played_ms: int, committed_ms: int | None = None) -> None: self.played_ms = max(self.played_ms, max(0, int(played_ms))) @@ -1103,25 +1105,54 @@ class DuplexSession: if text: self._response.assistant_text_buffer.append(text) + def record_generated_audio(self, response_id: str, *, frame_count: int, sample_rate_hz: int) -> int | None: + """Accumulate source frames, rounding only the cumulative duration. + + A late chunk cannot change a different response's playback cursor. + This records encoded audio, not transport delivery or client playback. + """ + if response_id != self.active_response_id: + return None + if type(frame_count) is not int or frame_count < 0: + raise ValueError("Audio frame_count must be a non-negative integer") + if type(sample_rate_hz) is not int or sample_rate_hz <= 0: + raise ValueError("Audio sample_rate_hz must be a positive integer") + playback = self._playback.current + if playback.sample_rate_hz is not None and playback.sample_rate_hz != sample_rate_hz: + raise ValueError("Audio sample rate changed within a response") + playback.sample_rate_hz = sample_rate_hz + playback.generated_frames += frame_count + playback.generated_ms = playback.generated_frames * 1000 // sample_rate_hz + return playback.generated_ms + def mark_audio_sent( self, duration_ms: int | None = None, *, + response_id: str | None = None, text_chars: int | None = None, audio_text_marks: list[dict[str, object]] | None = None, ) -> None: - playback = self._playback.current + """Record cumulative audio accepted by the output path. + + ``sent_ms`` includes queued/journaled output; only client acknowledgements + establish ``played_ms``. It is not confirmation of socket delivery. + """ + playback = self._playback.current if response_id is None else self._playback.by_response.get(response_id) + if playback is None: + return + is_current = response_id is None or response_id == self.active_response_id if duration_ms is not None: playback.generated_ms = max(playback.generated_ms, duration_ms) playback.sent_ms = max(playback.sent_ms, duration_ms) - if text_chars is not None and text_chars >= 0: + if is_current and text_chars is not None and text_chars >= 0: self._response.assistant_audio_text_marks.append( DuplexAssistantAudioTextMark( text_chars=int(text_chars), audio_end_ms=max(0, int(duration_ms)), ) ) - if audio_text_marks: + if is_current and audio_text_marks: for raw_mark in audio_text_marks: if not isinstance(raw_mark, dict): continue @@ -1135,7 +1166,8 @@ class DuplexSession: audio_end_ms=max(0, int(raw_audio_end_ms)), ) ) - self.turn_state = DuplexTurnState.ASSISTANT_PLAYING + if is_current: + self.turn_state = DuplexTurnState.ASSISTANT_PLAYING def _playback_cursor_for_response(self, response_id: str | None = None) -> DuplexPlaybackCursor: if response_id is None: diff --git a/vllm_omni/entrypoints/duplex/session_attachment.py b/vllm_omni/entrypoints/duplex/session_attachment.py index d68fb268b..242169a6f 100644 --- a/vllm_omni/entrypoints/duplex/session_attachment.py +++ b/vllm_omni/entrypoints/duplex/session_attachment.py @@ -275,12 +275,18 @@ class DuplexSessionAttachmentRegistry: payload: Mapping[str, object], *, journal: bool = True, + on_accepted: Callable[[], None] | None = None, ) -> JournalEntry | None: """Sequence and dispatch one event to the current attachment. The per-session lock keeps wire order equal to journal order without serializing unrelated sessions. A detached session still records replayable events, but has no transport side effect. + + ``on_accepted`` runs synchronously once the event is journaled, or + after a successful transport send when journaling is disabled. A later + send failure cannot undo acceptance into the journal. Detached, + non-journaled events do not invoke it. The callback must not raise. """ async with self._lock: state = self._require(session_id) @@ -291,8 +297,12 @@ class DuplexSessionAttachmentRegistry: entry = state.journal.record(payload) if journal else None attachment = state.attachment wire_payload = dict(entry.payload) if entry is not None else dict(payload) + if entry is not None and on_accepted is not None: + on_accepted() if attachment is not None: await attachment.send(wire_payload) + if entry is None and on_accepted is not None: + on_accepted() return entry async def acknowledge(self, session_id: str, sequence: int) -> int: diff --git a/vllm_omni/entrypoints/duplex/session_runner.py b/vllm_omni/entrypoints/duplex/session_runner.py index ed18d8591..8ee5060b0 100644 --- a/vllm_omni/entrypoints/duplex/session_runner.py +++ b/vllm_omni/entrypoints/duplex/session_runner.py @@ -7,6 +7,7 @@ import asyncio import binascii import json import uuid +from collections.abc import Callable from contextlib import suppress from copy import deepcopy @@ -99,7 +100,9 @@ class DuplexSessionRunnerMixin: if callable(close): await close(code=1000, reason=reason) - async def send_attachment_payload(payload: dict[str, object], *, journal: bool) -> None: + async def send_attachment_payload( + payload: dict[str, object], *, journal: bool, on_accepted: Callable[[], None] | None = None + ) -> None: assert session is not None session_id = session.session_id should_journal = journal and session_id not in self._resync_required_sessions @@ -108,6 +111,7 @@ class DuplexSessionRunnerMixin: session_id, payload, journal=should_journal, + on_accepted=on_accepted, ) except DuplexJournalOverflowError: first_overflow = session_id not in self._resync_required_sessions @@ -129,6 +133,7 @@ class DuplexSessionRunnerMixin: session_id, payload, journal=False, + on_accepted=on_accepted, ) if realtime_protocol is not None: @@ -155,23 +160,29 @@ class DuplexSessionRunnerMixin: event_emit_lock = asyncio.Lock() - async def send_outbound(payload: dict[str, object]) -> None: + async def send_outbound(payload: dict[str, object], *, on_accepted: Callable[[], None] | None = None) -> None: if not attachment_ready or session is None: await actor.send_json(dict(payload)) + if on_accepted is not None: + on_accepted() return if realtime_protocol is None: await send_attachment_payload( payload, journal=payload.get("type") not in {"session.created", "session.resumed"}, + on_accepted=on_accepted, ) return for projected in realtime_protocol.encode_outbound_event(payload): await send_attachment_payload( projected, journal=projected.get("type") not in {"session.created", "session.resumed"}, + # Item/part announcements may precede the audio. Only the + # actual audio event establishes acceptance of this chunk. + on_accepted=on_accepted if projected.get("type") == "response.audio.delta" else None, ) - async def emit_event(payload: dict[str, object]) -> None: + async def emit_event(payload: dict[str, object], *, on_accepted: Callable[[], None] | None = None) -> None: deferred_precreate_response = False async with event_emit_lock: accepted, deferred_overlap_payload = await self._apply_outbound_session_event( @@ -183,7 +194,7 @@ class DuplexSessionRunnerMixin: ) if not accepted: return - await send_outbound(payload) + await send_outbound(payload, on_accepted=on_accepted) if deferred_overlap_payload is not None: deferred_precreate_response = native.deferred_precreate_response native.deferred_precreate_response = False diff --git a/vllm_omni/entrypoints/openai/audio_utils_mixin.py b/vllm_omni/entrypoints/openai/audio_utils_mixin.py index 4382afc57..a13318ad4 100644 --- a/vllm_omni/entrypoints/openai/audio_utils_mixin.py +++ b/vllm_omni/entrypoints/openai/audio_utils_mixin.py @@ -29,7 +29,12 @@ import torch import torchaudio from vllm.logger import init_logger -from vllm_omni.entrypoints.openai.protocol.audio import DEFAULT_AUDIO_FORMAT, AudioResponse, CreateAudio +from vllm_omni.entrypoints.openai.protocol.audio import ( + DEFAULT_AUDIO_FORMAT, + AudioChunkMetadata, + AudioResponse, + CreateAudio, +) try: import soundfile @@ -244,7 +249,16 @@ class AudioMixin: audio_data = base64.b64encode(audio_data).decode("utf-8") - return AudioResponse(audio_data=audio_data, media_type=media_type) + return AudioResponse( + audio_data=audio_data, + media_type=media_type, + audio_metadata=AudioChunkMetadata( + format=response_format, + sample_rate_hz=int(sample_rate), + frame_count=int(audio_tensor.shape[0]), + channels=1 if audio_tensor.ndim == 1 else int(audio_tensor.shape[1]), + ), + ) @staticmethod def _resample_audio(audio_tensor: np.ndarray, source_rate: int, target_rate: int) -> np.ndarray: diff --git a/vllm_omni/entrypoints/openai/protocol/audio.py b/vllm_omni/entrypoints/openai/protocol/audio.py index 459f9d4ae..8a1b48304 100644 --- a/vllm_omni/entrypoints/openai/protocol/audio.py +++ b/vllm_omni/entrypoints/openai/protocol/audio.py @@ -421,9 +421,23 @@ class CreateAudio(BaseModel): arbitrary_types_allowed = True +class AudioChunkMetadata(BaseModel): + """Waveform dimensions after transforms, before encoding. + + Frames count samples per channel, not interleaved scalar samples or bytes. + For compressed formats this excludes any padding introduced by the codec. + """ + + format: str = Field(min_length=1, strict=True) + sample_rate_hz: int = Field(gt=0, strict=True) + frame_count: int = Field(ge=0, strict=True) + channels: int = Field(gt=0, strict=True) + + class AudioResponse(BaseModel): audio_data: bytes | str media_type: str + audio_metadata: AudioChunkMetadata | None = None # --- Batch Speech Models --- diff --git a/vllm_omni/entrypoints/openai/protocol/chat_completion.py b/vllm_omni/entrypoints/openai/protocol/chat_completion.py index 462e53af9..eaa15eea9 100644 --- a/vllm_omni/entrypoints/openai/protocol/chat_completion.py +++ b/vllm_omni/entrypoints/openai/protocol/chat_completion.py @@ -1,12 +1,33 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from typing import Any -from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionResponse, ChatCompletionStreamResponse +from pydantic import SerializeAsAny +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionResponse, + ChatCompletionResponseChoice, + ChatCompletionResponseStreamChoice, + ChatCompletionStreamResponse, +) + +from vllm_omni.entrypoints.openai.protocol.audio import AudioChunkMetadata + + +class OmniChatCompletionResponseStreamChoice(ChatCompletionResponseStreamChoice): + audio_metadata: AudioChunkMetadata | None = None + + +class OmniChatCompletionResponseChoice(ChatCompletionResponseChoice): + audio_metadata: AudioChunkMetadata | None = None class OmniChatCompletionStreamResponse(ChatCompletionStreamResponse): + choices: list[SerializeAsAny[ChatCompletionResponseStreamChoice]] modality: str | None = "text" metrics: dict[str, Any] | None = None class OmniChatCompletionResponse(ChatCompletionResponse): + choices: list[SerializeAsAny[ChatCompletionResponseChoice]] metrics: dict[str, Any] | None = None diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index 20bddaba6..484d5d170 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -34,7 +34,11 @@ from vllm_omni.entrypoints.openai.diffusion_request_utils import ( apply_normalized_diffusion_request_extra_args, normalize_diffusion_request_args, ) -from vllm_omni.entrypoints.openai.protocol.chat_completion import OmniChatCompletionResponse +from vllm_omni.entrypoints.openai.protocol.chat_completion import ( + OmniChatCompletionResponse, + OmniChatCompletionResponseChoice, + OmniChatCompletionResponseStreamChoice, +) from vllm_omni.entrypoints.utils import coerce_param_message_types from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniTextPrompt from vllm_omni.metrics import definitions as _metric_defs @@ -2860,18 +2864,20 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): for output in final_res.outputs: if stream: - choice_data = ChatCompletionResponseStreamChoice( + choice_data = OmniChatCompletionResponseStreamChoice( index=output.index, delta=DeltaMessage(role=role, content=audio_base64), + audio_metadata=audio_response.audio_metadata, logprobs=None, finish_reason=output.finish_reason, stop_reason=output.stop_reason, token_ids=(as_list(output.token_ids) if request.return_token_ids else None), ) else: - choice_data = ChatCompletionResponseChoice( + choice_data = OmniChatCompletionResponseChoice( index=output.index, message=ChatMessage(role=role, audio=audio_obj), + audio_metadata=audio_response.audio_metadata, logprobs=None, finish_reason="stop", stop_reason=output.stop_reason,