vllm-omni/vllm_omni/model_executor/models/personaplex/duplex/plugin.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""PersonaPlex full-duplex model plugin: engine policy and session policy in one class.

PersonaPlex is a pure-lockstep model. Engine policy is one 80 ms frame per
resumable Stage 0 append with greedy one-token decoding; session policy is
the bundled voice prompt and the persona text, resolved once at open into the
server-owned runtime configuration the worker replays as its prefill.
"""

from __future__ import annotations

import asyncio
from collections.abc import Mapping
from copy import deepcopy
from dataclasses import dataclass, field
from pathlib import PurePath
from typing import TYPE_CHECKING

from vllm.sampling_params import SamplingParams

from vllm_omni.engine.duplex.config import DuplexCapabilities, DuplexSessionConfig
from vllm_omni.engine.duplex.contracts import (
    DuplexAppendPlan,
    DuplexFence,
    DuplexOutputDecision,
)
from vllm_omni.engine.duplex.plugin import (
    DefaultDuplexModelSessionState,
    DuplexModelPlugin,
    DuplexRuntimeConfigError,
    EncodeAudio,
    reject_changed_runtime_value,
)
from vllm_omni.model_executor.common.duplex.payload import decode_pcm_f32le_payload
from vllm_omni.model_executor.models.personaplex.duplex.capabilities import personaplex_capabilities
from vllm_omni.model_executor.models.personaplex.duplex.config import (
    DEFAULT_PERSONA,
    DEFAULT_VOICE,
    FRAME_SIZE,
    SAMPLE_RATE,
)
from vllm_omni.model_executor.models.personaplex.duplex.data_plane import PersonaPlexDataPlaneSession
from vllm_omni.model_executor.models.personaplex.duplex.input import PersonaPlexPcmAppendBuffer

if TYPE_CHECKING:
    from vllm.config import ModelConfig

#: Runtime-config keys the server derives from the session fields; a client
#: cannot set them through ``extra_body``.
PRIVATE_RUNTIME_CONFIG_KEYS = frozenset(
    {
        "personaplex_prefill_slots",
        "personaplex_model_path",
        "personaplex_voice_prompt",
        "personaplex_persona",
    }
)


@dataclass
class PersonaPlexSessionState(DefaultDuplexModelSessionState):
    """Per-session model state: the framework flags plus the 80 ms framing buffer."""

    audio_buffer: PersonaPlexPcmAppendBuffer = field(default_factory=PersonaPlexPcmAppendBuffer)


def voice_name(value: object) -> str:
    """A bundled ``.pt`` basename (``NATF2.pt``); paths are refused so the worker only reads the checkpoint."""
    voice = value if isinstance(value, str) and value else DEFAULT_VOICE
    path = PurePath(voice)
    if path.name != voice or path.suffix != ".pt" or any(part == ".." for part in path.parts):
        raise DuplexRuntimeConfigError("PersonaPlex voice must be a bundled .pt basename", code="invalid_voice")
    return voice


def persona_text(value: object) -> str:
    return str(value) if isinstance(value, str) and value else DEFAULT_PERSONA


def prefill_slot_count(runtime_config: Mapping[str, object]) -> int:
    raw = runtime_config.get("personaplex_prefill_slots", 0)
    try:
        return max(0, int(raw))  # type: ignore[arg-type]
    except (TypeError, ValueError) as exc:
        raise ValueError("personaplex_prefill_slots must be a non-negative integer") from exc


class PersonaPlexDuplexPlugin(DuplexModelPlugin):
    """PersonaPlex-owned sampling policy, append planning, session state and output projection."""

    plugin_id = "personaplex"
    private_runtime_config_keys = PRIVATE_RUNTIME_CONFIG_KEYS
    #: One codec frame: the runner keeps a model turn clocked with these when the client pauses.
    silence_continuation_samples = FRAME_SIZE
    silence_continuation_sample_rate_hz = SAMPLE_RATE

    def __init__(self, encode_audio: EncodeAudio) -> None:
        super().__init__(encode_audio)
        self.data_plane = PersonaPlexDataPlaneSession(encode_audio)
        # The prefill length depends only on (model, voice, persona); the
        # voice bundle and tokenizer are read once per distinct triple.
        self._prefill_slots: dict[tuple[str, str, str], int] = {}
        self._prefill_lock = asyncio.Lock()

    # ---- engine policy (the resumable Stage 0 request) ----

    def configure_sampling_params(
        self,
        *,
        runtime_config: dict[str, object],
        defaults: tuple[object, ...],
    ) -> tuple[object, ...]:
        del runtime_config
        if not defaults:
            return defaults
        configured = list(defaults)
        stage0 = defaults[0]
        if isinstance(stage0, SamplingParams):
            # Greedy, one temporal token per frame: the native stepper takes
            # argmax for the text head and the depformer.
            stage0 = stage0.clone()
            stage0.temperature = 0.0
            stage0.top_k = 1
            stage0.max_tokens = 1
            configured[0] = stage0
        return tuple(configured)

    def plan_append(
        self,
        *,
        request_id: str,
        fence: DuplexFence,
        session_config: dict[str, object],
        runtime_config: dict[str, object],
        seq: int,
        turn_seq: int,
        payload: object,
        final: bool,
        sampling_params: object,
    ) -> DuplexAppendPlan:
        del sampling_params
        decode_pcm_f32le_payload(payload, sample_rate_hz=SAMPLE_RATE, exact_samples=FRAME_SIZE, model="PersonaPlex")
        normalized_payload = dict(payload)  # type: ignore[call-overload]
        # One scheduler slot per frame; the first append of an epoch also
        # carries the voice/persona prefill (a new epoch is a new Stage 0
        # request with fresh KV, so the worker replays it).
        prompt_slots = 1 + (prefill_slot_count(runtime_config) if seq <= 1 else 0)
        return DuplexAppendPlan(
            prompt={
                "prompt_token_ids": [0] * prompt_slots,
                "model_intermediate_buffer": {
                    "request_id": request_id,
                    "global_request_id": [fence.session_id],
                    "duplex": {
                        "data_plane": True,
                        "fence": fence,
                        "session_id": fence.session_id,
                        "epoch": fence.epoch,
                        "turn_id": fence.turn_id,
                        "seq": seq,
                        "turn_seq": turn_seq,
                        "mode": "append_audio_chunk",
                        "payload": normalized_payload,
                        "final": final,
                        "session_config": dict(session_config),
                        "runtime_config": dict(runtime_config),
                        "scheduler_token_budget": prompt_slots,
                    },
                },
            }
        )

    def decide_output(
        self,
        *,
        stage_id: int,
        final_stage_id: int,
        segment_finished: bool,
        segment_token_ids: tuple[int, ...],
        segment_output_metadata: dict[str, object],
        output: object,
    ) -> DuplexOutputDecision | None:
        # Always-clocked model: what the client hears comes from the final
        # stage data plane, never from a Stage 0 listen/speak decision.
        del stage_id, final_stage_id, segment_finished, segment_token_ids, segment_output_metadata, output
        return None

    # ---- session policy ----

    def create_session_state(self) -> PersonaPlexSessionState:
        return PersonaPlexSessionState()

    def capabilities(self, *, max_sessions: int) -> DuplexCapabilities:
        return personaplex_capabilities(max_sessions=max_sessions)

    async def prepare_runtime_config(
        self, config: DuplexSessionConfig, *, model_config: ModelConfig | None
    ) -> dict[str, object]:
        self.validate_client_extra_body(config.extra_body)
        model_path = getattr(model_config, "model", None)
        if not isinstance(model_path, str) or not model_path:
            raise DuplexRuntimeConfigError("PersonaPlex model path is unavailable", code="model_path_unavailable")
        voice = voice_name(config.voice)
        persona = persona_text(config.instructions)
        prefill_slots = await self._prefill_slots_for(model_path, voice, persona)
        return {
            "personaplex_model_path": model_path,
            "personaplex_voice_prompt": voice,
            "personaplex_persona": persona,
            "personaplex_prefill_slots": prefill_slots,
        }

    async def _prefill_slots_for(self, model_path: str, voice: str, persona: str) -> int:
        key = (model_path, voice, persona)
        cached = self._prefill_slots.get(key)
        if cached is not None:
            return cached
        from vllm_omni.model_executor.models.personaplex.duplex.stage0 import personaplex_prefill_slots

        async with self._prefill_lock:
            cached = self._prefill_slots.get(key)
            if cached is not None:
                return cached
            try:
                # Reads the voice bundle and the tokenizer: off the orchestrator loop.
                slots = await asyncio.to_thread(personaplex_prefill_slots, model_path, voice, persona)
            except Exception as exc:
                raise DuplexRuntimeConfigError(
                    f"PersonaPlex voice/persona prefill could not be prepared: {exc}",
                    code="prefill_unavailable",
                ) from exc
            self._prefill_slots[key] = slots
            return slots

    def runtime_config_for_update(
        self,
        config: DuplexSessionConfig,
        current: Mapping[str, object],
    ) -> dict[str, object]:
        self.validate_client_extra_body(config.extra_body)
        new_persona = persona_text(config.instructions)
        reject_changed_runtime_value(
            new_persona,
            current.get("personaplex_persona"),
            message="PersonaPlex persona (instructions) cannot be changed after the session is created",
            code="persona_update_unsupported",
        )
        new_voice = voice_name(config.voice)
        reject_changed_runtime_value(
            new_voice,
            current.get("personaplex_voice_prompt"),
            message="PersonaPlex voice cannot be changed after the session is created",
            code="voice_update_unsupported",
        )
        runtime_config = deepcopy(dict(current))
        runtime_config["personaplex_voice_prompt"] = new_voice
        runtime_config["personaplex_persona"] = new_persona
        return runtime_config


__all__ = [
    "PRIVATE_RUNTIME_CONFIG_KEYS",
    "PersonaPlexDuplexPlugin",
    "PersonaPlexSessionState",
    "persona_text",
    "prefill_slot_count",
    "voice_name",
]