137 lines
5.5 KiB
Python
137 lines
5.5 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""In-process duplex client: the :class:`DuplexClient` usage over a ``DuplexOmni``.
|
|
|
|
Example::
|
|
|
|
from vllm_omni.clients.inline_duplex import InlineDuplexClient
|
|
from vllm_omni.clients.minicpmo_4_5 import create_duplex_session_config
|
|
from vllm_omni.entrypoints.duplex_omni import DuplexOmni
|
|
|
|
omni = DuplexOmni(model="openbmb/MiniCPM-o-4_5", trust_remote_code=True)
|
|
cfg = create_duplex_session_config(ref_audio=audio_data_url(wav_bytes))
|
|
async with InlineDuplexClient(omni, model="openbmb/MiniCPM-o-4_5", config=cfg) as client:
|
|
await client.stream_pcm(pcm16)
|
|
await client.commit()
|
|
async for response in client.responses():
|
|
async for chunk in response.audio():
|
|
play(chunk)
|
|
break
|
|
|
|
There is no transport to drop, so this client has no reconnect, resume,
|
|
heartbeat or event-acknowledgement machinery. The ``DuplexOmni`` instance is
|
|
passed in by the caller; this module never imports server runtime code at
|
|
import time.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from typing import TYPE_CHECKING
|
|
|
|
from vllm_omni.clients.duplex import DuplexClientBase, DuplexConnectionError, SessionConfig
|
|
|
|
if TYPE_CHECKING:
|
|
from vllm_omni.entrypoints.duplex_omni import DuplexOmni, DuplexSessionHandle
|
|
from vllm_omni.protocol.duplex import RealtimeInputDefaults
|
|
|
|
__all__ = ["InlineDuplexClient"]
|
|
|
|
#: Acknowledgements are a websocket-journal concern; heartbeats are lease touches and go through.
|
|
_TRANSPORT_ONLY_EVENTS = frozenset({"session.event_ack"})
|
|
|
|
|
|
class InlineDuplexClient(DuplexClientBase):
|
|
"""Async client for one duplex session on an in-process :class:`DuplexOmni`."""
|
|
|
|
def __init__(
|
|
self,
|
|
omni: DuplexOmni,
|
|
*,
|
|
model: str,
|
|
config: SessionConfig | None = None,
|
|
handshake_timeout_s: float = 30.0,
|
|
) -> None:
|
|
super().__init__(model=model, config=config, handshake_timeout_s=handshake_timeout_s)
|
|
self.omni = omni
|
|
self._handle: DuplexSessionHandle | None = None
|
|
self._pump_task: asyncio.Task[None] | None = None
|
|
self._handshake_pending = False
|
|
|
|
# -- transport hooks ------------------------------------------------------
|
|
|
|
async def _open(self) -> None:
|
|
# The handle's first event is ``session.created``; the pump delivers
|
|
# it through ``_dispatch`` where the base handshake is waiting.
|
|
self._handle = await self.omni.open_session(self.config.to_session_payload(model=self.model))
|
|
# The base class follows up with the ``session.update`` handshake; the
|
|
# open above already carried that session object, so it is not resent.
|
|
self._handshake_pending = True
|
|
self._pump_task = asyncio.create_task(self._pump_events(), name="inline-duplex-client-pump")
|
|
|
|
def _input_defaults(self) -> RealtimeInputDefaults:
|
|
"""Audio format / sample-rate / VAD defaults declared by this session's object."""
|
|
from vllm_omni.protocol.duplex import RealtimeInputDefaults
|
|
|
|
session_payload = self.session_info or self.config.to_session_payload(model=self.model)
|
|
return RealtimeInputDefaults().with_session_payload(session_payload)
|
|
|
|
async def _send_command(self, payload: dict[str, object]) -> None:
|
|
handle = self._handle
|
|
if handle is None or handle.closed:
|
|
raise DuplexConnectionError("send failed: session is not open")
|
|
event_type = payload.get("type")
|
|
if event_type == "session.update" and self._handshake_pending:
|
|
self._handshake_pending = False
|
|
return
|
|
if event_type == "session.close":
|
|
await handle.close()
|
|
return
|
|
if event_type in _TRANSPORT_ONLY_EVENTS:
|
|
return
|
|
# Lazy import keeps ``vllm_omni.clients`` free of runtime imports at
|
|
# module load; the caller already holds a live DuplexOmni.
|
|
from vllm_omni.engine.duplex.commands import DuplexCommandError, command_from_realtime
|
|
from vllm_omni.engine.duplex.events import error_event
|
|
|
|
try:
|
|
command = command_from_realtime(payload, defaults=self._input_defaults())
|
|
except DuplexCommandError as exc:
|
|
await self._dispatch(
|
|
error_event(exc.code, str(exc), event_id=exc.event_id or payload.get("event_id")).to_realtime()
|
|
)
|
|
return
|
|
await handle.submit(command)
|
|
|
|
async def _teardown(self) -> None:
|
|
pump = self._pump_task
|
|
if pump is not None and not pump.done():
|
|
pump.cancel()
|
|
try:
|
|
await pump
|
|
except (asyncio.CancelledError, Exception): # noqa: BLE001
|
|
pass
|
|
handle = self._handle
|
|
if handle is not None and not handle.closed:
|
|
try:
|
|
await handle.close()
|
|
except Exception: # noqa: BLE001
|
|
pass
|
|
|
|
# -- internals ---------------------------------------------------------------
|
|
|
|
async def _pump_events(self) -> None:
|
|
handle = self._handle
|
|
assert handle is not None
|
|
reason = "closed"
|
|
try:
|
|
async for event in handle.events():
|
|
await self._dispatch(event.to_realtime())
|
|
except asyncio.CancelledError:
|
|
raise
|
|
except Exception as exc: # noqa: BLE001
|
|
reason = f"session failed: {exc}"
|
|
self._finalize(reason, expected=False)
|
|
return
|
|
self._finalize(reason, expected=True)
|