[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
@@ -1,19 +1,23 @@
|
||||
# tests/entrypoints/openai/test_serving_speech.py
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
from inspect import Signature, signature
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock, patch
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from fastapi import FastAPI, HTTPException, UploadFile
|
||||
from fastapi import FastAPI, HTTPException, Request, UploadFile
|
||||
from fastapi.params import File, Form
|
||||
from fastapi.responses import JSONResponse
|
||||
from fastapi.testclient import TestClient
|
||||
from pydantic import ValidationError
|
||||
from pytest_mock import MockerFixture
|
||||
from vllm.entrypoints.openai.engine.protocol import ErrorInfo, ErrorResponse
|
||||
|
||||
from vllm_omni.entrypoints.openai import api_server as api_server_module
|
||||
from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin
|
||||
from vllm_omni.entrypoints.openai.protocol.audio import CreateAudio, OpenAICreateSpeechRequest
|
||||
from vllm_omni.entrypoints.openai.serving_speech import (
|
||||
@@ -1052,4 +1056,33 @@ class TestAsyncOmniSupportedTasks:
|
||||
omni.stage_list = [stage]
|
||||
tasks = await omni.get_supported_tasks()
|
||||
assert "generate" in tasks
|
||||
assert "speech" in tasks
|
||||
|
||||
|
||||
def test_api_server_create_speech_wraps_error_response_status():
|
||||
handler = MagicMock()
|
||||
handler.create_speech = AsyncMock(
|
||||
return_value=ErrorResponse(
|
||||
error=ErrorInfo(message="bad request", type="BadRequestError", param=None, code=400),
|
||||
)
|
||||
)
|
||||
|
||||
app = FastAPI()
|
||||
app.state.openai_serving_speech = handler
|
||||
scope = {
|
||||
"type": "http",
|
||||
"app": app,
|
||||
"method": "POST",
|
||||
"path": "/v1/audio/speech",
|
||||
"headers": [],
|
||||
"query_string": b"",
|
||||
"client": ("127.0.0.1", 12345),
|
||||
"server": ("testserver", 80),
|
||||
"scheme": "http",
|
||||
}
|
||||
raw_request = Request(scope)
|
||||
request = OpenAICreateSpeechRequest(input="Hello")
|
||||
|
||||
response = asyncio.run(api_server_module.create_speech(request, raw_request))
|
||||
|
||||
assert isinstance(response, JSONResponse)
|
||||
assert response.status_code == 400
|
||||
@@ -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
|
||||
@@ -0,0 +1,253 @@
|
||||
"""Tests for SentenceSplitter used in streaming TTS input."""
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.entrypoints.openai.text_splitter import SentenceSplitter
|
||||
|
||||
pytestmark = [pytest.mark.openai, pytest.mark.speech]
|
||||
|
||||
|
||||
class TestSentenceSplitterEnglish:
|
||||
"""Tests for English sentence splitting."""
|
||||
|
||||
def test_single_sentence_no_boundary(self):
|
||||
splitter = SentenceSplitter()
|
||||
result = splitter.add_text("Hello world")
|
||||
assert result == []
|
||||
assert splitter.buffer == "Hello world"
|
||||
|
||||
def test_single_sentence_with_boundary(self):
|
||||
splitter = SentenceSplitter()
|
||||
result = splitter.add_text("Hello world. How are you?")
|
||||
assert len(result) == 1
|
||||
assert result[0] == "Hello world."
|
||||
|
||||
def test_multiple_sentences(self):
|
||||
splitter = SentenceSplitter()
|
||||
result = splitter.add_text("Hello. How are you? I am fine! ")
|
||||
assert len(result) == 3
|
||||
assert result[0] == "Hello."
|
||||
assert result[1] == "How are you?"
|
||||
assert result[2] == "I am fine!"
|
||||
|
||||
def test_exclamation_mark(self):
|
||||
splitter = SentenceSplitter()
|
||||
result = splitter.add_text("Wow, that is great! Tell me more.")
|
||||
assert len(result) == 1
|
||||
assert result[0] == "Wow, that is great!"
|
||||
|
||||
def test_question_mark(self):
|
||||
splitter = SentenceSplitter()
|
||||
result = splitter.add_text("Can you hear me? I hope so.")
|
||||
assert len(result) == 1
|
||||
assert result[0] == "Can you hear me?"
|
||||
|
||||
|
||||
class TestSentenceSplitterChinese:
|
||||
"""Tests for CJK sentence splitting."""
|
||||
|
||||
def test_chinese_period(self):
|
||||
splitter = SentenceSplitter()
|
||||
result = splitter.add_text("你好世界。你好吗")
|
||||
assert len(result) == 1
|
||||
assert result[0] == "你好世界。"
|
||||
|
||||
def test_chinese_exclamation(self):
|
||||
splitter = SentenceSplitter()
|
||||
result = splitter.add_text("太好了！谢谢你")
|
||||
assert len(result) == 1
|
||||
assert result[0] == "太好了！"
|
||||
|
||||
def test_chinese_question(self):
|
||||
splitter = SentenceSplitter()
|
||||
result = splitter.add_text("你是谁？我是小明")
|
||||
assert len(result) == 1
|
||||
assert result[0] == "你是谁？"
|
||||
|
||||
def test_chinese_comma_no_split(self):
|
||||
"""Chinese commas are clause-level and should not trigger a split."""
|
||||
splitter = SentenceSplitter()
|
||||
result = splitter.add_text("你好，世界")
|
||||
assert result == []
|
||||
assert splitter.buffer == "你好，世界"
|
||||
|
||||
def test_chinese_semicolon_no_split(self):
|
||||
"""Chinese semicolons are clause-level and should not trigger a split."""
|
||||
splitter = SentenceSplitter()
|
||||
result = splitter.add_text("第一点；第二点")
|
||||
assert result == []
|
||||
assert splitter.buffer == "第一点；第二点"
|
||||
|
||||
def test_chinese_multiple(self):
|
||||
splitter = SentenceSplitter()
|
||||
result = splitter.add_text("你好！你好吗？我很好。")
|
||||
assert len(result) == 3
|
||||
assert result[0] == "你好！"
|
||||
assert result[1] == "你好吗？"
|
||||
assert result[2] == "我很好。"
|
||||
|
||||
|
||||
class TestSentenceSplitterMixed:
|
||||
"""Tests for mixed-language sentence splitting."""
|
||||
|
||||
def test_mixed_english_chinese(self):
|
||||
splitter = SentenceSplitter()
|
||||
result = splitter.add_text("Hello世界。How are you? ")
|
||||
assert len(result) == 2
|
||||
assert result[0] == "Hello世界。"
|
||||
assert result[1] == "How are you?"
|
||||
|
||||
|
||||
class TestSentenceSplitterIncremental:
|
||||
"""Tests for incremental (multi-chunk) text input."""
|
||||
|
||||
def test_accumulation_across_chunks(self):
|
||||
splitter = SentenceSplitter()
|
||||
# First chunk: no boundary
|
||||
result1 = splitter.add_text("Hello ")
|
||||
assert result1 == []
|
||||
|
||||
# Second chunk: completes a sentence
|
||||
result2 = splitter.add_text("world. How")
|
||||
assert len(result2) == 1
|
||||
assert result2[0] == "Hello world."
|
||||
assert splitter.buffer == "How"
|
||||
|
||||
def test_word_by_word(self):
|
||||
splitter = SentenceSplitter()
|
||||
words = ["Hello, ", "how ", "are ", "you? ", "I ", "am ", "fine."]
|
||||
all_sentences = []
|
||||
for word in words:
|
||||
all_sentences.extend(splitter.add_text(word))
|
||||
|
||||
assert len(all_sentences) == 1
|
||||
assert all_sentences[0] == "Hello, how are you?"
|
||||
# "I am fine." stays in buffer (no trailing whitespace after period)
|
||||
|
||||
def test_three_chunks(self):
|
||||
splitter = SentenceSplitter()
|
||||
splitter.add_text("The quick brown ")
|
||||
splitter.add_text("fox jumps. ")
|
||||
result = splitter.add_text("Over the lazy dog. ")
|
||||
# "The quick brown fox jumps." should have been returned on second chunk
|
||||
# "Over the lazy dog." on third chunk
|
||||
assert len(result) == 1
|
||||
assert result[0] == "Over the lazy dog."
|
||||
|
||||
|
||||
class TestSentenceSplitterFlush:
|
||||
"""Tests for flush behavior."""
|
||||
|
||||
def test_flush_returns_remaining(self):
|
||||
splitter = SentenceSplitter()
|
||||
splitter.add_text("Hello world")
|
||||
result = splitter.flush()
|
||||
assert result == "Hello world"
|
||||
assert splitter.buffer == ""
|
||||
|
||||
def test_flush_empty_buffer(self):
|
||||
splitter = SentenceSplitter()
|
||||
result = splitter.flush()
|
||||
assert result is None
|
||||
|
||||
def test_flush_after_sentence(self):
|
||||
splitter = SentenceSplitter()
|
||||
splitter.add_text("Hello world. Remaining text")
|
||||
result = splitter.flush()
|
||||
assert result == "Remaining text"
|
||||
|
||||
def test_flush_whitespace_only(self):
|
||||
splitter = SentenceSplitter()
|
||||
splitter.add_text("Hello. ")
|
||||
# "Hello." extracted, buffer is " "
|
||||
result = splitter.flush()
|
||||
# Whitespace-only should return None
|
||||
assert result is None
|
||||
|
||||
def test_flush_clears_buffer(self):
|
||||
splitter = SentenceSplitter()
|
||||
splitter.add_text("some text")
|
||||
splitter.flush()
|
||||
assert splitter.buffer == ""
|
||||
# Second flush should return None
|
||||
assert splitter.flush() is None
|
||||
|
||||
|
||||
class TestSentenceSplitterEdgeCases:
|
||||
"""Edge case tests."""
|
||||
|
||||
def test_empty_input(self):
|
||||
splitter = SentenceSplitter()
|
||||
result = splitter.add_text("")
|
||||
assert result == []
|
||||
assert splitter.buffer == ""
|
||||
|
||||
def test_none_like_empty(self):
|
||||
"""Empty string should not affect buffer."""
|
||||
splitter = SentenceSplitter()
|
||||
splitter.add_text("Hello")
|
||||
splitter.add_text("")
|
||||
assert splitter.buffer == "Hello"
|
||||
|
||||
def test_only_punctuation(self):
|
||||
splitter = SentenceSplitter()
|
||||
result = splitter.add_text(". ")
|
||||
# "." is 1 char, below default min_sentence_length of 2
|
||||
# It will be carried forward
|
||||
assert result == []
|
||||
|
||||
def test_min_sentence_length(self):
|
||||
splitter = SentenceSplitter(min_sentence_length=10)
|
||||
result = splitter.add_text("Hi. Hello world. ")
|
||||
# "Hi." is 3 chars (< 10), so it gets carried to "Hello world."
|
||||
assert len(result) == 1
|
||||
assert "Hi." in result[0]
|
||||
assert "Hello world." in result[0]
|
||||
|
||||
def test_short_segments_are_carried_until_long_enough(self):
|
||||
splitter = SentenceSplitter(min_sentence_length=10)
|
||||
result = splitter.add_text("Hi. Ok. Hello there. ")
|
||||
assert result == ["Hi.Ok.Hello there."]
|
||||
assert splitter.buffer == ""
|
||||
|
||||
def test_min_sentence_length_zero(self):
|
||||
splitter = SentenceSplitter(min_sentence_length=0)
|
||||
result = splitter.add_text("A. B. ")
|
||||
assert len(result) == 2
|
||||
|
||||
def test_no_boundary_then_flush(self):
|
||||
splitter = SentenceSplitter()
|
||||
result = splitter.add_text("Hello world how are you")
|
||||
assert result == []
|
||||
flushed = splitter.flush()
|
||||
assert flushed == "Hello world how are you"
|
||||
|
||||
def test_consecutive_punctuation(self):
|
||||
splitter = SentenceSplitter()
|
||||
result = splitter.add_text("Really?! Yes, really. ")
|
||||
assert len(result) >= 1
|
||||
|
||||
def test_reuse_after_flush(self):
|
||||
"""Splitter can be reused after flush."""
|
||||
splitter = SentenceSplitter()
|
||||
splitter.add_text("First session.")
|
||||
splitter.flush()
|
||||
|
||||
result = splitter.add_text("Second session. More text")
|
||||
assert len(result) == 1
|
||||
assert result[0] == "Second session."
|
||||
assert splitter.buffer == "More text"
|
||||
|
||||
|
||||
class TestSentenceSplitterBufferLimit:
|
||||
"""Tests for buffer overflow protection."""
|
||||
|
||||
def test_buffer_overflow_raises(self):
|
||||
from vllm_omni.entrypoints.openai.text_splitter import _MAX_BUFFER_SIZE
|
||||
|
||||
splitter = SentenceSplitter()
|
||||
# Fill buffer just under the limit
|
||||
splitter.add_text("x" * (_MAX_BUFFER_SIZE - 1))
|
||||
# One more char should exceed the limit
|
||||
with pytest.raises(ValueError, match="exceeded maximum size"):
|
||||
splitter.add_text("xx")
|
||||
Reference in new issue
Block a user