[Feat][Qwen3-TTS] Support streaming audio output for websocket (#1719)
Signed-off-by: Sy03 <1370724210@qq.com> Signed-off-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> Signed-off-by: linyueqian <linyueqian@outlook.com> Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> Co-authored-by: linyueqian <linyueqian@outlook.com>
This commit is contained in:
12 files changed
+1719
-265
No files matched your search
@@ -0,0 +1,387 @@
|
||||
import asyncio
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
from fastapi import FastAPI, WebSocket
|
||||
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
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _build_test_app(speech_service=None, *, idle_timeout=30.0, config_timeout=10.0):
|
||||
if speech_service is None:
|
||||
speech_service = MagicMock(spec=OmniOpenAIServingSpeech)
|
||||
speech_service._generate_audio_bytes = AsyncMock(return_value=(b"RIFF" + b"\x00" * 32, "audio/wav"))
|
||||
speech_service._prepare_speech_generation = AsyncMock(return_value=("req-1", object(), {}))
|
||||
|
||||
async def mock_generate_pcm_chunks(_generator, _request_id):
|
||||
for chunk in (b"\x01\x02", b"\x03\x04\x05"):
|
||||
yield chunk
|
||||
|
||||
speech_service._generate_pcm_chunks = mock_generate_pcm_chunks
|
||||
speech_service.engine_client = MagicMock()
|
||||
speech_service.engine_client.abort = 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):
|
||||
app, speech_service = _build_test_app()
|
||||
|
||||
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. "})
|
||||
|
||||
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", "sentence_index": 0, "total_bytes": len(audio), "error": False}
|
||||
|
||||
ws.send_json({"type": "input.done"})
|
||||
session_done = ws.receive_json()
|
||||
assert session_done == {"type": "session.done", "total_sentences": 1}
|
||||
|
||||
assert speech_service._generate_audio_bytes.await_count == 1
|
||||
|
||||
def test_streaming_multiple_binary_frames(self):
|
||||
captured_requests = []
|
||||
|
||||
speech_service = MagicMock(spec=OmniOpenAIServingSpeech)
|
||||
speech_service._generate_audio_bytes = AsyncMock(return_value=(b"", "audio/wav"))
|
||||
speech_service.engine_client = MagicMock()
|
||||
speech_service.engine_client.abort = AsyncMock()
|
||||
|
||||
async def mock_prepare_speech_generation(request):
|
||||
captured_requests.append(request)
|
||||
return "req-stream", object(), {}
|
||||
|
||||
speech_service._prepare_speech_generation = mock_prepare_speech_generation
|
||||
|
||||
async def mock_generate_pcm_chunks(_generator, _request_id):
|
||||
for chunk in (b"\x01\x02", b"\x03\x04\x05", b"\x06"):
|
||||
yield 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. "})
|
||||
|
||||
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", "sentence_index": 0, "total_bytes": 6, "error": False}
|
||||
|
||||
ws.send_json({"type": "input.done"})
|
||||
assert ws.receive_json() == {"type": "session.done", "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 speech_service._generate_audio_bytes.await_count == 0
|
||||
|
||||
def test_flush_on_input_done(self):
|
||||
app, _ = _build_test_app()
|
||||
|
||||
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",
|
||||
"sentence_index": 0,
|
||||
"total_bytes": 36,
|
||||
"error": False,
|
||||
}
|
||||
assert ws.receive_json() == {"type": "session.done", "total_sentences": 1}
|
||||
|
||||
def test_invalid_streaming_config(self):
|
||||
app, _ = _build_test_app()
|
||||
|
||||
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):
|
||||
app, speech_service = _build_test_app()
|
||||
|
||||
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", "total_sentences": 0}
|
||||
|
||||
assert speech_service._generate_audio_bytes.await_count == 0
|
||||
|
||||
def test_multiple_sentences_increment_indices(self):
|
||||
app, _ = _build_test_app()
|
||||
|
||||
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. Second sentence. "})
|
||||
|
||||
first_start = ws.receive_json()
|
||||
assert first_start["sentence_index"] == 0
|
||||
ws.receive_bytes()
|
||||
assert ws.receive_json() == {
|
||||
"type": "audio.done",
|
||||
"sentence_index": 0,
|
||||
"total_bytes": 36,
|
||||
"error": False,
|
||||
}
|
||||
|
||||
second_start = ws.receive_json()
|
||||
assert second_start["sentence_index"] == 1
|
||||
ws.receive_bytes()
|
||||
assert ws.receive_json() == {
|
||||
"type": "audio.done",
|
||||
"sentence_index": 1,
|
||||
"total_bytes": 36,
|
||||
"error": False,
|
||||
}
|
||||
|
||||
ws.send_json({"type": "input.done"})
|
||||
assert ws.receive_json() == {"type": "session.done", "total_sentences": 2}
|
||||
|
||||
def test_unknown_message_type_keeps_session_open(self):
|
||||
app, _ = _build_test_app()
|
||||
|
||||
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. "})
|
||||
assert ws.receive_json()["type"] == "audio.start"
|
||||
ws.receive_bytes()
|
||||
assert ws.receive_json() == {
|
||||
"type": "audio.done",
|
||||
"sentence_index": 0,
|
||||
"total_bytes": 36,
|
||||
"error": False,
|
||||
}
|
||||
|
||||
ws.send_json({"type": "input.done"})
|
||||
assert ws.receive_json() == {"type": "session.done", "total_sentences": 1}
|
||||
|
||||
def test_config_timeout_closes_session(self):
|
||||
app, _ = _build_test_app(config_timeout=0.01)
|
||||
|
||||
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):
|
||||
speech_service = MagicMock(spec=OmniOpenAIServingSpeech)
|
||||
speech_service._generate_audio_bytes = AsyncMock(side_effect=RuntimeError("boom"))
|
||||
speech_service._prepare_speech_generation = AsyncMock(return_value=("req-err", object(), {}))
|
||||
speech_service._generate_pcm_chunks = AsyncMock()
|
||||
speech_service.engine_client = MagicMock()
|
||||
speech_service.engine_client.abort = AsyncMock()
|
||||
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. "})
|
||||
|
||||
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}
|
||||
|
||||
ws.send_json({"type": "input.done"})
|
||||
assert ws.receive_json() == {"type": "session.done", "total_sentences": 1}
|
||||
|
||||
def test_streaming_generation_error_marks_audio_done(self):
|
||||
speech_service = MagicMock(spec=OmniOpenAIServingSpeech)
|
||||
speech_service._generate_audio_bytes = AsyncMock(return_value=(b"", "audio/wav"))
|
||||
speech_service._prepare_speech_generation = AsyncMock(return_value=("req-stream-err", object(), {}))
|
||||
speech_service.engine_client = MagicMock()
|
||||
speech_service.engine_client.abort = AsyncMock()
|
||||
|
||||
async def mock_generate_pcm_chunks(_generator, _request_id):
|
||||
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. "})
|
||||
|
||||
assert ws.receive_json()["type"] == "audio.start"
|
||||
assert ws.receive_bytes() == b"\x01\x02"
|
||||
assert ws.receive_json() == {
|
||||
"type": "error",
|
||||
"message": "Generation failed for sentence 0: stream boom",
|
||||
}
|
||||
assert ws.receive_json() == {
|
||||
"type": "audio.done",
|
||||
"sentence_index": 0,
|
||||
"total_bytes": 2,
|
||||
"error": True,
|
||||
}
|
||||
|
||||
ws.send_json({"type": "input.done"})
|
||||
assert ws.receive_json() == {"type": "session.done", "total_sentences": 1}
|
||||
|
||||
def test_invalid_input_text_type_returns_validation_error(self):
|
||||
app, speech_service = _build_test_app()
|
||||
|
||||
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", "total_sentences": 0}
|
||||
|
||||
assert speech_service._generate_audio_bytes.await_count == 0
|
||||
|
||||
def test_input_text_message_too_large(self, monkeypatch):
|
||||
monkeypatch.setattr(streaming_speech_module, "_MAX_INPUT_TEXT_MESSAGE_SIZE", 32)
|
||||
app, speech_service = _build_test_app()
|
||||
|
||||
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", "total_sentences": 0}
|
||||
|
||||
assert speech_service._generate_audio_bytes.await_count == 0
|
||||
|
||||
def test_session_config_message_too_large(self, monkeypatch):
|
||||
monkeypatch.setattr(streaming_speech_module, "_MAX_CONFIG_MESSAGE_SIZE", 64)
|
||||
app, _ = _build_test_app()
|
||||
|
||||
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):
|
||||
speech_service = MagicMock(spec=OmniOpenAIServingSpeech)
|
||||
speech_service._generate_audio_bytes = AsyncMock(return_value=(b"", "audio/wav"))
|
||||
speech_service._prepare_speech_generation = AsyncMock(return_value=("req-abort", object(), {}))
|
||||
speech_service.engine_client = MagicMock()
|
||||
speech_service.engine_client.abort = AsyncMock()
|
||||
|
||||
async def mock_generate_pcm_chunks(_generator, _request_id):
|
||||
yield b"\x01\x02"
|
||||
|
||||
speech_service._generate_pcm_chunks = mock_generate_pcm_chunks
|
||||
handler = OmniStreamingSpeechHandler(speech_service=speech_service)
|
||||
|
||||
websocket = MagicMock()
|
||||
websocket.send_json = AsyncMock(side_effect=[None, WebSocketDisconnect()])
|
||||
websocket.send_bytes = AsyncMock(side_effect=WebSocketDisconnect())
|
||||
|
||||
config = 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.stream_audio = True
|
||||
|
||||
with pytest.raises(WebSocketDisconnect):
|
||||
asyncio.run(handler._generate_and_send(websocket, config, "Hello world.", 0))
|
||||
|
||||
speech_service.engine_client.abort.assert_awaited_once_with("req-abort")
|
||||
assert websocket.send_json.await_count == 2
|
||||
Reference in new issue
Block a user