994 lines
43 KiB
Python
994 lines
43 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
import asyncio
|
|
import base64
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
import torch
|
|
from fastapi import FastAPI, WebSocket
|
|
from pytest_mock import MockerFixture
|
|
from starlette.testclient import TestClient
|
|
from starlette.websockets import WebSocketDisconnect
|
|
|
|
from vllm_omni.entrypoints.openai import serving_speech_stream as streaming_speech_module
|
|
from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech
|
|
from vllm_omni.entrypoints.openai.serving_speech_stream import OmniStreamingSpeechHandler
|
|
from vllm_omni.model_executor.stage_input_processors.forced_aligner import ALIGNER_WORDS_KEY
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
def _fake_aligner_res(pairs, words):
|
|
"""Build a stand-in for the forced-aligner stage's pooling output.
|
|
|
|
Mirrors what rides the generator in production: ``res.outputs.data`` is a
|
|
``[n_words, 2]`` int32 tensor of ``[start_ms, end_ms]`` and the word strings
|
|
travel in ``additional_information``. Decoded by ``extract_word_timestamps``.
|
|
"""
|
|
return SimpleNamespace(
|
|
outputs=SimpleNamespace(data=torch.tensor(pairs, dtype=torch.int32), additional_information=None),
|
|
additional_information={ALIGNER_WORDS_KEY: list(words)},
|
|
)
|
|
|
|
|
|
def _build_test_app(
|
|
speech_service=None,
|
|
*,
|
|
idle_timeout=30.0,
|
|
config_timeout=10.0,
|
|
mocker: MockerFixture | None = None,
|
|
):
|
|
if speech_service is None:
|
|
assert mocker is not None
|
|
speech_service = mocker.MagicMock(spec=OmniOpenAIServingSpeech)
|
|
speech_service._generate_audio_bytes = mocker.AsyncMock(return_value=(b"RIFF" + b"\x00" * 32, "audio/wav"))
|
|
speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req-1", object(), {}))
|
|
speech_service.forced_aligner_enabled = False
|
|
|
|
async def mock_generate_pcm_chunks(
|
|
_generator,
|
|
_request_id,
|
|
*,
|
|
request_start_s=None,
|
|
request_arrival_ts=None,
|
|
include_sample_rate=False,
|
|
tts_params=None,
|
|
collect=None,
|
|
cumulative_audio=False,
|
|
):
|
|
for chunk in (b"\x01\x02", b"\x03\x04\x05"):
|
|
yield (chunk, 24000) if include_sample_rate else chunk
|
|
|
|
speech_service._generate_pcm_chunks = mock_generate_pcm_chunks
|
|
speech_service.engine_client = mocker.MagicMock()
|
|
speech_service.engine_client.abort = mocker.AsyncMock()
|
|
|
|
handler = OmniStreamingSpeechHandler(
|
|
speech_service=speech_service,
|
|
idle_timeout=idle_timeout,
|
|
config_timeout=config_timeout,
|
|
)
|
|
app = FastAPI()
|
|
|
|
@app.websocket("/v1/audio/speech/stream")
|
|
async def ws_endpoint(websocket: WebSocket):
|
|
await handler.handle_session(websocket)
|
|
|
|
return app, speech_service
|
|
|
|
|
|
class TestStreamingSpeechWebSocket:
|
|
def test_non_streaming_single_frame(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": "input.done"})
|
|
|
|
start = ws.receive_json()
|
|
assert start["type"] == "audio.start"
|
|
assert start["sentence_index"] == 0
|
|
assert start["sentence_text"] == "Hello world."
|
|
assert start["format"] == "wav"
|
|
|
|
audio = ws.receive_bytes()
|
|
assert audio.startswith(b"RIFF")
|
|
|
|
done = ws.receive_json()
|
|
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", "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 an utterance is in progress" 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 = []
|
|
captured_timing = {}
|
|
captured_tts_params = []
|
|
|
|
speech_service = mocker.MagicMock(spec=OmniOpenAIServingSpeech)
|
|
speech_service._generate_audio_bytes = mocker.AsyncMock(return_value=(b"", "audio/wav"))
|
|
speech_service.engine_client = mocker.MagicMock()
|
|
speech_service.engine_client.abort = mocker.AsyncMock()
|
|
speech_service.forced_aligner_enabled = False
|
|
|
|
async def mock_prepare_speech_generation(request, *, arrival_time=None):
|
|
captured_requests.append(request)
|
|
assert arrival_time is not None
|
|
captured_timing["prepare_arrival"] = arrival_time
|
|
return "req-stream", object(), {"_qwen3_tts_effective_max_tokens": [192]}
|
|
|
|
speech_service._prepare_speech_generation = mock_prepare_speech_generation
|
|
|
|
async def mock_generate_pcm_chunks(
|
|
_generator,
|
|
_request_id,
|
|
*,
|
|
request_start_s=None,
|
|
request_arrival_ts=None,
|
|
include_sample_rate=False,
|
|
tts_params=None,
|
|
):
|
|
assert request_start_s is not None
|
|
assert request_arrival_ts is not None
|
|
captured_timing["chunk_start"] = request_start_s
|
|
captured_timing["chunk_arrival"] = request_arrival_ts
|
|
captured_tts_params.append(tts_params)
|
|
for chunk in (b"\x01\x02", b"\x03\x04\x05", b"\x06"):
|
|
yield (chunk, 24000) if include_sample_rate else chunk
|
|
|
|
speech_service._generate_pcm_chunks = mock_generate_pcm_chunks
|
|
app, _ = _build_test_app(speech_service)
|
|
|
|
with TestClient(app) as client:
|
|
with client.websocket_connect("/v1/audio/speech/stream") as ws:
|
|
ws.send_json(
|
|
{
|
|
"type": "session.config",
|
|
"voice": "Vivian",
|
|
"stream_audio": True,
|
|
"response_format": "pcm",
|
|
"initial_codec_chunk_frames": 12,
|
|
}
|
|
)
|
|
ws.send_json({"type": "input.text", "text": "Hello world. "})
|
|
ws.send_json({"type": "input.done"})
|
|
|
|
start = ws.receive_json()
|
|
assert start["type"] == "audio.start"
|
|
assert start["format"] == "pcm"
|
|
assert start["sample_rate"] == 24000
|
|
|
|
assert ws.receive_bytes() == b"\x01\x02"
|
|
assert ws.receive_bytes() == b"\x03\x04\x05"
|
|
assert ws.receive_bytes() == b"\x06"
|
|
|
|
done = ws.receive_json()
|
|
assert done == {
|
|
"type": "audio.done",
|
|
"utterance_index": 0,
|
|
"sentence_index": 0,
|
|
"total_bytes": 6,
|
|
"error": False,
|
|
}
|
|
|
|
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
|
|
assert captured_requests[0].response_format == "pcm"
|
|
assert captured_requests[0].initial_codec_chunk_frames == 12
|
|
assert captured_timing["prepare_arrival"] == captured_timing["chunk_arrival"]
|
|
assert captured_timing["chunk_start"] > 0
|
|
assert captured_tts_params == [{"_qwen3_tts_effective_max_tokens": [192]}]
|
|
assert speech_service._generate_audio_bytes.await_count == 0
|
|
|
|
def test_word_timestamps_requires_configured_aligner(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",
|
|
"stream_audio": True,
|
|
"response_format": "pcm",
|
|
"word_timestamps": True,
|
|
}
|
|
)
|
|
ws.send_json({"type": "input.text", "text": "Hello world. "})
|
|
ws.send_json({"type": "input.done"})
|
|
|
|
error = ws.receive_json()
|
|
assert error["type"] == "error"
|
|
assert "without --forced-aligner" in error["message"]
|
|
|
|
def test_word_timestamps_emit_pipeline_json_frame(self, mocker: MockerFixture):
|
|
captured_requests = []
|
|
speech_service = mocker.MagicMock(spec=OmniOpenAIServingSpeech)
|
|
speech_service._generate_audio_bytes = mocker.AsyncMock(return_value=(b"", "audio/wav"))
|
|
speech_service.engine_client = mocker.MagicMock()
|
|
speech_service.engine_client.abort = mocker.AsyncMock()
|
|
speech_service.forced_aligner_enabled = True
|
|
|
|
async def mock_prepare_speech_generation(request, *, arrival_time=None):
|
|
captured_requests.append(request)
|
|
assert arrival_time is not None
|
|
return "req-stream", object(), {}
|
|
|
|
speech_service._prepare_speech_generation = mock_prepare_speech_generation
|
|
|
|
first_chunk = b"\x01" * 1000
|
|
second_chunk = b"\x02" * 1000
|
|
|
|
# The forced-aligner stage rides the same generator: its pooling output
|
|
# is surfaced via the ``collect`` channel once the audio has streamed.
|
|
async def mock_generate_pcm_chunks(
|
|
_generator,
|
|
_request_id,
|
|
*,
|
|
request_start_s=None,
|
|
request_arrival_ts=None,
|
|
include_sample_rate=False,
|
|
tts_params=None,
|
|
collect=None,
|
|
cumulative_audio=False,
|
|
):
|
|
assert request_start_s is not None
|
|
assert request_arrival_ts is not None
|
|
assert cumulative_audio is True
|
|
for chunk in (first_chunk, second_chunk):
|
|
yield (chunk, 1000) if include_sample_rate else chunk
|
|
if collect is not None:
|
|
collect["aligner_res"] = _fake_aligner_res([[0, 200], [200, 900]], ["Hello", "world"])
|
|
|
|
speech_service._generate_pcm_chunks = mock_generate_pcm_chunks
|
|
app, _ = _build_test_app(speech_service)
|
|
|
|
with TestClient(app) as client:
|
|
with client.websocket_connect("/v1/audio/speech/stream") as ws:
|
|
ws.send_json(
|
|
{
|
|
"type": "session.config",
|
|
"voice": "Vivian",
|
|
"stream_audio": True,
|
|
"response_format": "pcm",
|
|
"word_timestamps": True,
|
|
}
|
|
)
|
|
ws.send_json({"type": "input.text", "text": "Hello world. "})
|
|
ws.send_json({"type": "input.done"})
|
|
|
|
start = ws.receive_json()
|
|
assert start["type"] == "audio.start"
|
|
assert start["word_timestamps"] is True
|
|
|
|
# Audio streams first; timestamps are null until the sentence
|
|
# 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
|
|
assert chunk["sample_rate"] == 1000
|
|
assert base64.b64decode(chunk["audio_b64"]) == first_chunk
|
|
assert chunk["timestamps"] is None
|
|
|
|
chunk = ws.receive_json()
|
|
assert chunk["type"] == "audio.chunk"
|
|
assert chunk["chunk_id"] == 1
|
|
assert chunk["chunk_start_ms"] == 500
|
|
assert chunk["chunk_end_ms"] == 1000
|
|
assert chunk["sample_rate"] == 1000
|
|
assert base64.b64decode(chunk["audio_b64"]) == second_chunk
|
|
assert chunk["timestamps"] is None
|
|
|
|
# Final frame: empty audio carrying the whole-sentence timestamps.
|
|
chunk = ws.receive_json()
|
|
assert chunk["type"] == "audio.chunk"
|
|
assert chunk["chunk_id"] == 2
|
|
assert chunk["chunk_start_ms"] == 0
|
|
assert chunk["chunk_end_ms"] == 1000
|
|
assert chunk["sample_rate"] == 1000
|
|
assert base64.b64decode(chunk["audio_b64"]) == b""
|
|
assert chunk["timestamps"] == [
|
|
{"word": "Hello", "start_ms": 0, "end_ms": 200},
|
|
{"word": "world", "start_ms": 200, "end_ms": 900},
|
|
]
|
|
|
|
done = ws.receive_json()
|
|
assert done == {
|
|
"type": "audio.done",
|
|
"utterance_index": 0,
|
|
"sentence_index": 0,
|
|
"total_bytes": 2000,
|
|
"error": False,
|
|
}
|
|
|
|
assert captured_requests[0].word_timestamps is True
|
|
|
|
def test_word_timestamps_emit_word_dicts(self, mocker: MockerFixture):
|
|
# The streaming layer forwards the aligner's (already monotonic,
|
|
# non-overlapping) words as JSON dicts in the trailing frame.
|
|
speech_service = mocker.MagicMock(spec=OmniOpenAIServingSpeech)
|
|
speech_service._generate_audio_bytes = mocker.AsyncMock(return_value=(b"", "audio/wav"))
|
|
speech_service.engine_client = mocker.MagicMock()
|
|
speech_service.engine_client.abort = mocker.AsyncMock()
|
|
speech_service.forced_aligner_enabled = True
|
|
speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req", object(), {}))
|
|
|
|
async def mock_generate_pcm_chunks(
|
|
_generator,
|
|
_request_id,
|
|
*,
|
|
request_start_s=None,
|
|
request_arrival_ts=None,
|
|
include_sample_rate=False,
|
|
tts_params=None,
|
|
collect=None,
|
|
cumulative_audio=False,
|
|
):
|
|
chunk = b"\x01" * 1000
|
|
yield (chunk, 1000) if include_sample_rate else chunk
|
|
if collect is not None:
|
|
collect["aligner_res"] = _fake_aligner_res([[0, 1000], [1000, 1200]], ["Hello", "world"])
|
|
|
|
speech_service._generate_pcm_chunks = mock_generate_pcm_chunks
|
|
app, _ = _build_test_app(speech_service)
|
|
|
|
with TestClient(app) as client:
|
|
with client.websocket_connect("/v1/audio/speech/stream") as ws:
|
|
ws.send_json(
|
|
{
|
|
"type": "session.config",
|
|
"voice": "Vivian",
|
|
"stream_audio": True,
|
|
"response_format": "pcm",
|
|
"word_timestamps": True,
|
|
}
|
|
)
|
|
ws.send_json({"type": "input.text", "text": "Hello world. "})
|
|
ws.send_json({"type": "input.done"})
|
|
|
|
assert ws.receive_json()["type"] == "audio.start"
|
|
# Real-time audio chunk (timestamps null), then the timestamp frame.
|
|
assert ws.receive_json()["timestamps"] is None
|
|
final = ws.receive_json()
|
|
assert final["type"] == "audio.chunk"
|
|
timestamps = final["timestamps"]
|
|
assert timestamps == [
|
|
{"word": "Hello", "start_ms": 0, "end_ms": 1000},
|
|
{"word": "world", "start_ms": 1000, "end_ms": 1200},
|
|
]
|
|
for ts in timestamps:
|
|
assert ts["end_ms"] >= ts["start_ms"]
|
|
|
|
def test_flush_on_input_done(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 without punctuation"})
|
|
ws.send_json({"type": "input.done"})
|
|
|
|
assert ws.receive_json()["type"] == "audio.start"
|
|
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", "utterance_index": 0, "total_sentences": 1}
|
|
|
|
def test_invalid_streaming_config(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",
|
|
"stream_audio": True,
|
|
"response_format": "wav",
|
|
}
|
|
)
|
|
error = ws.receive_json()
|
|
assert error["type"] == "error"
|
|
assert "response_format='pcm'" in error["message"]
|
|
|
|
def test_empty_input_text_emits_no_audio(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": ""})
|
|
ws.send_json({"type": "input.done"})
|
|
|
|
assert ws.receive_json() == {"type": "session.done", "utterance_index": 0, "total_sentences": 0}
|
|
|
|
assert speech_service._generate_audio_bytes.await_count == 0
|
|
|
|
def test_multiple_sentences_are_buffered_into_one_request(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": "First sentence. "})
|
|
ws.send_json({"type": "input.text", "text": "Second sentence. "})
|
|
ws.send_json({"type": "input.done"})
|
|
|
|
start = ws.receive_json()
|
|
assert start["sentence_index"] == 0
|
|
assert start["sentence_text"] == "First sentence. Second sentence."
|
|
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", "utterance_index": 0, "total_sentences": 1}
|
|
|
|
def test_unknown_message_type_keeps_session_open(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": "unknown"})
|
|
|
|
error = ws.receive_json()
|
|
assert error == {"type": "error", "message": "Unknown message type: unknown"}
|
|
|
|
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",
|
|
"utterance_index": 0,
|
|
"sentence_index": 0,
|
|
"total_bytes": 36,
|
|
"error": False,
|
|
}
|
|
|
|
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)
|
|
|
|
with TestClient(app) as client:
|
|
with client.websocket_connect("/v1/audio/speech/stream") as ws:
|
|
error = ws.receive_json()
|
|
assert error == {"type": "error", "message": "Timeout waiting for session.config"}
|
|
|
|
def test_generation_error_marks_audio_done(self, mocker: MockerFixture):
|
|
speech_service = mocker.MagicMock(spec=OmniOpenAIServingSpeech)
|
|
speech_service._generate_audio_bytes = mocker.AsyncMock(side_effect=RuntimeError("boom"))
|
|
speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req-err", object(), {}))
|
|
speech_service._generate_pcm_chunks = mocker.AsyncMock()
|
|
speech_service.engine_client = mocker.MagicMock()
|
|
speech_service.engine_client.abort = mocker.AsyncMock()
|
|
speech_service.forced_aligner_enabled = False
|
|
app, _ = _build_test_app(speech_service)
|
|
|
|
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"
|
|
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", "utterance_index": 0, "total_sentences": 1}
|
|
|
|
def test_streaming_generation_error_marks_audio_done(self, mocker: MockerFixture):
|
|
speech_service = mocker.MagicMock(spec=OmniOpenAIServingSpeech)
|
|
speech_service._generate_audio_bytes = mocker.AsyncMock(return_value=(b"", "audio/wav"))
|
|
speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req-stream-err", object(), {}))
|
|
speech_service.engine_client = mocker.MagicMock()
|
|
speech_service.engine_client.abort = mocker.AsyncMock()
|
|
speech_service.forced_aligner_enabled = False
|
|
|
|
async def mock_generate_pcm_chunks(
|
|
_generator,
|
|
_request_id,
|
|
*,
|
|
request_start_s=None,
|
|
request_arrival_ts=None,
|
|
include_sample_rate=False,
|
|
tts_params=None,
|
|
):
|
|
yield b"\x01\x02"
|
|
raise RuntimeError("stream boom")
|
|
|
|
speech_service._generate_pcm_chunks = mock_generate_pcm_chunks
|
|
app, _ = _build_test_app(speech_service)
|
|
|
|
with TestClient(app) as client:
|
|
with client.websocket_connect("/v1/audio/speech/stream") as ws:
|
|
ws.send_json(
|
|
{
|
|
"type": "session.config",
|
|
"voice": "Vivian",
|
|
"stream_audio": True,
|
|
"response_format": "pcm",
|
|
}
|
|
)
|
|
ws.send_json({"type": "input.text", "text": "Hello world. "})
|
|
ws.send_json({"type": "input.done"})
|
|
|
|
assert ws.receive_json()["type"] == "audio.start"
|
|
assert ws.receive_bytes() == b"\x01\x02"
|
|
assert ws.receive_json() == {
|
|
"type": "error",
|
|
"message": "Generation failed for utterance 0, sentence 0: stream boom",
|
|
"partial_audio": True,
|
|
"action": "discard",
|
|
}
|
|
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", "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)
|
|
|
|
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": 123})
|
|
|
|
assert ws.receive_json() == {
|
|
"type": "error",
|
|
"message": "input.text requires a string value",
|
|
}
|
|
|
|
ws.send_json({"type": "input.done"})
|
|
assert ws.receive_json() == {"type": "session.done", "utterance_index": 0, "total_sentences": 0}
|
|
|
|
assert speech_service._generate_audio_bytes.await_count == 0
|
|
|
|
def test_input_text_message_too_large(self, monkeypatch, mocker: MockerFixture):
|
|
monkeypatch.setattr(streaming_speech_module, "_MAX_INPUT_TEXT_MESSAGE_SIZE", 32)
|
|
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": "x" * 128})
|
|
|
|
assert ws.receive_json() == {
|
|
"type": "error",
|
|
"message": "input.text message too large",
|
|
}
|
|
|
|
ws.send_json({"type": "input.done"})
|
|
assert ws.receive_json() == {"type": "session.done", "utterance_index": 0, "total_sentences": 0}
|
|
|
|
assert speech_service._generate_audio_bytes.await_count == 0
|
|
|
|
def test_session_config_message_too_large(self, monkeypatch, mocker: MockerFixture):
|
|
monkeypatch.setattr(streaming_speech_module, "_MAX_CONFIG_MESSAGE_SIZE", 64)
|
|
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", "ref_audio": "x" * 512})
|
|
|
|
assert ws.receive_json() == {
|
|
"type": "error",
|
|
"message": "session.config message too large",
|
|
}
|
|
|
|
def test_disconnect_aborts_streaming_request(self, mocker: MockerFixture):
|
|
speech_service = mocker.MagicMock(spec=OmniOpenAIServingSpeech)
|
|
speech_service._generate_audio_bytes = mocker.AsyncMock(return_value=(b"", "audio/wav"))
|
|
speech_service._prepare_speech_generation = mocker.AsyncMock(return_value=("req-abort", object(), {}))
|
|
speech_service.engine_client = mocker.MagicMock()
|
|
speech_service.engine_client.abort = mocker.AsyncMock()
|
|
speech_service.forced_aligner_enabled = False
|
|
|
|
async def mock_generate_pcm_chunks(
|
|
_generator,
|
|
_request_id,
|
|
*,
|
|
request_start_s=None,
|
|
request_arrival_ts=None,
|
|
include_sample_rate=False,
|
|
tts_params=None,
|
|
):
|
|
yield b"\x01\x02"
|
|
|
|
speech_service._generate_pcm_chunks = mock_generate_pcm_chunks
|
|
handler = OmniStreamingSpeechHandler(speech_service=speech_service)
|
|
|
|
websocket = mocker.MagicMock()
|
|
websocket.send_json = mocker.AsyncMock(side_effect=[None, WebSocketDisconnect()])
|
|
websocket.send_bytes = mocker.AsyncMock(side_effect=WebSocketDisconnect())
|
|
|
|
config = mocker.MagicMock()
|
|
config.model = None
|
|
config.voice = "Vivian"
|
|
config.task_type = None
|
|
config.language = None
|
|
config.instructions = None
|
|
config.response_format = "pcm"
|
|
config.speed = 1.0
|
|
config.max_new_tokens = None
|
|
config.initial_codec_chunk_frames = None
|
|
config.ref_audio = None
|
|
config.ref_text = None
|
|
config.x_vector_only_mode = None
|
|
config.speaker_embedding = None
|
|
config.stream_audio = True
|
|
config.word_timestamps = False
|
|
config.seed = None
|
|
config.non_streaming_mode = None
|
|
|
|
with pytest.raises(WebSocketDisconnect):
|
|
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
|
|
|
|
|
|
class TestWebSocketSentenceSplitting:
|
|
def test_sentence_granularity_emits_one_request_per_sentence(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",
|
|
"split_granularity": "sentence",
|
|
}
|
|
)
|
|
ws.send_json({"type": "input.text", "text": "Hello world. How are you? "})
|
|
ws.send_json({"type": "input.done"})
|
|
|
|
first = ws.receive_json()
|
|
assert first["sentence_index"] == 0
|
|
assert first["sentence_text"] == "Hello world."
|
|
ws.receive_bytes()
|
|
assert ws.receive_json()["type"] == "audio.done"
|
|
|
|
second = ws.receive_json()
|
|
assert second["sentence_index"] == 1
|
|
assert second["sentence_text"] == "How are you?"
|
|
ws.receive_bytes()
|
|
assert ws.receive_json()["type"] == "audio.done"
|
|
|
|
assert ws.receive_json() == {
|
|
"type": "session.done",
|
|
"utterance_index": 0,
|
|
"total_sentences": 2,
|
|
}
|
|
|
|
assert speech_service._generate_audio_bytes.await_count == 2
|
|
assert [call.args[0].input for call in speech_service._generate_audio_bytes.await_args_list] == [
|
|
"Hello world.",
|
|
"How are you?",
|
|
]
|
|
|
|
def test_sentence_granularity_emits_before_input_done(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",
|
|
"split_granularity": "sentence",
|
|
}
|
|
)
|
|
ws.send_json({"type": "input.text", "text": "First sentence. "})
|
|
|
|
start = ws.receive_json()
|
|
assert start["sentence_text"] == "First sentence."
|
|
ws.receive_bytes()
|
|
assert ws.receive_json()["type"] == "audio.done"
|
|
assert speech_service._generate_audio_bytes.await_count == 1
|
|
|
|
ws.send_json({"type": "input.done"})
|
|
assert ws.receive_json() == {
|
|
"type": "session.done",
|
|
"utterance_index": 0,
|
|
"total_sentences": 1,
|
|
}
|
|
|
|
def test_indic_danda_splits_without_latin_period(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",
|
|
"language": "Auto",
|
|
"split_granularity": "sentence",
|
|
}
|
|
)
|
|
ws.send_json({"type": "input.text", "text": "नमस्ते। कैसे हो?"})
|
|
ws.send_json({"type": "input.done"})
|
|
|
|
first = ws.receive_json()
|
|
assert first["sentence_text"] == "नमस्ते।"
|
|
ws.receive_bytes()
|
|
assert ws.receive_json()["type"] == "audio.done"
|
|
|
|
second = ws.receive_json()
|
|
assert second["sentence_text"] == "कैसे हो?"
|
|
ws.receive_bytes()
|
|
assert ws.receive_json()["type"] == "audio.done"
|
|
assert ws.receive_json()["total_sentences"] == 2
|
|
|
|
def test_session_config_rejected_after_a_split_unit_was_emitted(self, mocker: MockerFixture):
|
|
"""The splitter buffer is empty here, but the utterance is still open."""
|
|
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",
|
|
"split_granularity": "sentence",
|
|
}
|
|
)
|
|
ws.send_json({"type": "input.text", "text": "First sentence. "})
|
|
ws.receive_json()
|
|
ws.receive_bytes()
|
|
assert ws.receive_json()["type"] == "audio.done"
|
|
|
|
ws.send_json({"type": "session.config", "voice": "Serena"})
|
|
error = ws.receive_json()
|
|
assert error["type"] == "error"
|
|
assert "while an utterance is in progress" in error["message"]
|
|
|
|
ws.send_json({"type": "input.text", "text": "Second sentence. "})
|
|
assert ws.receive_json()["sentence_index"] == 1
|
|
ws.receive_bytes()
|
|
assert ws.receive_json()["type"] == "audio.done"
|
|
|
|
ws.send_json({"type": "input.done"})
|
|
assert ws.receive_json() == {
|
|
"type": "session.done",
|
|
"utterance_index": 0,
|
|
"total_sentences": 2,
|
|
}
|
|
|
|
# Reconfiguration is allowed again once the utterance closed.
|
|
ws.send_json({"type": "session.config", "voice": "Serena"})
|
|
ws.send_json({"type": "input.text", "text": "Third."})
|
|
ws.send_json({"type": "input.done"})
|
|
ws.receive_json()
|
|
ws.receive_bytes()
|
|
ws.receive_json()
|
|
ws.receive_json()
|
|
|
|
voices = [call.args[0].voice for call in speech_service._generate_audio_bytes.await_args_list]
|
|
assert voices == ["Vivian", "Vivian", "Serena"]
|
|
|
|
def test_seed_is_forwarded_to_speech_request(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", "seed": 42})
|
|
ws.send_json({"type": "input.text", "text": "Hello."})
|
|
ws.send_json({"type": "input.done"})
|
|
ws.receive_json()
|
|
ws.receive_bytes()
|
|
ws.receive_json()
|
|
ws.receive_json()
|
|
|
|
assert speech_service._generate_audio_bytes.await_args_list[0].args[0].seed == 42
|
|
|
|
|
|
class TestGeneratePcmChunksContract:
|
|
"""Guard: _generate_pcm_chunks must exist on OmniOpenAIServingSpeech.
|
|
|
|
The WebSocket handler calls speech_service._generate_pcm_chunks()
|
|
at runtime. If the method is removed, all WS TTS streaming breaks
|
|
with an AttributeError. This test catches that at CI time.
|
|
"""
|
|
|
|
def test_generate_pcm_chunks_defined(self):
|
|
assert hasattr(OmniOpenAIServingSpeech, "_generate_pcm_chunks")
|
|
assert asyncio.iscoroutinefunction(OmniOpenAIServingSpeech._generate_pcm_chunks) or callable(
|
|
OmniOpenAIServingSpeech._generate_pcm_chunks
|
|
)
|