vllm-omni/vllm_omni/experimental/fullduplex/client.py

"""Reusable Realtime WebSocket, PCM, and event helpers for MiniCPM-o demos."""

from __future__ import annotations

import asyncio
import base64
import json
import math
import time
import wave
from collections.abc import Callable
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit

from vllm_omni.metrics.definitions import compute_audio_rtf

try:
    import websockets
    from websockets.exceptions import ConnectionClosed
except ImportError as exc:  # pragma: no cover - example dependency
    raise SystemExit("Install websockets first: pip install websockets") from exc

PCM16_SAMPLE_RATE = 16_000
PCM16_BYTES_PER_SAMPLE = 2


def _rounded_ms(value: float) -> float:
    return round(float(value), 3)


def _interval_summary(values: list[float]) -> dict[str, float | int]:
    clean = sorted(_rounded_ms(value) for value in values if math.isfinite(value) and value >= 0)
    if not clean:
        return {"count": 0, "mean": 0.0, "p50": 0.0, "p95": 0.0, "max": 0.0}

    def nearest_rank(percentile: float) -> float:
        index = max(0, math.ceil(percentile * len(clean)) - 1)
        return clean[min(index, len(clean) - 1)]

    return {
        "count": len(clean),
        "mean": _rounded_ms(sum(clean) / len(clean)),
        "p50": nearest_rank(0.50),
        "p95": nearest_rank(0.95),
        "max": clean[-1],
    }


def summarize_session_request_metrics(
    request_metrics: list[dict[str, object]],
    *,
    session_id: str | None,
) -> dict[str, object]:
    """Average client-observed metrics across turns that emitted audio."""

    def mean(metric: str, *, digits: int = 3) -> float | None:
        values = [
            float(request[metric])
            for request in request_metrics
            if isinstance(request.get(metric), int | float) and math.isfinite(float(request[metric]))
        ]
        return round(sum(values) / len(values), digits) if values else None

    return {
        "session_id": session_id,
        "audio_turn_count": len(request_metrics),
        "mean_ttft_ms": mean("ttft_ms"),
        "mean_ttfp_ms": mean("ttfp_ms"),
        "mean_rtf": mean("rtf", digits=6),
    }


def _event_stage_metrics(event: dict[str, object]) -> dict[str, object] | None:
    candidates: list[object] = [event.get("vllm_omni")]
    metadata = event.get("metadata")
    if isinstance(metadata, dict):
        candidates.extend((metadata, metadata.get("vllm_omni")))
    response = event.get("response")
    if isinstance(response, dict):
        response_metadata = response.get("metadata")
        if isinstance(response_metadata, dict):
            candidates.extend((response_metadata, response_metadata.get("vllm_omni")))
    for candidate in candidates:
        if not isinstance(candidate, dict):
            continue
        stage_metrics = candidate.get("stage_metrics")
        if isinstance(stage_metrics, dict):
            return stage_metrics
    return None


def build_realtime_url(
    url: str,
    model: str,
    *,
    autostart: bool | None = None,
    session_id: str | None = None,
) -> str:
    """Add the explicit native-duplex query parameters to a Realtime URL."""
    parts = urlsplit(url)
    query = dict(parse_qsl(parts.query, keep_blank_values=True))
    query.setdefault("duplex", "1")
    query.setdefault("model", model)
    query.setdefault("minicpmo45_native_duplex", "1")
    if autostart is not None:
        query.setdefault("autostart", "1" if autostart else "0")
    if session_id:
        query.setdefault("session_id", session_id)
    return urlunsplit((parts.scheme, parts.netloc, parts.path, urlencode(query), parts.fragment))


def read_pcm16_wav(path: Path) -> bytes:
    """Read a mono, uncompressed, 16 kHz PCM16 WAV file."""
    with wave.open(str(path), "rb") as wav_file:
        if wav_file.getnchannels() != 1:
            raise ValueError("input WAV must be mono")
        if wav_file.getsampwidth() != PCM16_BYTES_PER_SAMPLE:
            raise ValueError("input WAV must be 16-bit PCM")
        if wav_file.getframerate() != PCM16_SAMPLE_RATE:
            raise ValueError("input WAV must be 16 kHz")
        if wav_file.getcomptype() != "NONE":
            raise ValueError("input WAV must be uncompressed PCM")
        return wav_file.readframes(wav_file.getnframes())


def write_pcm16_wav(path: Path, pcm16: bytes, *, sample_rate_hz: int) -> None:
    """Write mono PCM16 bytes as a WAV artifact."""
    with wave.open(str(path), "wb") as wav_file:
        wav_file.setnchannels(1)
        wav_file.setsampwidth(PCM16_BYTES_PER_SAMPLE)
        wav_file.setframerate(sample_rate_hz)
        wav_file.writeframes(pcm16)


