From 23bb4395bb334d21c482d0d2ba04c46738d5c214 Mon Sep 17 00:00:00 2001 From: iancarrasco-b10 Date: Tue, 28 Jul 2026 16:46:19 -0700 Subject: [PATCH] [Feature] Flush the speech WebSocket on input.done without closing it (#5517) Signed-off-by: iancarrasco-b10 --- docs/serving/speech_api.md | 32 +- .../examples/online_serving/text_to_speech.md | 9 +- .../online_serving/text_to_speech/README.md | 9 +- .../qwen3_tts/streaming_speech_client.py | 374 ++++++++++-------- .../openai_api/test_qwen3_tts_websocket.py | 138 +++++-- .../openai_api/test_serving_speech_stream.py | 218 +++++++++- vllm_omni/entrypoints/openai/api_server.py | 5 +- .../openai/serving_speech_stream.py | 205 ++++++---- 8 files changed, 688 insertions(+), 302 deletions(-) diff --git a/docs/serving/speech_api.md b/docs/serving/speech_api.md index 2bf630476..0deea7f8c 100644 --- a/docs/serving/speech_api.md +++ b/docs/serving/speech_api.md @@ -276,20 +276,42 @@ Client -> Server: | Message | Description | |---------|-------------| -| `{"type": "session.config", ...}` | Session configuration (sent once, first message) | +| `{"type": "session.config", ...}` | Session configuration (first message; may be resent between utterances to change it) | | `{"type": "input.text", "text": "..."}` | Text chunk | -| `{"type": "input.done"}` | End of input, flushes remaining buffer | +| `{"type": "input.done"}` | End of utterance: flushes the buffer and keeps the connection open | +| `{"type": "session.close"}` | End of connection | Server -> Client: | Message | Description | |---------|-------------| -| `{"type": "audio.start", "sentence_index": 0, "sentence_text": "...", "format": "pcm", "sample_rate": 24000}` | Audio generation starting for the buffered input | +| `{"type": "audio.start", "utterance_index": 0, "sentence_index": 0, "sentence_text": "...", "format": "pcm", "sample_rate": 24000}` | Audio generation starting for the buffered input | | Binary frame | Raw audio bytes (one or more PCM chunks when `stream_audio=true`) | -| `{"type": "audio.done", "sentence_index": 0, "total_bytes": 96000, "error": false}` | Audio complete for the buffered input | -| `{"type": "session.done", "total_sentences": N}` | Session complete | +| `{"type": "audio.done", "utterance_index": 0, "sentence_index": 0, "total_bytes": 96000, "error": false}` | Audio complete for the buffered input | +| `{"type": "session.done", "utterance_index": 0, "total_sentences": N}` | Flushed utterance complete | | `{"type": "error", "message": "..."}` | Non-fatal error | +### Flushing vs. Closing + +`input.done` is a flush, not a disconnect. The server synthesizes the buffered +text, emits `session.done`, and then waits on the same connection for the next +utterance, so a client that speaks repeatedly (for example one driven by an +upstream LLM) pays the WebSocket handshake once instead of once per utterance. + +* The session config is sticky. Send `input.text` again straight after + `session.done` to reuse it, or send another `session.config` first to change + voice, format, or reference audio. A `session.config` sent while text is + still buffered is rejected so no pending input is silently dropped. +* An utterance is the flush unit, not a linguistic one: it is whatever text was + buffered when `input.done` arrived, of any length, synthesized as one request. + `utterance_index` counts those flushes across the connection, so it tells you + which `input.done` a frame belongs to. `sentence_index` counts within one + flush and so pairs with `total_sentences`, which means every utterance reports + `sentence_index: 0` of `total_sentences: 1` (or `0` for an empty buffer). +* End the connection with `session.close`, or by closing the socket. An idle + connection is still closed after the server's idle timeout, which now also + applies to the gap between utterances. + ### Session Config Parameters All REST API parameters are supported, plus: diff --git a/docs/user_guide/examples/online_serving/text_to_speech.md b/docs/user_guide/examples/online_serving/text_to_speech.md index 5640845d8..c6a613723 100644 --- a/docs/user_guide/examples/online_serving/text_to_speech.md +++ b/docs/user_guide/examples/online_serving/text_to_speech.md @@ -422,12 +422,19 @@ curl -X POST http://localhost:8091/v1/audio/speech \ Raw PCM streaming requires `stream_format="audio"`, `response_format="pcm"`, and `async_chunk: true` on the stage config (default in `qwen3_tts.yaml`). `speed` is not supported when streaming. ### Streaming WebSocket -The `/v1/audio/speech/stream` endpoint accepts text incrementally, splits it at sentence boundaries, and emits one PCM stream per sentence: +The `/v1/audio/speech/stream` endpoint accepts text incrementally and synthesizes the buffered text as one continuous request on `input.done`: ```bash python qwen3_tts/streaming_speech_client.py --text "Hello world. How are you? I am fine." python qwen3_tts/streaming_speech_client.py --text "..." --simulate-stt --stt-delay 0.1 ``` +`input.done` flushes without closing, so repeating `--text` synthesizes several utterances over one connection: +```bash +python qwen3_tts/streaming_speech_client.py \ + --text "First utterance." \ + --text "Second utterance, same connection." +``` + ### Gradio demos ```bash ./qwen3_tts/run_gradio_demo.sh # CustomVoice (default) diff --git a/examples/online_serving/text_to_speech/README.md b/examples/online_serving/text_to_speech/README.md index 201cb5d78..7db0a1844 100644 --- a/examples/online_serving/text_to_speech/README.md +++ b/examples/online_serving/text_to_speech/README.md @@ -537,12 +537,19 @@ curl -X POST http://localhost:8091/v1/audio/speech \ Raw PCM streaming requires `stream_format="audio"`, `response_format="pcm"`, and `async_chunk: true` on the stage config (default in `qwen3_tts.yaml`). `speed` is not supported when streaming. ### Streaming WebSocket -The `/v1/audio/speech/stream` endpoint accepts text incrementally, splits it at sentence boundaries, and emits one PCM stream per sentence: +The `/v1/audio/speech/stream` endpoint accepts text incrementally and synthesizes the buffered text as one continuous request on `input.done`: ```bash python qwen3_tts/streaming_speech_client.py --text "Hello world. How are you? I am fine." python qwen3_tts/streaming_speech_client.py --text "..." --simulate-stt --stt-delay 0.1 ``` +`input.done` flushes without closing, so repeating `--text` synthesizes several utterances over one connection: +```bash +python qwen3_tts/streaming_speech_client.py \ + --text "First utterance." \ + --text "Second utterance, same connection." +``` + To receive word-level timestamps, launch the server with a forced aligner: ```bash vllm-omni serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \ diff --git a/examples/online_serving/text_to_speech/qwen3_tts/streaming_speech_client.py b/examples/online_serving/text_to_speech/qwen3_tts/streaming_speech_client.py index 92a40739e..e141a4c48 100644 --- a/examples/online_serving/text_to_speech/qwen3_tts/streaming_speech_client.py +++ b/examples/online_serving/text_to_speech/qwen3_tts/streaming_speech_client.py @@ -7,6 +7,12 @@ Usage: # Send full text at once python streaming_speech_client.py --text "Hello world. How are you? I am fine." + # Several utterances over one connection: input.done flushes each of them + # and the connection stays open, so only the first pays the handshake + python streaming_speech_client.py \ + --text "First utterance." \ + --text "Second utterance, same connection." + # Pick a built-in speaker for CustomVoice models python streaming_speech_client.py \ --text "打开电子枪，关闭扫描，切换到低倍模式。" \ @@ -64,18 +70,27 @@ except ImportError: raise SystemExit(1) -def save_audio_file(output_dir: str, sentence_index: int, response_format: str, chunks: list[bytes]) -> str: +def frame_basename(msg: dict) -> str: + """Name a file after the utterance and sentence a frame belongs to. + + Every utterance is sentence 0, so the utterance index is what keeps the + files of one connection from overwriting each other. + """ + return f"utterance_{msg['utterance_index']:03d}_sentence_{msg['sentence_index']:03d}" + + +def save_audio_file(output_dir: str, basename: str, response_format: str, chunks: list[bytes]) -> str: audio_bytes = b"".join(chunks) if response_format == "pcm": - raw_path = os.path.join(output_dir, f"sentence_{sentence_index:03d}.pcm") + raw_path = os.path.join(output_dir, f"{basename}.pcm") with open(raw_path, "wb") as f: f.write(audio_bytes) - wav_path = os.path.join(output_dir, f"sentence_{sentence_index:03d}.wav") + wav_path = os.path.join(output_dir, f"{basename}.wav") save_pcm_wav(wav_path, audio_bytes) return wav_path - filename = os.path.join(output_dir, f"sentence_{sentence_index:03d}.{response_format}") + filename = os.path.join(output_dir, f"{basename}.{response_format}") with open(filename, "wb") as f: f.write(audio_bytes) return filename @@ -95,183 +110,207 @@ def timestamp_words_text(timestamps: list[dict] | None) -> str: return " ".join(str(item.get("word", "")) for item in timestamps if item.get("word")) +async def send_utterance(ws, text: str, simulate_stt: bool, stt_delay: float) -> None: + """Send one utterance's text, then input.done to flush it.""" + if simulate_stt: + words = text.split(" ") + for i, word in enumerate(words): + chunk = word + (" " if i < len(words) - 1 else "") + await ws.send( + json.dumps( + { + "type": "input.text", + "text": chunk, + } + ) + ) + print(f" Sent: {chunk!r}") + await asyncio.sleep(stt_delay) + else: + await ws.send( + json.dumps( + { + "type": "input.text", + "text": text, + } + ) + ) + print(f"Sent full text: {text!r}") + + await ws.send(json.dumps({"type": "input.done"})) + print("Sent input.done") + + +async def receive_utterance( + ws, + output_dir: str, + response_format: str, + word_timestamps: bool, + save_chunks: bool, +) -> None: + """Consume frames until session.done marks the flushed utterance complete.""" + current_label = "utterance 0" + current_sentence_text = "" + current_chunks: list[bytes] = [] + current_timestamps: list[dict] = [] + # Distinguish `null` (aligner failed) from `[]` (silence / no tokens): + # the sentence is a failure only if no frame ever carried timestamps. + alignment_frame_seen = False + + while True: + message = await ws.recv() + + if isinstance(message, bytes): + current_chunks.append(message) + print(f" Received audio chunk for {current_label}: {len(message)} bytes") + else: + # JSON frame + msg = json.loads(message) + msg_type = msg.get("type") + + if msg_type == "audio.start": + current_label = f"utterance {msg['utterance_index']}" + current_sentence_text = msg.get("sentence_text", "") + current_chunks = [] + current_timestamps = [] + alignment_frame_seen = False + print(f" [{current_label}] Generating: {current_sentence_text!r}") + elif msg_type == "audio.chunk": + audio = base64.b64decode(msg["audio_b64"]) + current_chunks.append(audio) + chunk_timestamps = msg.get("timestamps") + # Sentence-level alignment sends a trailing timestamp-only + # frame with empty audio; don't write an empty chunk file. + if save_chunks and audio: + chunk_base = os.path.join( + output_dir, + f"{frame_basename(msg)}_chunk_{msg['chunk_id']:03d}", + ) + chunk_pcm = f"{chunk_base}.pcm" + chunk_wav = f"{chunk_base}.wav" + chunk_json = f"{chunk_base}_timestamps.json" + with open(chunk_pcm, "wb") as f: + f.write(audio) + save_pcm_wav(chunk_wav, audio, int(msg.get("sample_rate", 24000))) + with open(chunk_json, "w", encoding="utf-8") as f: + json.dump( + { + "utterance_index": msg["utterance_index"], + "sentence_index": msg["sentence_index"], + "chunk_id": msg["chunk_id"], + "chunk_start_ms": msg.get("chunk_start_ms"), + "chunk_end_ms": msg.get("chunk_end_ms"), + "sentence_text": current_sentence_text, + "covered_text": timestamp_words_text(chunk_timestamps), + "timestamps": chunk_timestamps, + }, + f, + ensure_ascii=False, + indent=2, + ) + if chunk_timestamps is None: + print(f" [{current_label}] chunk {msg['chunk_id']}: {len(audio)} bytes, timestamps unavailable") + else: + alignment_frame_seen = True + current_timestamps.extend(chunk_timestamps) + print( + f" [{current_label}] chunk {msg['chunk_id']}: " + f"{len(audio)} bytes, {len(chunk_timestamps)} timestamp(s)" + ) + if save_chunks and audio: + print(f" saved chunk -> {chunk_wav}") + elif msg_type == "audio.done": + filename = save_audio_file( + output_dir, + frame_basename(msg), + response_format, + current_chunks, + ) + if word_timestamps: + # `null` on failure, else the accumulated list (`[]` for silence). + sentence_timestamps = current_timestamps if alignment_frame_seen else None + ts_filename = os.path.join( + output_dir, + f"{frame_basename(msg)}_timestamps.json", + ) + with open(ts_filename, "w", encoding="utf-8") as f: + json.dump( + { + "utterance_index": msg["utterance_index"], + "sentence_index": msg["sentence_index"], + "sentence_text": current_sentence_text, + "covered_text": timestamp_words_text(sentence_timestamps), + "timestamps": sentence_timestamps, + }, + f, + ensure_ascii=False, + indent=2, + ) + print( + f" [{current_label}] Done" + f" bytes={msg.get('total_bytes', len(b''.join(current_chunks)))}" + f" error={msg.get('error', False)}" + f" -> {filename}" + ) + if word_timestamps: + print(f" [{current_label}] Timestamps -> {ts_filename}") + current_chunks = [] + current_timestamps = [] + alignment_frame_seen = False + elif msg_type == "session.done": + print(f"\nUtterance {msg['utterance_index']} complete: {msg['total_sentences']} sentence(s) generated") + return + elif msg_type == "error": + print(f" ERROR: {msg['message']}") + else: + print(f" Unknown message: {msg}") + + async def stream_tts( url: str, - text: str, + texts: list[str], config: dict, output_dir: str, simulate_stt: bool = False, stt_delay: float = 0.1, save_chunks: bool = False, ) -> None: - """Connect to the streaming TTS endpoint and process audio responses.""" + """Connect to the streaming TTS endpoint and process audio responses. + + Every text is synthesized over the same connection: input.done flushes an + utterance without closing, so only the first one pays the handshake. + """ os.makedirs(output_dir, exist_ok=True) + response_format = config.get("response_format", "wav") + word_timestamps = bool(config.get("word_timestamps", False)) async with websockets.connect(url) as ws: - # 1. Send session config + # 1. Send session config; it stays in effect for every utterance below. config_msg = {"type": "session.config", **config} await ws.send(json.dumps(config_msg)) print(f"Sent session config: {config}") - # 2. Send text (either all at once or word-by-word) - async def send_text(): - if simulate_stt: - words = text.split(" ") - for i, word in enumerate(words): - chunk = word + (" " if i < len(words) - 1 else "") - await ws.send( - json.dumps( - { - "type": "input.text", - "text": chunk, - } - ) - ) - print(f" Sent: {chunk!r}") - await asyncio.sleep(stt_delay) - else: - await ws.send( - json.dumps( - { - "type": "input.text", - "text": text, - } - ) - ) - print(f"Sent full text: {text!r}") - - # 3. Signal end of input - await ws.send(json.dumps({"type": "input.done"})) - print("Sent input.done") - - # Run sender and receiver concurrently - sender_task = asyncio.create_task(send_text()) - - response_format = config.get("response_format", "wav") - word_timestamps = bool(config.get("word_timestamps", False)) - current_sentence_index = 0 - current_sentence_text = "" - current_chunks: list[bytes] = [] - current_timestamps: list[dict] = [] - # Distinguish `null` (aligner failed) from `[]` (silence / no tokens): - # the sentence is a failure only if no frame ever carried timestamps. - alignment_frame_seen = False - - try: - while True: - message = await ws.recv() - - if isinstance(message, bytes): - current_chunks.append(message) - print(f" Received audio chunk for sentence {current_sentence_index}: {len(message)} bytes") - else: - # JSON frame - msg = json.loads(message) - msg_type = msg.get("type") - - if msg_type == "audio.start": - current_sentence_index = msg["sentence_index"] - current_sentence_text = msg.get("sentence_text", "") - current_chunks = [] - current_timestamps = [] - alignment_frame_seen = False - print(f" [sentence {msg['sentence_index']}] Generating: {current_sentence_text!r}") - elif msg_type == "audio.chunk": - audio = base64.b64decode(msg["audio_b64"]) - current_chunks.append(audio) - chunk_timestamps = msg.get("timestamps") - # Sentence-level alignment sends a trailing timestamp-only - # frame with empty audio; don't write an empty chunk file. - if save_chunks and audio: - chunk_base = os.path.join( - output_dir, - f"sentence_{msg['sentence_index']:03d}_chunk_{msg['chunk_id']:03d}", - ) - chunk_pcm = f"{chunk_base}.pcm" - chunk_wav = f"{chunk_base}.wav" - chunk_json = f"{chunk_base}_timestamps.json" - with open(chunk_pcm, "wb") as f: - f.write(audio) - save_pcm_wav(chunk_wav, audio, int(msg.get("sample_rate", 24000))) - with open(chunk_json, "w", encoding="utf-8") as f: - json.dump( - { - "sentence_index": msg["sentence_index"], - "chunk_id": msg["chunk_id"], - "chunk_start_ms": msg.get("chunk_start_ms"), - "chunk_end_ms": msg.get("chunk_end_ms"), - "sentence_text": current_sentence_text, - "covered_text": timestamp_words_text(chunk_timestamps), - "timestamps": chunk_timestamps, - }, - f, - ensure_ascii=False, - indent=2, - ) - if chunk_timestamps is None: - print( - f" [sentence {msg['sentence_index']}] chunk {msg['chunk_id']}: " - f"{len(audio)} bytes, timestamps unavailable" - ) - else: - alignment_frame_seen = True - current_timestamps.extend(chunk_timestamps) - print( - f" [sentence {msg['sentence_index']}] chunk {msg['chunk_id']}: " - f"{len(audio)} bytes, {len(chunk_timestamps)} timestamp(s)" - ) - if save_chunks and audio: - print(f" saved chunk -> {chunk_wav}") - elif msg_type == "audio.done": - filename = save_audio_file( - output_dir, - msg["sentence_index"], - response_format, - current_chunks, - ) - if word_timestamps: - # `null` on failure, else the accumulated list (`[]` for silence). - sentence_timestamps = current_timestamps if alignment_frame_seen else None - ts_filename = os.path.join( - output_dir, - f"sentence_{msg['sentence_index']:03d}_timestamps.json", - ) - with open(ts_filename, "w", encoding="utf-8") as f: - json.dump( - { - "sentence_index": msg["sentence_index"], - "sentence_text": current_sentence_text, - "covered_text": timestamp_words_text(sentence_timestamps), - "timestamps": sentence_timestamps, - }, - f, - ensure_ascii=False, - indent=2, - ) - print( - f" [sentence {msg['sentence_index']}] Done" - f" bytes={msg.get('total_bytes', len(b''.join(current_chunks)))}" - f" error={msg.get('error', False)}" - f" -> {filename}" - ) - if word_timestamps: - print(f" [sentence {msg['sentence_index']}] Timestamps -> {ts_filename}") - current_chunks = [] - current_timestamps = [] - alignment_frame_seen = False - elif msg_type == "session.done": - print(f"\nSession complete: {msg['total_sentences']} sentence(s) generated") - break - elif msg_type == "error": - print(f" ERROR: {msg['message']}") - else: - print(f" Unknown message: {msg}") - finally: - sender_task.cancel() + for text in texts: + # 2. Send text (either all at once or word-by-word), then flush. + sender_task = asyncio.create_task(send_utterance(ws, text, simulate_stt, stt_delay)) try: - await sender_task - except asyncio.CancelledError: - pass # Task cancellation is expected during shutdown + await receive_utterance( + ws, + output_dir=output_dir, + response_format=response_format, + word_timestamps=word_timestamps, + save_chunks=save_chunks, + ) + finally: + sender_task.cancel() + try: + await sender_task + except asyncio.CancelledError: + pass # Task cancellation is expected during shutdown + + # 3. Release the connection once there is nothing left to say. + await ws.send(json.dumps({"type": "session.close"})) + print("Sent session.close") print(f"\nAudio files saved to: {output_dir}/") @@ -286,7 +325,8 @@ def main(): parser.add_argument( "--text", required=True, - help="Text to synthesize", + action="append", + help="Text to synthesize; repeat to send several utterances over one connection", ) parser.add_argument( "--output-dir", @@ -385,7 +425,7 @@ def main(): asyncio.run( stream_tts( url=args.url, - text=args.text, + texts=args.text, config=config, output_dir=args.output_dir, simulate_stt=args.simulate_stt, diff --git a/tests/entrypoints/openai_api/test_qwen3_tts_websocket.py b/tests/entrypoints/openai_api/test_qwen3_tts_websocket.py index ecae21158..132d43858 100644 --- a/tests/entrypoints/openai_api/test_qwen3_tts_websocket.py +++ b/tests/entrypoints/openai_api/test_qwen3_tts_websocket.py @@ -33,26 +33,63 @@ tts_ws_server_params = [ ] -async def _run_ws_session(host: str, port: int, model: str) -> dict: - uri = f"ws://{host}:{port}/v1/audio/speech/stream" +def _session_config(model: str) -> str: + return json.dumps( + { + "type": "session.config", + "model": model, + "voice": "vivian", + "language": "English", + "response_format": "pcm", + "stream_audio": True, + } + ) + + +async def _collect_utterance(ws) -> dict: + """Read frames until session.done closes out one flushed utterance.""" starts: list[dict] = [] dones: list[dict] = [] chunk_lengths: dict[int, list[int]] = {} session_done: dict | None = None + while True: + message = await asyncio.wait_for(ws.recv(), timeout=180) + if isinstance(message, bytes): + if not starts: + raise AssertionError("Received audio bytes before audio.start") + sentence_index = starts[-1]["sentence_index"] + chunk_lengths.setdefault(sentence_index, []).append(len(message)) + continue + + payload = json.loads(message) + msg_type = payload.get("type") + if msg_type == "audio.start": + starts.append(payload) + chunk_lengths.setdefault(payload["sentence_index"], []) + elif msg_type == "audio.done": + dones.append(payload) + elif msg_type == "session.done": + session_done = payload + break + elif msg_type == "error": + raise AssertionError(f"WebSocket error: {payload['message']}") + else: + raise AssertionError(f"Unexpected WebSocket message: {payload}") + + return { + "starts": starts, + "dones": dones, + "chunk_lengths": chunk_lengths, + "session_done": session_done, + } + + +async def _run_ws_session(host: str, port: int, model: str) -> dict: + uri = f"ws://{host}:{port}/v1/audio/speech/stream" + async with websockets.connect(uri, max_size=None) as ws: - await ws.send( - json.dumps( - { - "type": "session.config", - "model": model, - "voice": "vivian", - "language": "English", - "response_format": "pcm", - "stream_audio": True, - } - ) - ) + await ws.send(_session_config(model)) await ws.send( json.dumps( { @@ -67,36 +104,23 @@ async def _run_ws_session(host: str, port: int, model: str) -> dict: ) await ws.send(json.dumps({"type": "input.done"})) - while True: - message = await asyncio.wait_for(ws.recv(), timeout=180) - if isinstance(message, bytes): - if not starts: - raise AssertionError("Received audio bytes before audio.start") - sentence_index = starts[-1]["sentence_index"] - chunk_lengths.setdefault(sentence_index, []).append(len(message)) - continue + return await _collect_utterance(ws) - payload = json.loads(message) - msg_type = payload.get("type") - if msg_type == "audio.start": - starts.append(payload) - chunk_lengths.setdefault(payload["sentence_index"], []) - elif msg_type == "audio.done": - dones.append(payload) - elif msg_type == "session.done": - session_done = payload - break - elif msg_type == "error": - raise AssertionError(f"WebSocket error: {payload['message']}") - else: - raise AssertionError(f"Unexpected WebSocket message: {payload}") - return { - "starts": starts, - "dones": dones, - "chunk_lengths": chunk_lengths, - "session_done": session_done, - } +async def _run_ws_reused_connection(host: str, port: int, model: str, texts: list[str]) -> list[dict]: + """Synthesize several utterances over a single connection.""" + uri = f"ws://{host}:{port}/v1/audio/speech/stream" + results: list[dict] = [] + + async with websockets.connect(uri, max_size=None) as ws: + await ws.send(_session_config(model)) + for text in texts: + await ws.send(json.dumps({"type": "input.text", "text": text})) + await ws.send(json.dumps({"type": "input.done"})) + results.append(await _collect_utterance(ws)) + await ws.send(json.dumps({"type": "session.close"})) + + return results class TestQwen3TTSWebSocket: @@ -113,6 +137,7 @@ class TestQwen3TTSWebSocket: session_done = result["session_done"] assert session_done is not None + assert session_done["utterance_index"] == 0 assert session_done["total_sentences"] == 1 assert len(starts) == 1 assert len(dones) == 1 @@ -131,3 +156,32 @@ class TestQwen3TTSWebSocket: assert total_bytes > 0 assert chunk_lengths[sentence_index], f"Expected binary PCM frames for sentence {sentence_index}" assert sum(chunk_lengths[sentence_index]) == total_bytes + + @pytest.mark.advanced_model + @pytest.mark.tts + @hardware_test(res={"cuda": "L4"}, num_cards=1) + @pytest.mark.parametrize("omni_server", tts_ws_server_params, indirect=True) + def test_input_done_flushes_without_closing_connection(self, omni_server) -> None: + # input.done must flush the buffer and leave the connection usable: + # the second utterance is synthesized without a second handshake. + texts = [ + "This is the first utterance sent over the websocket connection.", + "This is the second utterance, reusing the very same connection.", + ] + results = asyncio.run(_run_ws_reused_connection(omni_server.host, omni_server.port, omni_server.model, texts)) + + assert len(results) == len(texts) + for expected_index, result in enumerate(results): + assert result["session_done"] == { + "type": "session.done", + "utterance_index": expected_index, + "total_sentences": 1, + } + assert len(result["starts"]) == 1 + assert len(result["dones"]) == 1 + # utterance_index counts the flushes of the connection, while + # sentence_index stays within the flush it belongs to. + assert result["starts"][0]["utterance_index"] == expected_index + assert result["starts"][0]["sentence_index"] == 0 + assert result["dones"][0]["error"] is False + assert result["dones"][0]["total_bytes"] > 0 diff --git a/tests/entrypoints/openai_api/test_serving_speech_stream.py b/tests/entrypoints/openai_api/test_serving_speech_stream.py index 881717363..44e84b346 100644 --- a/tests/entrypoints/openai_api/test_serving_speech_stream.py +++ b/tests/entrypoints/openai_api/test_serving_speech_stream.py @@ -72,13 +72,162 @@ class TestStreamingSpeechWebSocket: assert audio.startswith(b"RIFF") done = ws.receive_json() - assert done == {"type": "audio.done", "sentence_index": 0, "total_bytes": len(audio), "error": False} + assert done == { + "type": "audio.done", + "utterance_index": 0, + "sentence_index": 0, + "total_bytes": len(audio), + "error": False, + } session_done = ws.receive_json() - assert session_done == {"type": "session.done", "total_sentences": 1} + assert session_done == {"type": "session.done", "utterance_index": 0, "total_sentences": 1} assert speech_service._generate_audio_bytes.await_count == 1 + def test_input_done_flushes_and_keeps_connection_open(self, mocker: MockerFixture): + app, speech_service = _build_test_app(mocker=mocker) + + with TestClient(app) as client: + with client.websocket_connect("/v1/audio/speech/stream") as ws: + ws.send_json({"type": "session.config", "voice": "Vivian"}) + + # The same connection serves several utterances; the config sent + # once at the top keeps applying and utterance_index rises. + for expected_index, text in enumerate(("First utterance. ", "Second utterance. ")): + ws.send_json({"type": "input.text", "text": text}) + ws.send_json({"type": "input.done"}) + + start = ws.receive_json() + assert start["type"] == "audio.start" + assert start["utterance_index"] == expected_index + assert start["sentence_text"] == text.strip() + assert ws.receive_bytes().startswith(b"RIFF") + assert ws.receive_json()["type"] == "audio.done" + assert ws.receive_json() == { + "type": "session.done", + "utterance_index": expected_index, + "total_sentences": 1, + } + + assert speech_service._generate_audio_bytes.await_count == 2 + assert [call.args[0].voice for call in speech_service._generate_audio_bytes.await_args_list] == [ + "Vivian", + "Vivian", + ] + + def test_sentence_index_stays_within_the_flushed_utterance(self, mocker: MockerFixture): + # sentence_index counts within one flush and utterance_index counts the + # flushes, so a late utterance never reports "sentence 2 of 1". + app, _ = _build_test_app(mocker=mocker) + + with TestClient(app) as client: + with client.websocket_connect("/v1/audio/speech/stream") as ws: + ws.send_json({"type": "session.config", "voice": "Vivian"}) + + for expected_index in range(3): + ws.send_json({"type": "input.text", "text": "Hello world. "}) + ws.send_json({"type": "input.done"}) + + start = ws.receive_json() + assert start["utterance_index"] == expected_index + assert start["sentence_index"] == 0 + ws.receive_bytes() + + done = ws.receive_json() + assert done["utterance_index"] == expected_index + assert done["sentence_index"] == 0 + + session_done = ws.receive_json() + assert session_done["utterance_index"] == expected_index + assert session_done["total_sentences"] == 1 + assert start["sentence_index"] < session_done["total_sentences"] + + def test_session_config_between_utterances_replaces_config(self, mocker: MockerFixture): + app, speech_service = _build_test_app(mocker=mocker) + + with TestClient(app) as client: + with client.websocket_connect("/v1/audio/speech/stream") as ws: + for expected_index, voice in enumerate(("Vivian", "Serena")): + ws.send_json({"type": "session.config", "voice": voice}) + ws.send_json({"type": "input.text", "text": "Hello world. "}) + ws.send_json({"type": "input.done"}) + + assert ws.receive_json()["type"] == "audio.start" + ws.receive_bytes() + assert ws.receive_json()["type"] == "audio.done" + assert ws.receive_json() == { + "type": "session.done", + "utterance_index": expected_index, + "total_sentences": 1, + } + + assert [call.args[0].voice for call in speech_service._generate_audio_bytes.await_args_list] == [ + "Vivian", + "Serena", + ] + + def test_session_config_rejected_while_input_is_buffered(self, mocker: MockerFixture): + app, speech_service = _build_test_app(mocker=mocker) + + with TestClient(app) as client: + with client.websocket_connect("/v1/audio/speech/stream") as ws: + ws.send_json({"type": "session.config", "voice": "Vivian"}) + ws.send_json({"type": "input.text", "text": "Hello world. "}) + ws.send_json({"type": "session.config", "voice": "Serena"}) + + error = ws.receive_json() + assert error["type"] == "error" + assert "while input is buffered" in error["message"] + + # The buffered text survives the rejected reconfiguration. + ws.send_json({"type": "input.done"}) + assert ws.receive_json()["sentence_text"] == "Hello world." + ws.receive_bytes() + assert ws.receive_json()["type"] == "audio.done" + assert ws.receive_json() == {"type": "session.done", "utterance_index": 0, "total_sentences": 1} + + assert speech_service._generate_audio_bytes.await_args_list[0].args[0].voice == "Vivian" + + def test_session_close_ends_connection(self, mocker: MockerFixture): + app, _ = _build_test_app(mocker=mocker) + + with TestClient(app) as client: + with client.websocket_connect("/v1/audio/speech/stream") as ws: + ws.send_json({"type": "session.config", "voice": "Vivian"}) + ws.send_json({"type": "input.text", "text": "Hello world. "}) + ws.send_json({"type": "input.done"}) + + assert ws.receive_json()["type"] == "audio.start" + ws.receive_bytes() + assert ws.receive_json()["type"] == "audio.done" + assert ws.receive_json() == {"type": "session.done", "utterance_index": 0, "total_sentences": 1} + + ws.send_json({"type": "session.close"}) + with pytest.raises(WebSocketDisconnect): + ws.receive_json() + + def test_idle_timeout_closes_reused_connection(self, mocker: MockerFixture): + app, _ = _build_test_app(idle_timeout=0.05, mocker=mocker) + + with TestClient(app) as client: + with client.websocket_connect("/v1/audio/speech/stream") as ws: + ws.send_json({"type": "session.config", "voice": "Vivian"}) + ws.send_json({"type": "input.text", "text": "Hello world. "}) + ws.send_json({"type": "input.done"}) + + assert ws.receive_json()["type"] == "audio.start" + ws.receive_bytes() + assert ws.receive_json()["type"] == "audio.done" + assert ws.receive_json() == {"type": "session.done", "utterance_index": 0, "total_sentences": 1} + + # Holding a connection open is not free: an idle client still + # gets timed out after the flush. + assert ws.receive_json() == { + "type": "error", + "message": "Idle timeout: no message received", + } + def test_streaming_multiple_binary_frames(self, mocker: MockerFixture): captured_requests = [] @@ -125,9 +274,15 @@ class TestStreamingSpeechWebSocket: assert ws.receive_bytes() == b"\x06" done = ws.receive_json() - assert done == {"type": "audio.done", "sentence_index": 0, "total_bytes": 6, "error": False} + assert done == { + "type": "audio.done", + "utterance_index": 0, + "sentence_index": 0, + "total_bytes": 6, + "error": False, + } - assert ws.receive_json() == {"type": "session.done", "total_sentences": 1} + assert ws.receive_json() == {"type": "session.done", "utterance_index": 0, "total_sentences": 1} assert len(captured_requests) == 1 assert captured_requests[0].stream is True @@ -210,6 +365,8 @@ class TestStreamingSpeechWebSocket: # is fully aligned. chunk = ws.receive_json() assert chunk["type"] == "audio.chunk" + assert chunk["utterance_index"] == 0 + assert chunk["sentence_index"] == 0 assert chunk["chunk_id"] == 0 assert chunk["chunk_start_ms"] == 0 assert chunk["chunk_end_ms"] == 500 @@ -240,7 +397,13 @@ class TestStreamingSpeechWebSocket: ] done = ws.receive_json() - assert done == {"type": "audio.done", "sentence_index": 0, "total_bytes": 2000, "error": False} + assert done == { + "type": "audio.done", + "utterance_index": 0, + "sentence_index": 0, + "total_bytes": 2000, + "error": False, + } assert captured_requests[0].word_timestamps is True assert mock_align.await_count == 1 @@ -309,11 +472,12 @@ class TestStreamingSpeechWebSocket: assert ws.receive_bytes() assert ws.receive_json() == { "type": "audio.done", + "utterance_index": 0, "sentence_index": 0, "total_bytes": 36, "error": False, } - assert ws.receive_json() == {"type": "session.done", "total_sentences": 1} + assert ws.receive_json() == {"type": "session.done", "utterance_index": 0, "total_sentences": 1} def test_invalid_streaming_config(self, mocker: MockerFixture): app, _ = _build_test_app(mocker=mocker) @@ -341,7 +505,7 @@ class TestStreamingSpeechWebSocket: ws.send_json({"type": "input.text", "text": ""}) ws.send_json({"type": "input.done"}) - assert ws.receive_json() == {"type": "session.done", "total_sentences": 0} + assert ws.receive_json() == {"type": "session.done", "utterance_index": 0, "total_sentences": 0} assert speech_service._generate_audio_bytes.await_count == 0 @@ -361,11 +525,12 @@ class TestStreamingSpeechWebSocket: ws.receive_bytes() assert ws.receive_json() == { "type": "audio.done", + "utterance_index": 0, "sentence_index": 0, "total_bytes": 36, "error": False, } - assert ws.receive_json() == {"type": "session.done", "total_sentences": 1} + assert ws.receive_json() == {"type": "session.done", "utterance_index": 0, "total_sentences": 1} def test_unknown_message_type_keeps_session_open(self, mocker: MockerFixture): app, _ = _build_test_app(mocker=mocker) @@ -384,12 +549,13 @@ class TestStreamingSpeechWebSocket: ws.receive_bytes() assert ws.receive_json() == { "type": "audio.done", + "utterance_index": 0, "sentence_index": 0, "total_bytes": 36, "error": False, } - assert ws.receive_json() == {"type": "session.done", "total_sentences": 1} + assert ws.receive_json() == {"type": "session.done", "utterance_index": 0, "total_sentences": 1} def test_config_timeout_closes_session(self, mocker: MockerFixture): app, _ = _build_test_app(config_timeout=0.01, mocker=mocker) @@ -416,10 +582,19 @@ class TestStreamingSpeechWebSocket: ws.send_json({"type": "input.done"}) assert ws.receive_json()["type"] == "audio.start" - assert ws.receive_json() == {"type": "error", "message": "Generation failed for sentence 0: boom"} - assert ws.receive_json() == {"type": "audio.done", "sentence_index": 0, "total_bytes": 0, "error": True} + assert ws.receive_json() == { + "type": "error", + "message": "Generation failed for utterance 0, sentence 0: boom", + } + assert ws.receive_json() == { + "type": "audio.done", + "utterance_index": 0, + "sentence_index": 0, + "total_bytes": 0, + "error": True, + } - assert ws.receive_json() == {"type": "session.done", "total_sentences": 1} + assert ws.receive_json() == {"type": "session.done", "utterance_index": 0, "total_sentences": 1} def test_streaming_generation_error_marks_audio_done(self, mocker: MockerFixture): speech_service = mocker.MagicMock(spec=OmniOpenAIServingSpeech) @@ -453,16 +628,17 @@ class TestStreamingSpeechWebSocket: assert ws.receive_bytes() == b"\x01\x02" assert ws.receive_json() == { "type": "error", - "message": "Generation failed for sentence 0: stream boom", + "message": "Generation failed for utterance 0, sentence 0: stream boom", } assert ws.receive_json() == { "type": "audio.done", + "utterance_index": 0, "sentence_index": 0, "total_bytes": 2, "error": True, } - assert ws.receive_json() == {"type": "session.done", "total_sentences": 1} + assert ws.receive_json() == {"type": "session.done", "utterance_index": 0, "total_sentences": 1} def test_invalid_input_text_type_returns_validation_error(self, mocker: MockerFixture): app, speech_service = _build_test_app(mocker=mocker) @@ -478,7 +654,7 @@ class TestStreamingSpeechWebSocket: } ws.send_json({"type": "input.done"}) - assert ws.receive_json() == {"type": "session.done", "total_sentences": 0} + assert ws.receive_json() == {"type": "session.done", "utterance_index": 0, "total_sentences": 0} assert speech_service._generate_audio_bytes.await_count == 0 @@ -497,7 +673,7 @@ class TestStreamingSpeechWebSocket: } ws.send_json({"type": "input.done"}) - assert ws.receive_json() == {"type": "session.done", "total_sentences": 0} + assert ws.receive_json() == {"type": "session.done", "utterance_index": 0, "total_sentences": 0} assert speech_service._generate_audio_bytes.await_count == 0 @@ -550,7 +726,15 @@ class TestStreamingSpeechWebSocket: config.word_timestamps = False with pytest.raises(WebSocketDisconnect): - asyncio.run(handler._generate_and_send(websocket, config, "Hello world.", 0)) + asyncio.run( + handler._generate_and_send( + websocket, + config, + "Hello world.", + utterance_index=0, + sentence_index=0, + ) + ) speech_service.engine_client.abort.assert_awaited_once_with("req-abort") assert websocket.send_json.await_count == 2 diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 44225b69b..6d088f65e 100755 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -1572,8 +1572,9 @@ async def delete_voice(name: str, raw_request: Request): async def streaming_speech(websocket: WebSocket): """WebSocket endpoint for streaming text input TTS. - Accepts text incrementally, splits at sentence boundaries, and - returns audio per sentence. See serving_speech_stream.py for protocol. + Accepts text incrementally and returns audio for the buffered text on + input.done, which flushes without closing the connection. See + serving_speech_stream.py for protocol. """ handler = getattr(websocket.app.state, "openai_streaming_speech", None) if handler is None: diff --git a/vllm_omni/entrypoints/openai/serving_speech_stream.py b/vllm_omni/entrypoints/openai/serving_speech_stream.py index 28acf8679..102cfda45 100644 --- a/vllm_omni/entrypoints/openai/serving_speech_stream.py +++ b/vllm_omni/entrypoints/openai/serving_speech_stream.py @@ -3,26 +3,46 @@ Accepts text incrementally via WebSocket, buffers it until input.done, and generates audio once for the buffered input using the existing TTS pipeline. +input.done is a flush, not a close: it ends the current utterance and the +connection stays open, so the next utterance reuses the same connection +instead of paying another WebSocket handshake. A connection ends on +session.close, on the idle timeout, or when the client closes the socket. +The session config is sticky across flushes and can be replaced by sending +another session.config between utterances. + +"Utterance" here names the flush unit rather than any linguistic unit: it is +whatever text the client had buffered when it sent input.done, from a word to +several paragraphs, synthesized as a single request. + Protocol: Client -> Server: - {"type": "session.config", ...} # Session config (sent once first) + {"type": "session.config", ...} # Session config (first message; repeatable) {"type": "input.text", "text": "..."} # Text chunks - {"type": "input.done"} # End of input + {"type": "input.done"} # End of utterance, flush and keep connection open + {"type": "session.close"} # End of connection Server -> Client (default, word_timestamps=false): - {"type": "audio.start", "sentence_index": 0, "sentence_text": "...", "format": "wav"} + {"type": "audio.start", "utterance_index": 0, "sentence_index": 0, + "sentence_text": "...", "format": "wav"} ... - {"type": "audio.done", "sentence_index": 0} - {"type": "session.done", "total_sentences": N} + {"type": "audio.done", "utterance_index": 0, "sentence_index": 0} + {"type": "session.done", "utterance_index": 0, "total_sentences": N} {"type": "error", "message": "..."} + # session.done ends the flushed utterance, not the connection. An + # utterance is just the flush unit: whatever text was buffered when + # input.done arrived, of any length. utterance_index counts those + # flushes across the connection, while sentence_index counts within + # one of them and so pairs with total_sentences. Server -> Client (when word_timestamps=true): - {"type": "audio.start", "sentence_index": 0, "sentence_text": "...", "format": "pcm"} - {"type": "audio.chunk", "sentence_index": 0, "chunk_id": 0, "audio_b64": "", "timestamps": null} + {"type": "audio.start", "utterance_index": 0, "sentence_index": 0, + "sentence_text": "...", "format": "pcm"} + {"type": "audio.chunk", "utterance_index": 0, "sentence_index": 0, "chunk_id": 0, + "audio_b64": "", "timestamps": null} ... {"type": "audio.chunk", "audio_b64": "", "timestamps": [{"word", "start_ms", "end_ms"}, ...]} - {"type": "audio.done", "sentence_index": 0} + {"type": "audio.done", "utterance_index": 0, "sentence_index": 0} # Audio is JSON base64 PCM (not binary). A trailing empty-audio chunk carries the # full sentence-relative alignment. timestamps: list = aligned, [] = silence, null = failed. """ @@ -57,9 +77,11 @@ _MAX_INPUT_TEXT_MESSAGE_SIZE = 128 * 1024 class OmniStreamingSpeechHandler: """Handles WebSocket sessions for streaming text-input TTS. - Each WebSocket connection is an independent session. Text arrives - incrementally, is buffered until input.done, and audio is generated once - for the buffered input using the existing OmniOpenAIServingSpeech pipeline. + A connection carries one or more utterances. Text arrives incrementally, + is buffered until input.done, and audio is generated once for the + buffered input using the existing OmniOpenAIServingSpeech pipeline. The + connection outlives each utterance so a client can keep synthesizing on + it without reconnecting. Args: speech_service: The existing TTS serving instance (reused for @@ -79,54 +101,64 @@ class OmniStreamingSpeechHandler: self._config_timeout = config_timeout async def handle_session(self, websocket: WebSocket) -> None: - """Main session loop for a single WebSocket connection.""" + """Main loop for a single WebSocket connection. + + Serves any number of utterances: text is buffered until input.done, + which flushes it as one TTS request and then leaves the connection + open for the next one. Rejecting a message is only fatal before the + first valid session.config; afterwards the error is reported and the + connection survives. + """ await websocket.accept() + config: StreamingSpeechSessionConfig | None = None + text_parts: list[str] = [] + utterance_index = 0 + try: - # 1. Wait for session.config - config = await self._receive_config(websocket) - if config is None: - return # Error already sent, connection closing - - # Validate model if specified - if config.model and hasattr(self._speech_service, "_check_model"): - error = await self._speech_service._check_model( - OpenAICreateSpeechRequest(input="ping", model=config.model) - ) - if error is not None: - await self._send_error(websocket, str(error)) - return - - text_parts: list[str] = [] - - # 2. Receive text chunks until input.done while True: try: raw = await asyncio.wait_for( websocket.receive_text(), - timeout=self._idle_timeout, + timeout=self._config_timeout if config is None else self._idle_timeout, ) except asyncio.TimeoutError: - await self._send_error(websocket, "Idle timeout: no message received") + if config is None: + await self._send_error(websocket, "Timeout waiting for session.config") + else: + await self._send_error(websocket, "Idle timeout: no message received") return - if len(raw) > _MAX_INPUT_TEXT_MESSAGE_SIZE: - await self._send_error(websocket, "input.text message too large") - continue - - try: - msg = json.loads(raw) - except json.JSONDecodeError: - await self._send_error(websocket, "Invalid JSON message") - continue - - if not isinstance(msg, dict): - await self._send_error(websocket, "WebSocket messages must be JSON objects") + msg = await self._parse_message(websocket, raw) + if msg is None: + if config is None: + return # Malformed handshake, connection closing continue msg_type = msg.get("type") - if msg_type == "input.text": + if msg_type == "session.config": + if text_parts: + await self._send_error( + websocket, + "session.config cannot be applied while input is buffered; send input.done first", + ) + continue + new_config = await self._build_config(websocket, msg) + if new_config is None: + if config is None: + return # Error already sent, connection closing + continue # Keep serving with the previous config + config = new_config + + elif config is None: + await self._send_error( + websocket, + f"Expected session.config, got: {msg_type}", + ) + return + + elif msg_type == "input.text": text = msg.get("text", "") if not isinstance(text, str): await self._send_error(websocket, "input.text requires a string value") @@ -135,17 +167,31 @@ class OmniStreamingSpeechHandler: elif msg_type == "input.done": full_text = "".join(text_parts).strip() + text_parts.clear() total_sentences = 0 if full_text: - await self._generate_and_send(websocket, config, full_text, 0) + # However long the buffered text is, the pipeline takes + # it as one request, so every flush is sentence 0 of 1. + await self._generate_and_send( + websocket, + config, + full_text, + utterance_index=utterance_index, + sentence_index=0, + ) total_sentences = 1 await websocket.send_json( { "type": "session.done", + "utterance_index": utterance_index, "total_sentences": total_sentences, } ) + utterance_index += 1 + + elif msg_type == "session.close": + await websocket.close() return else: @@ -163,44 +209,50 @@ class OmniStreamingSpeechHandler: except Exception: logger.debug("Failed to send error to streaming speech client", exc_info=True) - async def _receive_config(self, websocket: WebSocket) -> StreamingSpeechSessionConfig | None: - """Wait for and validate the session.config message.""" - try: - raw = await asyncio.wait_for( - websocket.receive_text(), - timeout=self._config_timeout, - ) - except asyncio.TimeoutError: - await self._send_error(websocket, "Timeout waiting for session.config") - return None + async def _parse_message(self, websocket: WebSocket, raw: str) -> dict | None: + """Decode one client message, or report why it was rejected. - if len(raw) > _MAX_CONFIG_MESSAGE_SIZE: - await self._send_error(websocket, "session.config message too large") + Size limits are per message type: session.config carries ref_audio + payloads and gets the larger budget. + """ + if len(raw) > max(_MAX_CONFIG_MESSAGE_SIZE, _MAX_INPUT_TEXT_MESSAGE_SIZE): + await self._send_error(websocket, "WebSocket message too large") return None try: msg = json.loads(raw) except json.JSONDecodeError: - await self._send_error(websocket, "Invalid JSON in session.config") + await self._send_error(websocket, "Invalid JSON message") return None if not isinstance(msg, dict): - await self._send_error(websocket, "session.config must be a JSON object") + await self._send_error(websocket, "WebSocket messages must be JSON objects") return None - if msg.get("type") != "session.config": - await self._send_error( - websocket, - f"Expected session.config, got: {msg.get('type')}", - ) + if msg.get("type") == "session.config": + if len(raw) > _MAX_CONFIG_MESSAGE_SIZE: + await self._send_error(websocket, "session.config message too large") + return None + elif len(raw) > _MAX_INPUT_TEXT_MESSAGE_SIZE: + await self._send_error(websocket, "input.text message too large") return None + return msg + + async def _build_config(self, websocket: WebSocket, msg: dict) -> StreamingSpeechSessionConfig | None: + """Validate a session.config message and its model.""" try: config = StreamingSpeechSessionConfig(**{k: v for k, v in msg.items() if k != "type"}) except ValidationError as e: await self._send_error(websocket, f"Invalid session config: {e}") return None + if config.model and hasattr(self._speech_service, "_check_model"): + error = await self._speech_service._check_model(OpenAICreateSpeechRequest(input="ping", model=config.model)) + if error is not None: + await self._send_error(websocket, str(error)) + return None + return config async def _generate_and_send( @@ -208,9 +260,15 @@ class OmniStreamingSpeechHandler: websocket: WebSocket, config: StreamingSpeechSessionConfig, sentence_text: str, + *, + utterance_index: int, sentence_index: int, ) -> None: - """Generate audio for a single sentence and send it over WebSocket.""" + """Generate audio for a single sentence and send it over WebSocket. + + ``utterance_index`` identifies the flush this sentence belongs to and + ``sentence_index`` its position inside that flush. + """ response_format = config.response_format or "wav" # Reject unmet word-timestamps preconditions early with a clear reason. @@ -253,6 +311,7 @@ class OmniStreamingSpeechHandler: start_payload = { "type": "audio.start", + "utterance_index": utterance_index, "sentence_index": sentence_index, "sentence_text": sentence_text, "format": response_format, @@ -277,6 +336,7 @@ class OmniStreamingSpeechHandler: request_id=request_id, generator=generator, sentence_text=sentence_text, + utterance_index=utterance_index, sentence_index=sentence_index, language=config.language, ) @@ -298,13 +358,22 @@ class OmniStreamingSpeechHandler: raise except Exception as e: generation_failed = True - logger.error("Generation failed for sentence %d: %s", sentence_index, e) - await self._send_error(websocket, f"Generation failed for sentence {sentence_index}: {e}") + logger.error( + "Generation failed for utterance %d, sentence %d: %s", + utterance_index, + sentence_index, + e, + ) + await self._send_error( + websocket, + f"Generation failed for utterance {utterance_index}, sentence {sentence_index}: {e}", + ) finally: try: await websocket.send_json( { "type": "audio.done", + "utterance_index": utterance_index, "sentence_index": sentence_index, "total_bytes": total_bytes, "error": generation_failed, @@ -320,6 +389,7 @@ class OmniStreamingSpeechHandler: request_id: str, generator, sentence_text: str, + utterance_index: int, sentence_index: int, language: str | None = None, ) -> int: @@ -350,6 +420,7 @@ class OmniStreamingSpeechHandler: await websocket.send_json( { "type": "audio.chunk", + "utterance_index": utterance_index, "sentence_index": sentence_index, "chunk_id": chunk_id, "chunk_start_ms": chunk_start_ms,