vllm-omni/vllm_omni/experimental/fullduplex/personaplex/data_plane.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

from __future__ import annotations

from collections.abc import Callable, Iterator, Mapping
from dataclasses import dataclass

import numpy as np

from vllm_omni.experimental.fullduplex.engine.contracts import (
    duplex_resource_request_belongs_to_session,
)

EncodeAudio = Callable[[object, int, str, float | None], str | None]


@dataclass(frozen=True, slots=True)
class PersonaPlexDataPlaneContext:
    epoch: int = 0
    turn_id: int = 0
    active_response_turn_id: int | None = None
    active_response_id: str | None = None
    auto_responds: bool = True
    response_format: str = "wav"
    speed: float | None = None
    modalities: tuple[str, ...] = ("audio", "text")


@dataclass(slots=True)
class _RequestCursor:
    audio_samples: int = 0
    text: str = ""
    terminal: bool = False


class PersonaPlexDataPlaneSession:
    """Project cumulative staged PersonaPlex output into Realtime deltas."""

    def __init__(self, encode_audio: EncodeAudio) -> None:
        self._encode_audio = encode_audio
        self._requests: dict[str, _RequestCursor] = {}

    def begin_request(self, request_id: str) -> None:
        self._requests.setdefault(request_id, _RequestCursor()).terminal = False

    def is_terminal(self, request_id: str | None) -> bool:
        if request_id is None:
            return False
        state = self._requests.get(request_id)
        return state is not None and state.terminal

    def mark_terminal(self, request_id: str) -> None:
        self._requests.setdefault(request_id, _RequestCursor()).terminal = True

    def close_stream(self, request_id: str) -> None:
        self._requests.pop(request_id, None)

    def close_session(self, session_id: str, *, active_request_id: str | None = None) -> None:
        if active_request_id is not None:
            self._requests.pop(active_request_id, None)
        for request_id in list(self._requests):
            if duplex_resource_request_belongs_to_session(request_id, session_id):
                self._requests.pop(request_id, None)

    def project(
        self,
        result: object,
        *,
        context: PersonaPlexDataPlaneContext | None = None,
    ) -> Iterator[dict[str, object]]:
        if not isinstance(result, dict):
            return
        outputs = result.get("data_plane_outputs")
        if not isinstance(outputs, list):
            return
        context = context or PersonaPlexDataPlaneContext()
        for output in outputs:
            projected = self._project_output(output, context=context)
            if projected is not None:
                yield projected

    def _project_output(
        self,
        output: object,
        *,
        context: PersonaPlexDataPlaneContext,
    ) -> dict[str, object] | None:
        output, completion = _unwrap_output(output)
        request_id = getattr(output, "request_id", None)
        if not isinstance(request_id, str) or not request_id:
            request_id = None
        state = self._requests.setdefault(request_id, _RequestCursor()) if request_id is not None else _RequestCursor()
        multimodal = _multimodal_output(output, completion)
        audio = _audio_value(multimodal)
        audio_delta = _slice_audio(audio, state.audio_samples)
        audio_samples = _num_samples(audio)
        if audio_samples is not None:
            state.audio_samples = audio_samples

        text = _text_value(multimodal, completion)
        text_delta = _text_delta(text, state.text)
        if text:
            state.text = text

        sample_rate_hz = _sample_rate(multimodal)
        encoded = self._encode_audio(audio_delta, sample_rate_hz, context.response_format, context.speed)
        if not encoded and not text_delta:
            return None
        delta_samples = _num_samples(audio_delta) or 0
        return {
            "supported": True,
            "stage_role": "tts",
            "is_listen": False,
            "data_plane_request_id": request_id,
            "text": text_delta,
            "audio_data": encoded or "",
            "audio_format": context.response_format,
            "sample_rate_hz": sample_rate_hz,
            "audio_duration_ms": round(delta_samples * 1000 / max(1, sample_rate_hz)),
            "end_of_turn": False,
            "uses_model_runner_scheduler": True,
            "runner_kv_backed": True,
            "runtime_impl": "scheduler_data_plane",
            "owned_runtime": False,
        }


def _unwrap_output(output: object) -> tuple[object, object | None]:
    inner = getattr(output, "request_output", None)
    if inner is not None and inner is not output:
        output = inner
    outputs = getattr(output, "outputs", None)
    completion = outputs[0] if isinstance(outputs, list) and outputs else None
    return output, completion


def _multimodal_output(output: object, completion: object | None) -> dict[str, object]:
    for candidate in (
        getattr(output, "multimodal_output", None),
        getattr(completion, "multimodal_output", None) if completion is not None else None,
    ):
        if isinstance(candidate, Mapping) and candidate:
            return dict(candidate)
    return {}


def _audio_value(multimodal: dict[str, object]) -> object | None:
    value = next(
        (multimodal[key] for key in ("audio", "model_outputs", "latent") if key in multimodal),
        None,
    )
    if isinstance(value, list) and len(value) == 1:
        return value[0]
    return value


def _text_value(multimodal: dict[str, object], completion: object | None) -> str:
    for candidate in (
        multimodal.get("text"),
        multimodal.get("llm_output_text"),
        getattr(completion, "text", None) if completion is not None else None,
    ):
        if isinstance(candidate, str) and candidate:
            return candidate
    return ""


def _text_delta(text: str, previous: str) -> str:
    if not text:
        return ""
    if text == previous:
        return ""
    if previous and text.startswith(previous):
        return text[len(previous) :]
    return text


def _slice_audio(audio: object | None, offset: int) -> object | None:
    samples = _num_samples(audio)
    if samples is None or samples <= 0:
        return None
    if offset <= 0 or samples < offset:
        return audio
    if samples == offset:
        return None
    try:
        import torch

        if isinstance(audio, torch.Tensor):
            return audio.reshape(-1)[offset:].contiguous()
    except Exception:
        pass
    return np.asarray(audio, dtype=np.float32).reshape(-1)[offset:]


def _num_samples(audio: object | None) -> int | None:
    if audio is None:
        return None
    try:
        import torch

        if isinstance(audio, torch.Tensor):
            return int(audio.numel())
    except Exception:
        pass
    try:
        return int(np.asarray(audio, dtype=np.float32).size)
    except Exception:
        return None


def _sample_rate(multimodal: dict[str, object]) -> int:
    value = multimodal.get("sr", multimodal.get("sample_rate_hz", 24000))
    if isinstance(value, list) and value:
        value = value[0]
    if hasattr(value, "item"):
        try:
            value = value.item()
        except Exception:
            value = 24000
    return int(value) if isinstance(value, int | float) else 24000


__all__ = [
    "PersonaPlexDataPlaneContext",
    "PersonaPlexDataPlaneSession",
]