Nick CaoandCodex
65699fe94c
Signed-off-by: Nick Cao <ncao@redhat.com> Co-authored-by: Codex <noreply@openai.com>
290 lines
10 KiB
Python
290 lines
10 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
import io
|
|
import json
|
|
import wave
|
|
from types import SimpleNamespace
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
import pytest
|
|
from vllm.sampling_params import RequestOutputKind, SamplingParams
|
|
|
|
from tests.helpers.serving_chat import build_serving_chat
|
|
from vllm_omni.entrypoints.openai.realtime.connection import (
|
|
SAMPLE_RATE_HZ,
|
|
OpenAIFullDuplexConnection,
|
|
_ResolvedResponse,
|
|
)
|
|
from vllm_omni.entrypoints.openai.realtime.session import ActiveResponse
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
class _FakeRenderer:
|
|
def get_tokenizer(self) -> object:
|
|
return object()
|
|
|
|
def __init__(self) -> None:
|
|
self.raw_prompt_lengths: list[int] = []
|
|
self.tokenization_params: list[Any] = []
|
|
self.engine_inputs: list[dict[str, Any]] = []
|
|
self.conversations: list[list[dict[str, Any]]] = []
|
|
|
|
async def render_chat_async(
|
|
self,
|
|
conversations: list[list[dict[str, Any]]],
|
|
_chat_params: Any,
|
|
tok_params: Any,
|
|
):
|
|
self.conversations.append(conversations[0])
|
|
raw_prompt_length = 2 * len(conversations[0])
|
|
self.raw_prompt_lengths.append(raw_prompt_length)
|
|
self.tokenization_params.append(tok_params)
|
|
# Model-side expansion makes the rendered prompt much longer than its text tokens.
|
|
engine_input = {"prompt_token_ids": [0] * (10 + 20 * raw_prompt_length)}
|
|
self.engine_inputs.append(engine_input)
|
|
return conversations, [engine_input]
|
|
|
|
|
|
class _FakeWebSocket:
|
|
def __init__(self) -> None:
|
|
self.messages: list[str] = []
|
|
|
|
async def send_text(self, message: str) -> None:
|
|
self.messages.append(message)
|
|
|
|
|
|
def _make_connection(
|
|
*,
|
|
max_model_len: int = 100,
|
|
websocket: Any = None,
|
|
engine: Any = None,
|
|
chat_handler: Any = None,
|
|
) -> tuple[OpenAIFullDuplexConnection, _FakeRenderer]:
|
|
renderer = _FakeRenderer()
|
|
model_config = SimpleNamespace(max_model_len=max_model_len, multimodal_config=None)
|
|
|
|
async def preprocess_chat(request: Any, messages: list[dict[str, Any]], **kwargs: Any):
|
|
tok_params = kwargs.get("tok_params")
|
|
if tok_params is None:
|
|
tok_params = request.build_tok_params(model_config)
|
|
return await renderer.render_chat_async(
|
|
[messages],
|
|
kwargs.get("default_template_kwargs"),
|
|
tok_params,
|
|
)
|
|
|
|
connection = OpenAIFullDuplexConnection(
|
|
websocket=websocket,
|
|
engine=engine if engine is not None else SimpleNamespace(model_config=model_config),
|
|
model_name="test-model",
|
|
chat_handler=chat_handler
|
|
if chat_handler is not None
|
|
else SimpleNamespace(
|
|
renderer=renderer,
|
|
chat_template=None,
|
|
chat_template_content_format="auto",
|
|
_effective_chat_template_kwargs=lambda _request: {},
|
|
_preprocess_chat=preprocess_chat,
|
|
_fix_minicpmo45_audio_stream_output_kinds=lambda params, _modalities: params,
|
|
),
|
|
)
|
|
return connection, renderer
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_preflight_truncates_by_rendered_token_count() -> None:
|
|
connection, renderer = _make_connection()
|
|
|
|
first = SimpleNamespace(
|
|
id="first",
|
|
type="message",
|
|
role="user",
|
|
content=[SimpleNamespace(type="input_text", text="first")],
|
|
)
|
|
second = SimpleNamespace(
|
|
id="second",
|
|
type="message",
|
|
role="user",
|
|
content=[SimpleNamespace(type="input_text", text="second")],
|
|
)
|
|
items = [first, second]
|
|
response = _ResolvedResponse(
|
|
input=items,
|
|
instructions=None,
|
|
modalities=["text"],
|
|
max_output_tokens="inf",
|
|
tools=None,
|
|
tool_choice="none",
|
|
metadata=None,
|
|
)
|
|
|
|
engine_input = await connection._prepare_engine_input_with_auto_truncation(response)
|
|
|
|
assert renderer.raw_prompt_lengths == [4, 2, 2]
|
|
assert [len(prompt["prompt_token_ids"]) for prompt in renderer.engine_inputs] == [90, 50, 50]
|
|
assert all(params.max_total_tokens is None for params in renderer.tokenization_params)
|
|
assert all(params.max_output_tokens == 0 for params in renderer.tokenization_params)
|
|
assert items == [second]
|
|
assert engine_input is renderer.engine_inputs[-1]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_audio_prompt_uses_standard_chat_audio_content() -> None:
|
|
audio = b"\x00\x00\x00\x40"
|
|
connection, renderer = _make_connection(websocket=_FakeWebSocket())
|
|
connection.session.input_audio_buffer.extend(audio)
|
|
item = await connection._commit_audio_buffer()
|
|
|
|
assert item is not None
|
|
assert base64.b64decode(item.content[0].audio) == audio
|
|
|
|
await connection._build_full_prompt(items=[item])
|
|
|
|
audio_content = renderer.conversations[0][0]["content"][0]["input_audio"]
|
|
assert audio_content["format"] == "wav"
|
|
with wave.open(io.BytesIO(base64.b64decode(audio_content["data"]))) as wav:
|
|
assert wav.getframerate() == SAMPLE_RATE_HZ
|
|
assert wav.getnchannels() == 1
|
|
assert wav.getsampwidth() == 2
|
|
assert wav.readframes(2) == audio
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_generated_audio_is_resampled_to_realtime_rate() -> None:
|
|
source_rate = 32_000
|
|
waveform = np.sin(2 * np.pi * 440 * np.arange(source_rate) / source_rate).astype(np.float32)
|
|
outputs: list[Any] = []
|
|
for chunk in np.split(waveform, 2):
|
|
multimodal_output = {"audio": chunk, "sample_rate": source_rate}
|
|
outputs.append(
|
|
SimpleNamespace(
|
|
final_output_type="audio",
|
|
multimodal_output=multimodal_output,
|
|
outputs=[SimpleNamespace(multimodal_output=multimodal_output)],
|
|
)
|
|
)
|
|
|
|
async def generate(**_kwargs):
|
|
for output in outputs:
|
|
yield output
|
|
|
|
model_config = SimpleNamespace(max_model_len=100, multimodal_config=None)
|
|
engine = SimpleNamespace(
|
|
model_config=model_config,
|
|
default_sampling_params_list=[],
|
|
generate=generate,
|
|
)
|
|
websocket = _FakeWebSocket()
|
|
connection, _ = _make_connection(websocket=websocket, engine=engine)
|
|
response = _ResolvedResponse(
|
|
input=[],
|
|
instructions=None,
|
|
modalities=["audio"],
|
|
max_output_tokens="inf",
|
|
tools=None,
|
|
tool_choice="none",
|
|
metadata=None,
|
|
)
|
|
active = ActiveResponse(response_id="resp_test", request_id="req_test")
|
|
connection.session.active_response = active
|
|
|
|
await connection._run_response(active.response_id, response, {"prompt_token_ids": []})
|
|
|
|
events = [json.loads(message) for message in websocket.messages]
|
|
audio_deltas = [event for event in events if event["type"] == "response.output_audio.delta"]
|
|
pcm = b"".join(base64.b64decode(event["delta"]) for event in audio_deltas)
|
|
assert active.item_id is not None
|
|
assert len(pcm) == SAMPLE_RATE_HZ * 2
|
|
assert connection.session.item_duration_ms[active.item_id] == pytest.approx(1000)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_minicpmo45_realtime_keeps_thinker_final_only() -> None:
|
|
model_arch = "MiniCPMO45OmniForConditionalGeneration"
|
|
stages = [
|
|
SimpleNamespace(engine_args=SimpleNamespace(model_arch=model_arch, model_stage="llm")),
|
|
SimpleNamespace(engine_args=SimpleNamespace(model_arch=model_arch, model_stage="tts")),
|
|
]
|
|
submitted_params: list[list[SamplingParams]] = []
|
|
|
|
async def generate(**kwargs):
|
|
submitted_params.append(kwargs["sampling_params_list"])
|
|
yield SimpleNamespace(final_output_type="text", outputs=[])
|
|
|
|
model_config = SimpleNamespace(max_model_len=100, multimodal_config=None)
|
|
chat_handler = build_serving_chat()
|
|
engine = chat_handler.engine_client
|
|
engine.model_config = model_config
|
|
engine.stage_configs = stages
|
|
engine.default_sampling_params_list = [SamplingParams(), SamplingParams()]
|
|
engine.generate = generate
|
|
websocket = _FakeWebSocket()
|
|
connection, _ = _make_connection(
|
|
websocket=websocket,
|
|
engine=engine,
|
|
chat_handler=chat_handler,
|
|
)
|
|
response = _ResolvedResponse(
|
|
input=[],
|
|
instructions=None,
|
|
modalities=["audio"],
|
|
max_output_tokens="inf",
|
|
tools=None,
|
|
tool_choice="none",
|
|
metadata=None,
|
|
)
|
|
active = ActiveResponse(response_id="resp_test", request_id="req_test")
|
|
connection.session.active_response = active
|
|
|
|
await connection._run_response(active.response_id, response, {"prompt_token_ids": []})
|
|
|
|
assert len(submitted_params) == 1
|
|
assert [params.output_kind for params in submitted_params[0]] == [
|
|
RequestOutputKind.FINAL_ONLY,
|
|
RequestOutputKind.DELTA,
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("audio_end_ms", "expected"),
|
|
[(0, ""), (250, "0"), (500, "01"), (999, "012"), (1000, "0123")],
|
|
)
|
|
def test_truncate_transcript_uses_item_audio_duration(audio_end_ms: float, expected: str) -> None:
|
|
connection, _ = _make_connection()
|
|
connection._tokenizer = SimpleNamespace(decode=lambda ids, **_kwargs: "".join(map(str, ids)))
|
|
connection.session.item_duration_ms["item"] = 1000
|
|
connection.session.item_token_ids["item"] = [0, 1, 2, 3]
|
|
|
|
assert connection._truncate_transcript("item", audio_end_ms) == expected
|
|
|
|
|
|
def test_truncate_transcript_falls_back_without_tokens_or_duration() -> None:
|
|
connection, _ = _make_connection()
|
|
connection._tokenizer = SimpleNamespace(decode=lambda ids, **_kwargs: "".join(map(str, ids)))
|
|
connection.session.item_token_ids["missing_duration"] = [0, 1, 2, 3]
|
|
connection.session.item_duration_ms["missing_tokens"] = 1000
|
|
|
|
assert connection._truncate_transcript("missing_duration", 500) == ""
|
|
assert connection._truncate_transcript("missing_tokens", 500) == ""
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_rejected_response_create_does_not_cancel_active_response() -> None:
|
|
websocket = _FakeWebSocket()
|
|
connection, _ = _make_connection(max_model_len=5, websocket=websocket)
|
|
active_response = SimpleNamespace(response_id="active", request_id="active-request")
|
|
connection.session.active_response = active_response
|
|
|
|
await connection._handle_response_create(SimpleNamespace(event_id="evt_create", response=None))
|
|
|
|
assert connection.session.active_response is active_response
|
|
assert not connection._response_cancel_event.is_set()
|
|
assert len(websocket.messages) == 1
|
|
assert "exceeds the model's input token limit" in websocket.messages[0]
|