[Feature] Flush the speech WebSocket on input.done without closing it (#5517)

Signed-off-by: iancarrasco-b10 <ian.carrasco@baseten.co>
This commit is contained in:
iancarrasco-b10 authored and GitHub committed 2026-07-28 23:46:19 +00:00
1 parent 0feb92bdfa
commit 23bb4395bb
8 files changed
+688 -302

No files matched your search

+27 -5
View File
@@ -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:
@@ -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)
@@ -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 \
@@ -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,
@@ -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
@@ -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
+3 -2
View File
@@ -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:
@@ -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"}
<binary frame: audio bytes>
...
{"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": "<base64 PCM>", "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": "<base64 PCM>", "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,