async def wait_for(
    predicate: Callable[[], bool],
    *,
    timeout_s: float,
    label: str,
) -> None:
    """Wait for a collector predicate without coupling to a scenario runner."""
    deadline = time.monotonic() + timeout_s
    while time.monotonic() < deadline:
        if predicate():
            return
        await asyncio.sleep(0.02)
    raise TimeoutError(f"Timed out waiting for {label}")


@dataclass
class RealtimeEventCollector:
    """Collect server events and decode response audio by response identity."""

    events: list[dict[str, object]] = field(default_factory=list)
    event_received_at_s: list[float] = field(default_factory=list)
    response_audio: dict[str, list[bytes]] = field(default_factory=dict)
    response_ids: list[str] = field(default_factory=list)
    output_sample_rate_hz: int = 24_000

    @staticmethod
    def response_id(event: dict[str, object]) -> str | None:
        response_id = event.get("response_id")
        if isinstance(response_id, str):
            return response_id
        response = event.get("response")
        if isinstance(response, dict):
            response_id = response.get("id")
            if isinstance(response_id, str):
                return response_id
        return None

    def add(self, event: dict[str, object], *, received_at_s: float | None = None) -> None:
        received_at = time.monotonic() if received_at_s is None else float(received_at_s)
        stored_event = dict(event)
        stored_event.setdefault("_client_received_at_s", received_at)
        self.events.append(stored_event)
        self.event_received_at_s.append(received_at)
        response_id = self.response_id(stored_event)
        event_type = stored_event.get("type")
        if event_type == "response.created" and response_id and response_id not in self.response_ids:
            self.response_ids.append(response_id)
        if event_type == "response.audio.delta":
            delta = stored_event.get("delta") or stored_event.get("audio")
            if isinstance(delta, str) and response_id:
                try:
                    self.response_audio.setdefault(response_id, []).append(base64.b64decode(delta))
                except ValueError:
                    pass
            sample_rate_hz = stored_event.get("sample_rate_hz")
            if isinstance(sample_rate_hz, int) and sample_rate_hz > 0:
                self.output_sample_rate_hz = sample_rate_hz

    def count(self, event_type: str) -> int:
        return sum(event.get("type") == event_type for event in self.events)

    def audio_bytes(self, response_id: str | None = None) -> bytes:
        if response_id is not None:
            return b"".join(self.response_audio.get(response_id, ()))
        return b"".join(
            chunk for response_id in self.response_ids for chunk in self.response_audio.get(response_id, ())
        )

    def errors(self) -> list[dict[str, object]]:
        return [event for event in self.events if event.get("type") == "error"]

    def first_received_at(
        self,
        *event_types: str,
        after_s: float = 0.0,
    ) -> float | None:
        for event, received_at_s in zip(self.events, self.event_received_at_s, strict=True):
            if received_at_s >= after_s and event.get("type") in event_types:
                return received_at_s
        return None

    def last_received_at(self, event_type: str) -> float | None:
        for event, received_at_s in zip(
            reversed(self.events),
            reversed(self.event_received_at_s),
            strict=True,
        ):
            if event.get("type") == event_type:
                return received_at_s
        return None

    def timing_summary(
        self,
        *,
        after_s: float,
        input_committed_at_s: float | None = None,
        response_id: str | None = None,
        measurement_origin: dict[str, str] | None = None,
    ) -> dict[str, object]:
        """Summarize engine token metrics and client-observed audio cadence."""
        stage0_metrics: dict[str, object] | None = None
        response_created_at_s: float | None = None
        first_text_received_at_s: float | None = None
        audio_received_at_s: list[float] = []
        cumulative_audio_ms: list[float] = []
        for event, received_at_s in zip(self.events, self.event_received_at_s, strict=True):
            if received_at_s < after_s:
                continue
            event_response_id = self.response_id(event)
            if response_id is not None and event_response_id != response_id:
                continue
            if event.get("type") == "response.created" and response_created_at_s is None:
                response_created_at_s = received_at_s
            if (
                event.get("type")
                in {
                    "response.audio_transcript.delta",
                    "response.output_text.delta",
                    "response.text.delta",
                }
                and isinstance(event.get("delta"), str)
                and bool(event["delta"])
                and first_text_received_at_s is None
            ):
                first_text_received_at_s = received_at_s

            stage_metrics = _event_stage_metrics(event)
            stage0 = stage_metrics.get("0") if isinstance(stage_metrics, dict) else None
            if isinstance(stage0, dict):
                stage0_metrics = stage0

            if event.get("type") != "response.audio.delta" or (
                response_id is not None and event_response_id != response_id
            ):
                continue
            delta = event.get("delta") or event.get("audio")
            if not isinstance(delta, str) or not delta:
                continue
            audio_received_at_s.append(received_at_s)
            metadata = event.get("metadata")
            duration_ms = metadata.get("audio_duration_ms") if isinstance(metadata, dict) else None
            if isinstance(duration_ms, int | float) and math.isfinite(float(duration_ms)):
                cumulative_audio_ms.append(max(0.0, float(duration_ms)))

        result: dict[str, object] = {}
        if stage0_metrics is not None:
            raw_itls = stage0_metrics.get("vllm_itls_ms")
            itls = (
                [float(value) for value in raw_itls if isinstance(value, int | float)]
                if isinstance(raw_itls, list)
                else []
            )
            result["stage0_tokens"] = {
                "source": "engine_stage_metrics",
                "output_token_count": int(stage0_metrics.get("num_tokens_out") or 0),
                "ttft_ms": float(stage0_metrics.get("vllm_ttft_ms") or 0.0),
                "tpot_ms": float(stage0_metrics.get("vllm_tpot_ms") or 0.0),
                "inter_token_interval_ms": _interval_summary(itls),
            }

        if audio_received_at_s:
            intervals_ms = [
                (current - previous) * 1000.0 for previous, current in zip(audio_received_at_s, audio_received_at_s[1:])
            ]
            chunk_durations_ms: list[float] = []
            previous_duration_ms = 0.0
            for duration_ms in cumulative_audio_ms:
                chunk_durations_ms.append(
                    duration_ms - previous_duration_ms if duration_ms >= previous_duration_ms else duration_ms
                )
                previous_duration_ms = duration_ms
            interval_summary = _interval_summary(intervals_ms)
            result["audio_output"] = {
                "source": "client_monotonic_receive",
                "chunk_count": len(audio_received_at_s),
                "response_created_to_first_audio_ms": (
                    _rounded_ms((audio_received_at_s[0] - response_created_at_s) * 1000.0)
                    if response_created_at_s is not None
                    else None
                ),
                "commit_to_first_audio_ms": (
                    _rounded_ms((audio_received_at_s[0] - input_committed_at_s) * 1000.0)
                    if input_committed_at_s is not None
                    else None
                ),
                "inter_chunk_interval_ms": interval_summary,
                "chunk_duration_ms": _interval_summary(chunk_durations_ms),
                "max_chunk_gap_ms": interval_summary["max"],
            }
            request_started_at_s = input_committed_at_s if input_committed_at_s is not None else response_created_at_s
            if request_started_at_s is not None:
                audio_duration_ms = (
                    max(cumulative_audio_ms)
                    if cumulative_audio_ms
                    else len(self.audio_bytes(response_id))
                    * 1000.0
                    / (self.output_sample_rate_hz * PCM16_BYTES_PER_SAMPLE)
                )
                audio_generation_ms = max(
                    0.0,
                    (audio_received_at_s[-1] - request_started_at_s) * 1000.0,
                )
                result["request_metrics"] = {
                    "source": "client_monotonic_receive",
                    "measurement_origin": measurement_origin
                    or {
                        "ttft": "input_audio_buffer.commit client send to first non-empty text delta",
                        "ttfp": "input_audio_buffer.commit client send to first audio packet",
                        "rtf": "commit-to-last-audio receive time divided by emitted audio duration",
                    },
                    "ttft_ms": (
                        _rounded_ms((first_text_received_at_s - request_started_at_s) * 1000.0)
                        if first_text_received_at_s is not None
                        else None
                    ),
                    "ttfp_ms": _rounded_ms((audio_received_at_s[0] - request_started_at_s) * 1000.0),
                    "rtf": round(
                        compute_audio_rtf(
                            audio_generation_ms / 1000.0,
                            audio_duration_ms / 1000.0,
                        ),
                        6,
                    )
                    if audio_duration_ms > 0
                    else None,
                    "audio_generation_ms": _rounded_ms(audio_generation_ms),
                    "audio_duration_ms": _rounded_ms(audio_duration_ms),
                }
        return result


