127 lines
4.5 KiB
Python
127 lines
4.5 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
import pytest
|
|
|
|
from tests.entrypoints.openai_api.test_duplex_handler import FakeChatService, TimedWebSocket
|
|
from vllm_omni.engine.duplex.messages import DuplexFence
|
|
from vllm_omni.entrypoints.async_omni import AsyncOmni
|
|
from vllm_omni.entrypoints.duplex.protocol import DuplexSession, DuplexSessionConfig
|
|
from vllm_omni.entrypoints.duplex.serving import OmniDuplexSessionHandler
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
class FenceRecordingEngine:
|
|
output_modalities = ["text", "audio"]
|
|
|
|
def __init__(self) -> None:
|
|
self.calls: list[tuple[str, DuplexFence | None]] = []
|
|
self.signal_next_fences: list[DuplexFence | None] = []
|
|
|
|
async def open_duplex_session_async(self, session_id, *, fence=None, **kwargs):
|
|
del session_id, kwargs
|
|
self.calls.append(("open", fence))
|
|
return {"ok": True}
|
|
|
|
async def append_duplex_input_async(self, session_id, *, fence=None, **kwargs):
|
|
del session_id, kwargs
|
|
self.calls.append(("append", fence))
|
|
return {"ok": True}
|
|
|
|
async def signal_duplex_turn_async(
|
|
self,
|
|
session_id,
|
|
*,
|
|
event,
|
|
fence=None,
|
|
next_fence=None,
|
|
timeout=None,
|
|
):
|
|
del session_id, event, timeout
|
|
self.calls.append(("signal", fence))
|
|
self.signal_next_fences.append(next_fence)
|
|
return {"ok": True}
|
|
|
|
async def close_duplex_session_async(self, session_id, *, fence=None, **kwargs):
|
|
del session_id, kwargs
|
|
self.calls.append(("close", fence))
|
|
return {"ok": True}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_omni_forwards_fence_for_all_runtime_operations():
|
|
engine = FenceRecordingEngine()
|
|
app = object.__new__(AsyncOmni)
|
|
app.engine = engine
|
|
app.request_states = {}
|
|
app._final_output_handler = lambda: None
|
|
fence = DuplexFence("sid", epoch=1, turn_id=2)
|
|
|
|
await app.open_duplex_session_async("sid", fence=fence)
|
|
await app.append_duplex_input_async(
|
|
"sid", mode="append_audio_chunk", payload={}, fence=fence, collect_outputs=False
|
|
)
|
|
await app.signal_duplex_turn_async("sid", event="turn.end", fence=fence)
|
|
await app.close_duplex_session_async("sid", reason="done", fence=fence)
|
|
|
|
assert engine.calls == [(operation, fence) for operation in ("open", "append", "signal", "close")]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_openai_handler_passes_current_session_fence_for_runtime_controls():
|
|
engine = FenceRecordingEngine()
|
|
handler = OmniDuplexSessionHandler(chat_service=FakeChatService(engine))
|
|
session = DuplexSession(
|
|
session_id="sid",
|
|
config=DuplexSessionConfig(extra_body={"native_duplex": True}),
|
|
epoch=1,
|
|
turn_id=2,
|
|
)
|
|
ws = TimedWebSocket()
|
|
fence = DuplexFence("sid", epoch=1, turn_id=2)
|
|
|
|
await handler._open_runtime_session(session, ws.send_json)
|
|
await handler._signal_runtime_session(session, "turn.end", send_json=ws.send_json)
|
|
await handler._close_runtime_session(session, reason="done", send_json=ws.send_json)
|
|
|
|
assert engine.calls == [(operation, fence) for operation in ("open", "signal", "close")]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_runtime_signal_uses_fence_captured_before_session_turn_advances():
|
|
engine = FenceRecordingEngine()
|
|
handler = OmniDuplexSessionHandler(chat_service=FakeChatService(engine))
|
|
session = DuplexSession(session_id="sid", config=DuplexSessionConfig(extra_body={"native_duplex": True}))
|
|
captured_fence = DuplexFence("sid", epoch=0, turn_id=0)
|
|
session.turn_id = 1
|
|
|
|
await handler._signal_runtime_session(
|
|
session,
|
|
"input.commit",
|
|
fence=captured_fence,
|
|
)
|
|
|
|
assert engine.calls == [("signal", captured_fence)]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_openai_handler_forwards_cancel_fence_through_async_omni_facade():
|
|
engine = FenceRecordingEngine()
|
|
app = object.__new__(AsyncOmni)
|
|
app.engine = engine
|
|
handler = OmniDuplexSessionHandler(chat_service=FakeChatService(app))
|
|
session = DuplexSession(session_id="sid-cancel", config=DuplexSessionConfig(extra_body={"native_duplex": True}))
|
|
cancelled_fence = DuplexFence("sid-cancel", epoch=2, turn_id=3)
|
|
next_fence = DuplexFence("sid-cancel", epoch=3, turn_id=3)
|
|
|
|
assert await handler._signal_runtime_session(
|
|
session,
|
|
"input.cancel",
|
|
fence=cancelled_fence,
|
|
next_fence=next_fence,
|
|
)
|
|
|
|
assert engine.calls == [("signal", cancelled_fence)]
|
|
assert engine.signal_next_fences == [next_fence]
|