149 lines
5.1 KiB
Python
149 lines
5.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Mapping
|
|
from typing import TYPE_CHECKING, TypedDict
|
|
|
|
if TYPE_CHECKING:
|
|
from vllm_omni.engine.duplex.contracts import DuplexFence
|
|
|
|
|
|
class DuplexIntermediateBuffer(TypedDict, total=False):
|
|
"""Structured keys carried through ``model_intermediate_buffer``.
|
|
|
|
The buffer remains a dict for scheduler and msgspec compatibility, but
|
|
duplex-specific producers and consumers should use the helpers in this
|
|
module instead of scattering nested string keys across serving, runner, and
|
|
model code.
|
|
"""
|
|
|
|
request_id: str
|
|
global_request_id: list[str]
|
|
prompt_token_ids: list[int]
|
|
llm_output_token_ids: list[int]
|
|
llm_output_text: list[str]
|
|
stream_output: bool
|
|
native_duplex: bool
|
|
ids: dict[str, object]
|
|
hidden_states: dict[str, object]
|
|
codes: dict[str, object]
|
|
meta: dict[str, object]
|
|
duplex: dict[str, object]
|
|
omni_payload: object
|
|
waveform: object
|
|
mel_spec: object
|
|
|
|
|
|
def build_duplex_append_prompt(
|
|
*,
|
|
request_id: str,
|
|
fence: DuplexFence,
|
|
session_config: Mapping[str, object],
|
|
runtime_config: Mapping[str, object],
|
|
seq: int,
|
|
turn_seq: int,
|
|
payload: object,
|
|
final: bool,
|
|
prompt_token_ids: list[int],
|
|
model_fields: Mapping[str, object],
|
|
) -> dict[str, object]:
|
|
"""Wrap model-selected tokens and payload in the shared append envelope.
|
|
|
|
Models own token selection and extra worker fields; the framework owns
|
|
request identity, sequencing and shallow config snapshots.
|
|
"""
|
|
return {
|
|
"prompt_token_ids": prompt_token_ids,
|
|
"model_intermediate_buffer": {
|
|
"request_id": request_id,
|
|
"global_request_id": [fence.session_id],
|
|
"duplex": {
|
|
**model_fields,
|
|
"data_plane": True,
|
|
"fence": fence,
|
|
"session_id": fence.session_id,
|
|
"epoch": fence.epoch,
|
|
"seq": seq,
|
|
"turn_id": fence.turn_id,
|
|
"turn_seq": turn_seq,
|
|
"mode": "append_audio_chunk",
|
|
"payload": payload,
|
|
"final": final,
|
|
"session_config": dict(session_config),
|
|
"runtime_config": dict(runtime_config),
|
|
"scheduler_token_budget": len(prompt_token_ids),
|
|
},
|
|
},
|
|
}
|
|
|
|
|
|
def build_duplex_intermediate_buffer(
|
|
*,
|
|
request_id: str,
|
|
prompt_token_ids: list[int] | None = None,
|
|
output_token_ids: list[int] | None = None,
|
|
output_text: str | None = None,
|
|
stream_output: bool = False,
|
|
native_duplex: bool = False,
|
|
) -> DuplexIntermediateBuffer:
|
|
buffer: DuplexIntermediateBuffer = {
|
|
"global_request_id": [str(request_id)],
|
|
"ids": {},
|
|
}
|
|
if prompt_token_ids is not None:
|
|
prompt_ids = [int(token_id) for token_id in prompt_token_ids]
|
|
buffer["prompt_token_ids"] = prompt_ids
|
|
buffer["ids"]["prompt"] = prompt_ids
|
|
if output_token_ids is not None:
|
|
output_ids = [int(token_id) for token_id in output_token_ids]
|
|
buffer["llm_output_token_ids"] = output_ids
|
|
buffer["ids"]["output"] = output_ids
|
|
if output_text is not None:
|
|
buffer["llm_output_text"] = [output_text]
|
|
if stream_output:
|
|
buffer["stream_output"] = True
|
|
if native_duplex:
|
|
buffer["native_duplex"] = True
|
|
return buffer
|
|
|
|
|
|
def set_ref_audio(buffer: DuplexIntermediateBuffer, waveform: object, sample_rate_hz: int) -> None:
|
|
buffer.setdefault("codes", {})["ref"] = waveform
|
|
buffer.setdefault("meta", {})["ref_audio_sr"] = int(sample_rate_hz)
|
|
|
|
|
|
def set_tts_handoff(buffer: DuplexIntermediateBuffer, token_ids: object | None, hidden_states: object | None) -> None:
|
|
"""Store the AR-to-TTS handoff used by the full-duplex stage bridge."""
|
|
if token_ids is not None:
|
|
buffer.setdefault("ids", {})["tts"] = token_ids
|
|
if hidden_states is not None:
|
|
buffer.setdefault("hidden_states", {})["tts"] = hidden_states
|
|
|
|
|
|
def get_tts_handoff(info: dict[str, object]) -> tuple[object | None, object | None]:
|
|
"""Read the canonical handoff, including the legacy flat aliases."""
|
|
ids_info = info.get("ids")
|
|
hidden_info = info.get("hidden_states")
|
|
token_ids = ids_info.get("tts") if isinstance(ids_info, dict) else None
|
|
hidden_states = hidden_info.get("tts") if isinstance(hidden_info, dict) else None
|
|
return (
|
|
info.get("tts_token_ids") if token_ids is None else token_ids,
|
|
info.get("tts_hidden_states") if hidden_states is None else hidden_states,
|
|
)
|
|
|
|
|
|
def get_stream_request_key(info: dict[str, object]) -> str:
|
|
key = info.get("global_request_id") or info.get("request_id") or info.get("_omni_req_id")
|
|
if isinstance(key, (list, tuple)):
|
|
key = key[0] if key else None
|
|
if isinstance(key, bytes):
|
|
key = key.decode("utf-8", errors="replace")
|
|
if key is None:
|
|
raise ValueError(
|
|
"Duplex streaming handoff requires a stable request id; "
|
|
"expected global_request_id, request_id, or _omni_req_id."
|
|
)
|
|
return str(key)
|