class RealtimeDuplexClient:
    """Small async client used by the user demo and reusable smoke probes."""

    def __init__(self, url: str, *, max_size: int = 64 * 1024 * 1024) -> None:
        self.url = url
        self.max_size = max_size
        self.events = RealtimeEventCollector()
        self._ws: Any = None
        self._reader_task: asyncio.Task[None] | None = None

    async def __aenter__(self) -> RealtimeDuplexClient:
        self._ws = await websockets.connect(self.url, max_size=self.max_size)
        self._reader_task = asyncio.create_task(self._read_events())
        return self

    async def __aexit__(self, exc_type, exc, traceback) -> None:
        if self._ws is not None:
            await self._ws.close()
        if self._reader_task is not None:
            self._reader_task.cancel()
            try:
                await self._reader_task
            except asyncio.CancelledError:
                pass

    async def _read_events(self) -> None:
        try:
            while True:
                raw = await self._ws.recv()
                if not isinstance(raw, str):
                    continue
                event = json.loads(raw)
                if isinstance(event, dict):
                    self.events.add(event)
        except ConnectionClosed:
            return

    async def send(self, event: dict[str, object]) -> None:
        await self._ws.send(json.dumps(event))

    async def configure(
        self,
        model: str,
        *,
        output_audio_format: str = "pcm16",
        ref_audio: str | None = None,
        instructions: str | None = None,
        initial_user_text: str | None = None,
        native_duplex: bool = True,
        auto_response: bool = True,
        temperature: float | None = None,
        extra_body: dict[str, object] | None = None,
        session_id: str | None = None,
        timeout_s: float = 20.0,
    ) -> None:
        session_extra_body = dict(extra_body or {})
        if native_duplex:
            session_extra_body.update(
                {
                    "auto_response": auto_response,
                    "minicpmo45_native_duplex": True,
                    "force_listen_count": 0,
                }
            )
        else:
            session_extra_body["minicpmo45_native_duplex"] = False
        session: dict[str, object] = {
            "model": model,
            "modalities": ["audio", "text"],
            "input_audio_format": "pcm16",
            "output_audio_format": output_audio_format,
            "turn_detection": None,
            "overlap_policy": "listen_only",
            "playback_commit_policy": "ack_only",
            "extra_body": session_extra_body,
        }
        if temperature is not None:
            session["temperature"] = float(temperature)
        if ref_audio is not None:
            session["ref_audio"] = ref_audio
        if instructions is not None:
            session["instructions"] = instructions
        if initial_user_text is not None:
            extra_body = session["extra_body"]
            assert isinstance(extra_body, dict)
            extra_body["duplex_initial_user_text"] = initial_user_text
        if session_id:
            session["session_id"] = session_id
        await self.send({"type": "session.update", "session": session})
        await wait_for(
            lambda: self.events.count("session.created") > 0,
            timeout_s=timeout_s,
            label="session.created",
        )

    async def stream_pcm16(
        self,
        pcm16: bytes,
        *,
        chunk_ms: int = 200,
        realtime: bool = True,
    ) -> None:
        chunk_bytes = max(
            PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE * chunk_ms // 1000,
            PCM16_BYTES_PER_SAMPLE,
        )
        audio_end_ms = 0
        for offset in range(0, len(pcm16), chunk_bytes):
            chunk = pcm16[offset : offset + chunk_bytes]
            duration_ms = len(chunk) * 1000 // (PCM16_SAMPLE_RATE * PCM16_BYTES_PER_SAMPLE)
            audio_end_ms += duration_ms
            await self.send(
                {
                    "type": "input_audio_buffer.append",
                    "audio": base64.b64encode(chunk).decode("ascii"),
                    "input_audio_format": "pcm16",
                    "sample_rate_hz": PCM16_SAMPLE_RATE,
                    "duration_ms": duration_ms,
                    "audio_end_ms": audio_end_ms,
                }
            )
            if realtime:
                await asyncio.sleep(duration_ms / 1000)

    async def commit(self) -> None:
        await self.send({"type": "input_audio_buffer.commit", "final": True})

    async def acknowledge_playback(self) -> None:
        for response_id in self.events.response_ids:
            pcm16 = self.events.audio_bytes(response_id)
            if not pcm16:
                continue
            played_ms = len(pcm16) * 1000 // (self.events.output_sample_rate_hz * PCM16_BYTES_PER_SAMPLE)
            await self.send(
                {
                    "type": "playback.ack",
                    "response_id": response_id,
                    "item_id": f"item_{response_id}",
                    "played_ms": played_ms,
                    "committed_ms": played_ms,
                }
            )

    async def close_session(self, *, timeout_s: float = 20.0) -> None:
        await self.send({"type": "session.close"})
        await wait_for(
            lambda: self.events.count("session.closed") > 0,
            timeout_s=timeout_s,
            label="session.closed",
        )