diff --git a/runtime/src/speech_to_speech/LLM/README.md b/runtime/src/speech_to_speech/LLM/README.md deleted file mode 100644 index 2a5a9d68dd2693a319ef24e9a260e94eda143096..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/LLM/README.md +++ /dev/null @@ -1,139 +0,0 @@ -# LLM Summary - -## Available LLM backends (`--llm_backend`) - -Runtime-supported values in `s2s_pipeline.py`: - -- `transformers` → `language_model.py` (Transformers backend) -- `mlx-lm` → `language_model.py` (MLX backend) -- `responses-api` → `responses_api_language_model.py` - -## Usage - -### 1) Transformers (`--llm_backend transformers`) - -- Handler: `LanguageModelHandler` -- Typical use: local GPU/CPU inference using Hugging Face Transformers -- Backend-specific args prefix: `--llm_*` -- Shared args (from base): `--model_name`, `--chat_size`, `--init_chat_prompt`, `--enable_lang_prompt` - -```bash -python s2s_pipeline.py \ - --llm_backend transformers \ - --model_name Qwen/Qwen3-4B-Instruct-2507 \ - --llm_device cuda \ - --llm_torch_dtype float16 \ - --llm_gen_max_new_tokens 128 -``` - -Common options: -- `--llm_gen_min_new_tokens` -- `--llm_gen_temperature` -- `--llm_gen_do_sample` -- `--chat_size` -- `--init_chat_prompt` - -### 2) MLX-LM (`--llm_backend mlx-lm`) - -- Handler: `LanguageModelHandler` -- Typical use: Apple Silicon local inference -- Backend-specific args prefix: same as Transformers (`--llm_*`) - -```bash -python s2s_pipeline.py \ - --llm_backend mlx-lm \ - --model_name mlx-community/Qwen3-4B-Instruct-2507-bf16 \ - --llm_device mps \ - --llm_gen_max_new_tokens 128 -``` - -Common options: -- `--llm_gen_temperature` -- `--llm_gen_do_sample` -- `--chat_size` -- `--init_chat_prompt` - -### 3) OpenAI-compatible API (`--llm_backend responses-api`) - -- Handler: `ResponsesApiModelHandler` -- Typical use: remote model serving via OpenAI-compatible endpoints -- Backend-specific args prefix: `--responses_api_*` -- Shared args (from base): `--model_name`, `--chat_size`, `--init_chat_prompt`, `--enable_lang_prompt` - -```bash -python s2s_pipeline.py \ - --llm_backend responses-api \ - --model_name gpt-5.4-mini \ - --responses_api_api_key YOUR_API_KEY \ - --responses_api_base_url https://api.example.com/v1 \ - --responses_api_stream true -``` - -Common options: -- `--chat_size` -- `--init_chat_prompt` -- `--user_role` - -## LLM Behavior - -When STT is set to language auto-detection (`--language auto`), LLM handlers can receive `(text, language_code)` and prepend a language control instruction like: - -- `Please reply to my message in .` - -This helps the assistant respond in the detected language. The behavior is opt-in via `--enable_lang_prompt` (shared across all backends); it defaults to `False`. - -## Setup - -### CUDA setup - -```bash -python s2s_pipeline.py \ - --llm_backend transformers \ - --model_name microsoft/Phi-3-mini-4k-instruct -``` - -### Local Mac setup - -```bash -python s2s_pipeline.py \ - --local_mac_optimal_settings \ - --model_name mlx-community/Qwen3-4B-Instruct-2507-bf16 -``` - -`--local_mac_optimal_settings` already sets `--llm_backend mlx-lm` and will default the model to `mlx-community/Qwen3-4B-Instruct-2507-bf16` if not overridden. - -### Realtime (OpenAI-compatible) setup - -Run the server in realtime mode, then connect with the realtime client: - -```bash -# 1. Start the pipeline in realtime mode -python s2s_pipeline.py \ - --mode realtime \ - --llm_backend mlx-lm \ - --model_name mlx-community/Qwen3-4B-Instruct-2507-bf16 \ - --ws_host 0.0.0.0 \ - --ws_port 8765 - -# 2. Connect with the realtime client -python listen_and_play_realtime.py --host 127.0.0.1 --port 8765 -``` - -Or with `--local_mac_optimal_settings` on Apple Silicon: - -```bash -python s2s_pipeline.py \ - --local_mac_optimal_settings \ - --mode realtime \ - --ws_host 0.0.0.0 \ - --ws_port 8765 -``` - -### Remote API setup - -```bash -python s2s_pipeline.py \ - --llm_backend responses-api \ - --model_name gpt-5.4-mini \ - --responses_api_api_key YOUR_API_KEY -``` diff --git a/runtime/src/speech_to_speech/LLM/__init__.py b/runtime/src/speech_to_speech/LLM/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/runtime/src/speech_to_speech/LLM/audio_input_notifier.py b/runtime/src/speech_to_speech/LLM/audio_input_notifier.py deleted file mode 100644 index 0733bbc90cc11926fe8be980ed7185cde49751dc..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/LLM/audio_input_notifier.py +++ /dev/null @@ -1,82 +0,0 @@ -from __future__ import annotations - -import logging -from collections.abc import Iterator -from queue import Queue -from threading import Event - -from speech_to_speech.api.openai_realtime.runtime_config import RuntimeConfig -from speech_to_speech.baseHandler import BaseHandler -from speech_to_speech.pipeline.events import AudioInputCompletedEvent -from speech_to_speech.pipeline.handler_types import LLMIn -from speech_to_speech.pipeline.messages import GenerateResponseRequest, VADAudio -from speech_to_speech.pipeline.queue_types import TextEventItem -from speech_to_speech.pipeline.speculative_turns import SpeculativeTurnTracker - -logger = logging.getLogger(__name__) - - -class AudioInputNotifier(BaseHandler[VADAudio, LLMIn]): - """Bridge final VAD audio directly into the LLM stage.""" - - def setup( - self, - runtime_config: RuntimeConfig | None = None, - should_listen: Event | None = None, - sample_rate: int = 16000, - speculative_turns: SpeculativeTurnTracker | None = None, - text_output_queue: Queue[TextEventItem] | None = None, - ) -> None: - self.runtime_config = runtime_config - self.should_listen = should_listen - self.sample_rate = sample_rate - self.speculative_turns = speculative_turns - self.text_output_queue = text_output_queue - - def should_process_input(self, item: VADAudio) -> bool: - if item.mode == "progressive": - return False - if self.speculative_turns is None or item.turn_id is None or item.turn_revision is None: - return True - return self.speculative_turns.is_latest_after_pending_reopen(item.turn_id, item.turn_revision) - - def process(self, vad_audio: VADAudio) -> Iterator[LLMIn]: - audio_duration_s = len(vad_audio.audio) / self.sample_rate if self.sample_rate else 0.0 - logger.info( - "Audio input completed: %.3fs turn=%s rev=%s", - audio_duration_s, - vad_audio.turn_id, - vad_audio.turn_revision, - ) - - # Realtime mode has no setup-level RuntimeConfig. Route the completed - # audio through RealtimeService so response lifecycle and usage - # bookkeeping happen before the LLM request is queued. - if self.runtime_config is None and self.text_output_queue is not None: - self.text_output_queue.put( - AudioInputCompletedEvent( - audio=vad_audio.audio, - audio_sample_rate=self.sample_rate, - audio_duration_s=audio_duration_s, - turn_id=vad_audio.turn_id, - turn_revision=vad_audio.turn_revision, - speech_stopped_at_s=vad_audio.created_at_s, - ) - ) - return - - runtime_config = vad_audio.runtime_config or self.runtime_config - if runtime_config is None: - logger.error("AudioInputNotifier received audio without RuntimeConfig; dropping LLM request") - if self.should_listen is not None: - self.should_listen.set() - return - - yield GenerateResponseRequest( - runtime_config=runtime_config, - audio=vad_audio.audio, - audio_sample_rate=self.sample_rate, - turn_id=vad_audio.turn_id, - turn_revision=vad_audio.turn_revision, - speech_stopped_at_s=vad_audio.created_at_s, - ) diff --git a/runtime/src/speech_to_speech/LLM/base_openai_compatible_language_model.py b/runtime/src/speech_to_speech/LLM/base_openai_compatible_language_model.py deleted file mode 100644 index 638834cf0d99bc42d9874154734bc4bdad2dcdc4..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/LLM/base_openai_compatible_language_model.py +++ /dev/null @@ -1,833 +0,0 @@ -from __future__ import annotations - -import base64 -import io -import ipaddress -import logging -import os -import wave -from abc import ABC, abstractmethod -from collections.abc import Callable, Generator, Iterator -from typing import Any, Literal, Optional -from urllib.parse import urlparse - -import httpx -import numpy as np -from nltk import sent_tokenize -from openai import OpenAI -from openai.types.realtime.conversation_item import ( - RealtimeConversationItemAssistantMessage, - RealtimeConversationItemFunctionCall, -) -from openai.types.realtime.realtime_conversation_item_assistant_message import ( - Content as AssistantContent, -) -from openai.types.responses import ResponseFunctionToolCall -from pydantic import BaseModel, ConfigDict, Field - -from speech_to_speech.baseHandler import BaseHandler -from speech_to_speech.LLM.chat import ( - Chat, - ChatItemError, - SupportedItem, - build_active_chat, - make_system_message, - make_user_audio_message, - make_user_message, -) -from speech_to_speech.LLM.compaction_prompt import CompactGenerateFn, build_compactor -from speech_to_speech.LLM.text_prompt import build_text_system_prompt -from speech_to_speech.LLM.utils import remove_unspeechable, resolve_auto_language -from speech_to_speech.LLM.voice_prompt import build_voice_system_prompt -from speech_to_speech.pipeline.cancel_scope import CancelScope -from speech_to_speech.pipeline.handler_types import LLMIn, LLMOut -from speech_to_speech.pipeline.messages import ( - EndOfResponse, - LLMResponseChunk, - TokenUsage, -) -from speech_to_speech.pipeline.speculative_turns import SpeculativeTurnTracker -from speech_to_speech.utils.utils import is_out_of_band, response_wants_audio - -logger = logging.getLogger(__name__) - -# About 18–24 seconds of default SDK backoff before warmup fails. -WARMUP_MAX_RETRIES = 6 - - -# ── Normalised provider events ──────────────────────────────────────────────── -# Each backend's stream/response is mapped to this small vocabulary so the shared -# speech-pipeline logic (sentence batching, cancellation, history, token usage) -# lives in one place. Subclasses differ only in how they produce these events. - - -class TextDelta(BaseModel): - """Incremental assistant text. Always RAW (unfiltered); the base applies - ``remove_unspeechable`` for the audio path.""" - - text: str - - -class AssistantMessage(BaseModel): - """A complete assistant turn to write back to history.""" - - content: list[AssistantContent] - - -class ToolCall(BaseModel): - """A complete function tool call (``call_id`` / ``id`` already regenerated).""" - - item: ResponseFunctionToolCall - - -class Usage(BaseModel): - """Token accounting for the turn.""" - - input_tokens: int - output_tokens: int - - -ProviderEvent = TextDelta | AssistantMessage | ToolCall | Usage -SerializeFn = Callable[[Chat], Any] -RequestFn = Callable[[Any, dict[str, Any]], Any] -EventIteratorFn = Callable[[Any], Iterator[ProviderEvent]] - - -class _Turn(BaseModel): - """Per-request context threaded through generation (immutable for the turn).""" - - model_config = ConfigDict(arbitrary_types_allowed=True) - - language_code: Optional[str] - gen: int | None - runtime_config: Any - response: Any - turn_id: str | None - turn_revision: int | None - speech_stopped_at_s: float | None - wants_audio: bool - - -class _GenState(BaseModel): - """Mutable accumulators collected while consuming a turn's events.""" - - model_config = ConfigDict(arbitrary_types_allowed=True) - - tools: list[ResponseFunctionToolCall] = Field(default_factory=list) - pending: list[SupportedItem] = Field(default_factory=list) - recorded_item_ids: set[str] = Field(default_factory=set) - recorded_call_ids: set[str] = Field(default_factory=set) - clean_text: str = "" # filtered text, kept only for the debug log - input_tokens: int = 0 - output_tokens: int = 0 - - -class BaseOpenAICompatibleHandler(BaseHandler[LLMIn, LLMOut], ABC): - """Shared lifecycle for OpenAI-compatible LLM backends (Responses & Chat - Completions). - - Subclasses implement four hooks — :meth:`warmup`, - :meth:`_build_compaction_generate_fn`, :meth:`_serialize`, :meth:`_request`, - :meth:`_iter_events` and :meth:`_build_optional_kwargs` — and inherit the - request/response orchestration: speculative-turn gating, cancellation, - sentence batching, text-only vs audio handling, history write-back, token - usage, out-of-band handling and error termination. - """ - - # ── setup ───────────────────────────────────────────────────────────────── - - def setup( - self, - model_name: str = "gpt-5.4-mini", - device: str = "cuda", - gen_kwargs: dict[str, Any] = {}, - base_url: Optional[str] = None, - api_key: Optional[str] = None, - stream: bool = True, - user_role: str = "user", - cancel_scope: CancelScope | None = None, - speculative_turns: SpeculativeTurnTracker | None = None, - disable_thinking: bool = True, - reasoning_effort: Optional[str] = None, - request_timeout_s: float = 20.0, - stream_batch_sentences: int = 3, - enable_lang_prompt: bool = False, - compact_history: bool = False, - audio_max_tokens: int = 256, - audio_temperature: float = 0.0, - audio_content_type: Literal["input_audio", "audio_url"] = "input_audio", - audio_history_turns: int = 1, - **_kwargs: Any, - ) -> None: - self.cancel_scope = cancel_scope - self.speculative_turns = speculative_turns - self.model_name = model_name - self.stream = stream - self.stream_batch_sentences = max(1, stream_batch_sentences) - self.enable_lang_prompt = enable_lang_prompt - self.gen_kwargs = dict(gen_kwargs) - self.audio_max_tokens = audio_max_tokens - self.audio_temperature = audio_temperature - if audio_content_type not in {"input_audio", "audio_url"}: - raise ValueError("audio_content_type must be either 'input_audio' or 'audio_url'.") - self.audio_content_type = audio_content_type - self.audio_history_turns = max(0, audio_history_turns) - self.request_timeout_s = float(request_timeout_s) - self.request_timeout = httpx.Timeout( - self.request_timeout_s, - connect=min(10.0, self.request_timeout_s), - ) - - self.user_role = user_role - if ( - api_key is None - and not os.environ.get("OPENAI_API_KEY") - and base_url is not None - and self._is_local_base_url(base_url) - ): - api_key = "none" - self.client = OpenAI(api_key=api_key, base_url=base_url) - self._extra_body = self._build_extra_body(base_url, disable_thinking, reasoning_effort) - self.compactor = build_compactor(self._build_compaction_generate_fn()) if compact_history else None - self.warmup() - - @staticmethod - def _is_official_openai(base_url: Optional[str]) -> bool: - """Whether ``base_url`` points at the official OpenAI server. - - Normalises a trailing slash so ``https://api.openai.com/v1/`` is also - recognised; the official server rejects the provider-specific extra_body - keys we send to vLLM / the HF router. - """ - if base_url is None: - return False - return base_url.rstrip("/") == "https://api.openai.com/v1" - - @staticmethod - def _is_local_base_url(base_url: str) -> bool: - """Whether *base_url* points at localhost or a loopback IP address.""" - host = urlparse(base_url).hostname - if host is None: - return False - if host.rstrip(".").lower() == "localhost": - return True - try: - return ipaddress.ip_address(host).is_loopback - except ValueError: - return False - - @classmethod - def _build_extra_body( - cls, - base_url: Optional[str], - disable_thinking: bool, - reasoning_effort: Optional[str], - ) -> Optional[dict[str, Any]]: - """Build the provider-specific ``extra_body`` used to disable reasoning. - - Providers differ in how reasoning is turned off: vLLM/Qwen honour - ``chat_template_kwargs.enable_thinking=false``, while others (e.g. GLM via - the HF router) ignore that and require ``reasoning_effort='none'``. A - non-empty ``reasoning_effort`` therefore takes precedence; otherwise we fall - back to the chat-template flag. None of this applies to the official - OpenAI server, which rejects unknown extra_body keys. - """ - if base_url is None or cls._is_official_openai(base_url): - return None - if reasoning_effort: - return {"reasoning_effort": reasoning_effort} - if disable_thinking: - return {"chat_template_kwargs": {"enable_thinking": False}} - return None - - # ── subclass hooks ──────────────────────────────────────────────────────-- - - @abstractmethod - def warmup(self) -> None: - """Issue a cheap request so the model/connection is ready before serving.""" - ... - - @abstractmethod - def _build_compaction_generate_fn(self) -> CompactGenerateFn: - """Return a ``(system, user) -> text`` fn used to compact long histories.""" - ... - - @abstractmethod - def _serialize(self, active_chat: Chat) -> Any: - """Serialise the chat to the backend's request payload (input/messages).""" - ... - - @abstractmethod - def _request(self, api_input: Any, optional_kwargs: dict[str, Any], runtime_config: Any) -> Any: - """Issue the create() call and return the response or stream.""" - ... - - @abstractmethod - def _iter_stream_events(self, api_response: Any) -> Iterator[ProviderEvent]: - """Map a streaming response to normalised :data:`ProviderEvent`s.""" - ... - - @abstractmethod - def _iter_response_events(self, api_response: Any) -> Iterator[ProviderEvent]: - """Map a non-streaming response to normalised :data:`ProviderEvent`s.""" - ... - - def _iter_events(self, api_response: Any) -> Iterator[ProviderEvent]: - """Dispatch to the stream/non-stream mapper. ``self.stream`` is the single - source of truth (it set the request's ``stream=`` flag), so the response - type always matches it.""" - if self.stream: - yield from self._iter_stream_events(api_response) - else: - yield from self._iter_response_events(api_response) - - @abstractmethod - def _build_optional_kwargs(self, req_tools: Any, req_tool_choice: Any) -> dict[str, Any]: - """Build the per-request tools/tool_choice kwargs in the backend's shape.""" - ... - - # ── audio-input protocol hooks ─────────────────────────────────────────── - - def _serialize_audio(self, active_chat: Chat) -> Any: - """Serialize an audio turn using the selected backend's native protocol.""" - return self._serialize(active_chat) - - def _build_audio_optional_kwargs( - self, - response: Any, - req_tools: Any, - req_tool_choice: Any, - ) -> dict[str, Any]: - """Build audio request parameters in the selected backend's shape.""" - kwargs = self._build_optional_kwargs(req_tools, req_tool_choice) - max_tokens = getattr(response, "max_output_tokens", None) if response is not None else None - kwargs.setdefault("max_tokens", max_tokens or self.audio_max_tokens) - kwargs.setdefault("temperature", self.audio_temperature) - return kwargs - - def _request_audio(self, api_input: Any, optional_kwargs: dict[str, Any], runtime_config: Any) -> Any: - return self._request(api_input, optional_kwargs, runtime_config) - - def _iter_audio_events(self, api_response: Any) -> Iterator[ProviderEvent]: - yield from self._iter_events(api_response) - - @staticmethod - def _audio_to_wav_base64(audio: np.ndarray, sample_rate: int) -> str: - """Encode a mono 16-bit WAV payload without touching the filesystem.""" - audio_array = np.asarray(audio) - if audio_array.ndim > 1: - audio_array = np.mean(audio_array, axis=1) - if np.issubdtype(audio_array.dtype, np.floating): - pcm = (np.clip(audio_array, -1.0, 1.0) * 32767.0).astype(" bool: - return self.speculative_turns is None or self.speculative_turns.is_latest(turn_id, turn_revision) - - def _generation_is_stale(self, gen: int | None) -> bool: - return gen is not None and self.cancel_scope is not None and self.cancel_scope.is_stale(gen) - - def _turn_output_allowed(self, turn_id: str | None, turn_revision: int | None) -> bool: - if self.speculative_turns is None: - return True - return self.speculative_turns.is_latest_after_reopen_grace(turn_id, turn_revision) - - def _apply_config( - self, - chat: Chat, - instructions: Optional[str], - wants_audio: bool = True, - ) -> None: - if instructions: - builder = build_voice_system_prompt if wants_audio else build_text_system_prompt - full_instructions = builder(instructions) - chat.add_item(make_system_message(full_instructions)) - - # ── output helpers ──────────────────────────────────────────────────────-- - - def _chunk( - self, - turn: _Turn, - *, - text: str = "", - tools: list[ResponseFunctionToolCall] | None = None, - language_code: Optional[str] = None, - ) -> LLMResponseChunk: - return LLMResponseChunk( - text=text, - language_code=language_code if language_code is not None else turn.language_code, - tools=tools or [], - runtime_config=turn.runtime_config, - response=turn.response, - turn_id=turn.turn_id, - turn_revision=turn.turn_revision, - speech_stopped_at_s=turn.speech_stopped_at_s, - cancel_generation=turn.gen, - ) - - def _record_tool_call(self, state: _GenState, turn: _Turn, item: ResponseFunctionToolCall) -> Iterator[LLMOut]: - """Emit a tool call, persisting it (and any assistant text seen so far) - to history *before* it is forwarded to the client. - - The function_call must already exist in the conversation by the time the - client returns its ``function_call_output``; otherwise a fast client - races ahead of the deferred end-of-turn write-back and the output is - rejected ("No function_call with call_id ... found"), which makes the - model re-issue the same tool call. The call lands in ``_pending_tool_calls`` - (not serialized until its output pairs it), so eager recording is safe. - - Out-of-band turns never touch the default conversation, and a stale turn - records nothing (it is not forwarded to the client either).""" - state.tools.append(item) - fc_item = RealtimeConversationItemFunctionCall( - type="function_call", - name=item.name, - arguments=item.arguments, - call_id=item.call_id, - id=item.id, - status=item.status, - ) - if self._generation_is_stale(turn.gen) or not self._turn_output_allowed(turn.turn_id, turn.turn_revision): - logger.info("LLM generation cancelled (stale speculative turn)") - return - if not is_out_of_band(turn.response): - # Flush assistant text accumulated before this call first (so history - # order matches what the client received), then persist the call — - # all before the chunk leaves for the client. - chat = turn.runtime_config.chat - for pending_item in state.pending: - recorded = chat.add_item(pending_item) - if recorded.id is not None: - state.recorded_item_ids.add(recorded.id) - state.pending.clear() - recorded_call = chat.add_item(fc_item) - if recorded_call.id is not None: - state.recorded_item_ids.add(recorded_call.id) - state.recorded_call_ids.add(item.call_id) - yield self._chunk(turn, tools=[item]) - - # ── consumption ─────────────────────────────────────────────────────────-- - - def _consume_streaming( - self, - events: Iterator[ProviderEvent], - state: _GenState, - turn: _Turn, - ) -> Generator[LLMOut, None, bool]: - cancelled = False - printable_text = "" - sentence_batch: list[str] = [] - - def _flush(batch: list[str]) -> Iterator[LLMOut]: - if not batch: - return - if not self._turn_output_allowed(turn.turn_id, turn.turn_revision): - logger.info("LLM generation cancelled (stale speculative turn)") - return - yield self._chunk(turn, text=" ".join(batch)) - - for event in events: - if self._generation_is_stale(turn.gen) or not self._turn_is_latest(turn.turn_id, turn.turn_revision): - logger.info("LLM generation cancelled (interruption)") - cancelled = True - break - - if isinstance(event, Usage): - state.input_tokens = event.input_tokens - state.output_tokens = event.output_tokens - elif isinstance(event, AssistantMessage): - state.pending.append( - RealtimeConversationItemAssistantMessage(type="message", role="assistant", content=event.content) - ) - elif isinstance(event, ToolCall): - # Flush any pending spoken text before emitting the tool call. - if printable_text.strip(): - sentence_batch.append(printable_text.strip()) - printable_text = "" - if sentence_batch: - if not self._turn_output_allowed(turn.turn_id, turn.turn_revision): - logger.info("LLM generation cancelled (stale speculative turn)") - cancelled = True - break - yield from _flush(sentence_batch) - sentence_batch = [] - yield from self._record_tool_call(state, turn, event.item) - elif isinstance(event, TextDelta): - if not turn.wants_audio: - # Text-only: forward verbatim. Keep every character (no - # remove_unspeechable, which strips TTS-unfriendly symbols) and - # don't sentence-split (sent_tokenize collapses newlines/markdown). - state.clean_text += event.text - if event.text: - if not self._turn_output_allowed(turn.turn_id, turn.turn_revision): - logger.info("LLM generation cancelled (stale speculative turn)") - cancelled = True - break - yield self._chunk(turn, text=event.text) - continue - new_text = remove_unspeechable(event.text) - state.clean_text += new_text - printable_text += new_text - sentences = sent_tokenize(printable_text) - if len(sentences) > 1: - for s in sentences[:-1]: - sentence_batch.append(s) - if len(sentence_batch) >= self.stream_batch_sentences: - if not self._turn_output_allowed(turn.turn_id, turn.turn_revision): - logger.info("LLM generation cancelled (stale speculative turn)") - cancelled = True - break - yield from _flush(sentence_batch) - sentence_batch = [] - if cancelled: - break - printable_text = sentences[-1] - - if not cancelled: - if printable_text.strip(): - sentence_batch.append(printable_text.strip()) - if sentence_batch: - if self._generation_is_stale(turn.gen): - logger.info("LLM generation cancelled (interruption)") - else: - logger.debug(f"Clean text: {state.clean_text}") - yield from _flush(sentence_batch) - logger.info(f"Tools: {state.tools}") - return ( - not cancelled - and not self._generation_is_stale(turn.gen) - and self._turn_is_latest(turn.turn_id, turn.turn_revision) - and self._turn_output_allowed(turn.turn_id, turn.turn_revision) - ) - - def _consume_nonstreaming( - self, - events: Iterator[ProviderEvent], - state: _GenState, - turn: _Turn, - ) -> Generator[LLMOut, None, bool]: - if self._generation_is_stale(turn.gen) or not self._turn_is_latest(turn.turn_id, turn.turn_revision): - logger.info("LLM generation cancelled (interruption)") - return False - for event in events: - if isinstance(event, Usage): - state.input_tokens = event.input_tokens - state.output_tokens = event.output_tokens - elif isinstance(event, AssistantMessage): - state.pending.append( - RealtimeConversationItemAssistantMessage(type="message", role="assistant", content=event.content) - ) - elif isinstance(event, ToolCall): - yield from self._record_tool_call(state, turn, event.item) - elif isinstance(event, TextDelta): - # Text-only keeps every character verbatim; audio strips - # TTS-unfriendly symbols via remove_unspeechable. - spoken = event.text if not turn.wants_audio else remove_unspeechable(event.text) - state.clean_text += spoken - out = spoken if not turn.wants_audio else spoken.strip() - if ( - out - and not self._generation_is_stale(turn.gen) - and self._turn_output_allowed(turn.turn_id, turn.turn_revision) - ): - yield self._chunk(turn, text=out) - logger.debug(f"Clean text: {state.clean_text}") - logger.info(f"Tools: {state.tools}") - return ( - not self._generation_is_stale(turn.gen) - and self._turn_is_latest(turn.turn_id, turn.turn_revision) - and self._turn_output_allowed(turn.turn_id, turn.turn_revision) - ) - - # ── orchestration ───────────────────────────────────────────────────────── - - def _generate( - self, - active_chat: Chat, - original_chat: Chat, - turn: _Turn, - optional_kwargs: dict[str, Any], - *, - serialize_fn: SerializeFn | None = None, - request_fn: RequestFn | None = None, - event_iterator_fn: EventIteratorFn | None = None, - transactional_user_message_id: str | None = None, - history_commit_fn: Callable[[], None] | None = None, - ) -> Generator[LLMOut, None, bool]: - api_response: Any = None - state = _GenState() - error_message: str | None = None - generation_completed = False - history_committed = False - transaction_rolled_back = False - consumed_image_ids: set[str] = set() - - def rollback_transaction() -> None: - nonlocal transaction_rolled_back - if transactional_user_message_id is None or history_committed or transaction_rolled_back: - return - original_chat.rollback_generation( - transactional_user_message_id, - item_ids=state.recorded_item_ids, - call_ids=state.recorded_call_ids, - ) - transaction_rolled_back = True - - try: - try: - api_input = (serialize_fn or self._serialize)(active_chat) - # Images the model actually sees this turn; only these are stripped on - # write-back, so an image a fast client injects mid-generation for the - # next turn survives (it is not in this serialized snapshot). - consumed_image_ids = active_chat.image_message_ids() - if not api_input: - # Nothing to send: empty `instructions` and no `input` (in the response, - # the default conversation, or the out-of-band context). The provider - # would reject this; fail with a clear message instead of an opaque error. - error_message = "Cannot generate a response: no instructions and no input were provided." - else: - if request_fn is not None: - api_response = request_fn(api_input, optional_kwargs) - else: - api_response = self._request(api_input, optional_kwargs, turn.runtime_config) - if api_response is not None: - events = (event_iterator_fn or self._iter_events)(api_response) - if self.stream: - generation_completed = yield from self._consume_streaming(events, state, turn) - else: - generation_completed = yield from self._consume_nonstreaming(events, state, turn) - except httpx.ReadTimeout: - logger.warning( - "OpenAI API read timed out after %.1fs; ending the current response", - self.request_timeout_s, - ) - if not self._generation_is_stale(turn.gen) and self._turn_output_allowed( - turn.turn_id, turn.turn_revision - ): - # Canned apology carries no language_code (mirrors the prior handlers). - yield LLMResponseChunk( - text="Wow I'm a bit slow today, could you repeat that?", - runtime_config=turn.runtime_config, - response=turn.response, - turn_id=turn.turn_id, - turn_revision=turn.turn_revision, - speech_stopped_at_s=turn.speech_stopped_at_s, - cancel_generation=turn.gen, - ) - except Exception as exc: - # Any other generation failure must still terminate the response: record - # the error and fall through to the EndOfResponse below. Without this the - # exception would escape process() and no EndOfResponse would be emitted, - # leaving st.in_response stuck and locking every subsequent response. - logger.exception("LLM generation failed; ending the current response") - if error_message is None: - error_message = f"Language model generation failed: {exc}" - - can_commit = ( - error_message is None - and generation_completed - and not self._generation_is_stale(turn.gen) - and self._turn_is_latest(turn.turn_id, turn.turn_revision) - and self._turn_output_allowed(turn.turn_id, turn.turn_revision) - ) - if can_commit: - try: - # Out-of-band responses emit output and usage but never write back to the - # default conversation (their context was a throwaway chat). - if not is_out_of_band(turn.response): - # Tool calls (and any assistant text preceding them) were already - # written eagerly in _record_tool_call; only trailing items remain. - for item in state.pending: - recorded = original_chat.add_item(item) - if recorded.id is not None: - state.recorded_item_ids.add(recorded.id) - original_chat.strip_images(consumed_image_ids) - if history_commit_fn is not None: - history_commit_fn() - original_chat.trim_if_needed(self.compactor) - history_committed = True - except Exception as exc: - logger.exception("LLM history commit failed; rolling back the current response") - error_message = f"Language model history commit failed: {exc}" - - rollback_transaction() - if history_committed and (state.input_tokens or state.output_tokens): - yield TokenUsage( - input_tokens=state.input_tokens, - output_tokens=state.output_tokens, - turn_id=turn.turn_id, - turn_revision=turn.turn_revision, - ) - yield EndOfResponse( - turn_id=turn.turn_id, - turn_revision=turn.turn_revision, - cancel_generation=turn.gen, - error=error_message, - ) - return history_committed - finally: - if api_response is not None and hasattr(api_response, "close"): - try: - api_response.close() - except Exception: - pass - rollback_transaction() - - def _process_audio(self, request: LLMIn) -> Iterator[LLMOut]: - """Process an audio-input turn through the selected backend protocol.""" - assert request.audio is not None - runtime_config = request.runtime_config - response = request.response - turn_id = request.turn_id - turn_revision = request.turn_revision - speech_stopped_at_s = request.speech_stopped_at_s - if not self._turn_is_latest(turn_id, turn_revision): - logger.info("Skipping stale LLM request for turn=%s rev=%s", turn_id, turn_revision) - yield EndOfResponse(turn_id=turn_id, turn_revision=turn_revision) - return - - original_chat = runtime_config.chat - if is_out_of_band(response): - try: - active_chat = build_active_chat(original_chat, response) - except ChatItemError as exc: - logger.info("Out-of-band response rejected: %s", exc) - yield EndOfResponse(turn_id=turn_id, turn_revision=turn_revision, error=str(exc)) - return - else: - active_chat = original_chat.copy() - - language_code = request.language_code - instructions = ( - response.instructions if response and response.instructions else runtime_config.session.instructions - ) or "" - req_tools = response.tools if response and response.tools else runtime_config.session.tools - req_tool_choice = ( - response.tool_choice if response and response.tool_choice else runtime_config.session.tool_choice - ) - wants_audio = response_wants_audio(response) - self._apply_config(active_chat, instructions, wants_audio) - language_code, lang_name = resolve_auto_language(language_code) - if lang_name and self.enable_lang_prompt: - active_chat.add_item(make_user_message(f"Please reply to my message in {lang_name}.")) - - audio_b64 = self._audio_to_wav_base64(request.audio, request.audio_sample_rate) - audio_message = active_chat.add_item(make_user_audio_message(audio_b64)) - optional_kwargs = self._build_audio_optional_kwargs(response, req_tools, req_tool_choice) - - transactional_user_message_id: str | None = None - history_commit_fn: Callable[[], None] | None = None - if not is_out_of_band(response): - provisional_message = make_user_audio_message(audio_b64) - provisional_message.id = audio_message.id - original_chat.add_item(provisional_message) - assert provisional_message.id is not None - transactional_user_message_id = provisional_message.id - - def commit_audio_history() -> None: - original_chat.compact_audio_history(self.audio_history_turns) - - history_commit_fn = commit_audio_history - - # CancelScope.is_stale(gen) is checked when the stream iterator advances; a - # blocked read inside httpx cannot be aborted by cancel_scope.cancel() from - # the websocket router. Mitigations: request_timeout_s / ReadTimeout. - gen = self.cancel_scope.generation if self.cancel_scope else None - turn = _Turn( - language_code=language_code, - gen=gen, - runtime_config=runtime_config, - response=response, - turn_id=turn_id, - turn_revision=turn_revision, - speech_stopped_at_s=speech_stopped_at_s, - wants_audio=wants_audio, - ) - yield from self._generate( - active_chat, - original_chat, - turn, - optional_kwargs, - serialize_fn=self._serialize_audio, - request_fn=lambda api_input, kwargs: self._request_audio(api_input, kwargs, runtime_config), - event_iterator_fn=self._iter_audio_events, - transactional_user_message_id=transactional_user_message_id, - history_commit_fn=history_commit_fn, - ) - - def process(self, request: LLMIn) -> Iterator[LLMOut]: - """Process a language model request and yield LLMResponseChunks.""" - if request.audio is not None: - yield from self._process_audio(request) - return - - runtime_config = request.runtime_config - response = request.response - turn_id = request.turn_id - turn_revision = request.turn_revision - speech_stopped_at_s = request.speech_stopped_at_s - if not self._turn_is_latest(turn_id, turn_revision): - logger.info("Skipping stale LLM request for turn=%s rev=%s", turn_id, turn_revision) - yield EndOfResponse(turn_id=turn_id, turn_revision=turn_revision) - return - - original_chat = runtime_config.chat - if is_out_of_band(response): - try: - active_chat = build_active_chat(original_chat, response) - except ChatItemError as exc: - logger.info("Out-of-band response rejected: %s", exc) - yield EndOfResponse(turn_id=turn_id, turn_revision=turn_revision, error=str(exc)) - return - else: - active_chat = original_chat.copy() - language_code = request.language_code - instructions = ( - response.instructions if response and response.instructions else runtime_config.session.instructions - ) or "" - req_tools = response.tools if response and response.tools else runtime_config.session.tools - req_tool_choice = ( - response.tool_choice if response and response.tool_choice else runtime_config.session.tool_choice - ) - wants_audio = response_wants_audio(response) - self._apply_config(active_chat, instructions, wants_audio) - language_code, lang_name = resolve_auto_language(language_code) - if lang_name and self.enable_lang_prompt: - active_chat.add_item(make_user_message(f"Please reply to my message in {lang_name}.")) - - optional_kwargs = self._build_optional_kwargs(req_tools, req_tool_choice) - - # CancelScope.is_stale(gen) is checked when the stream iterator advances; a - # blocked read inside httpx cannot be aborted by cancel_scope.cancel() from - # the websocket router. Mitigations: request_timeout_s / ReadTimeout. - gen = self.cancel_scope.generation if self.cancel_scope else None - - turn = _Turn( - language_code=language_code, - gen=gen, - runtime_config=runtime_config, - response=response, - turn_id=turn_id, - turn_revision=turn_revision, - speech_stopped_at_s=speech_stopped_at_s, - wants_audio=wants_audio, - ) - yield from self._generate(active_chat, original_chat, turn, optional_kwargs) - - @property - def timing_log_level(self) -> int: - return logging.INFO - - def should_log_timing(self, output: LLMOut) -> bool: - return isinstance(output, LLMResponseChunk) and self.last_time > self.min_time_to_debug diff --git a/runtime/src/speech_to_speech/LLM/chat.py b/runtime/src/speech_to_speech/LLM/chat.py deleted file mode 100644 index 0f2a772be5c283d5b9a9a263e3572a660dba2992..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/LLM/chat.py +++ /dev/null @@ -1,835 +0,0 @@ -from __future__ import annotations - -import json -import logging -import threading -from collections.abc import Callable -from typing import Any, Literal, Union - -from openai.types.realtime import ConversationItem -from openai.types.realtime.conversation_item import ( - RealtimeConversationItemAssistantMessage, - RealtimeConversationItemFunctionCall, - RealtimeConversationItemFunctionCallOutput, - RealtimeConversationItemSystemMessage, - RealtimeConversationItemUserMessage, -) -from openai.types.realtime.realtime_conversation_item_assistant_message import ( - Content as AssistantContent, -) -from openai.types.realtime.realtime_conversation_item_system_message import Content as SystemContent -from openai.types.realtime.realtime_conversation_item_user_message import Content as UserContent -from openai.types.realtime.realtime_response_create_params import RealtimeResponseCreateParams -from openai.types.responses.response_input_image_param import ResponseInputImageParam -from openai.types.responses.response_input_message_content_list_param import ( - ResponseInputMessageContentListParam, -) -from openai.types.responses.response_input_param import ( - FunctionCallOutput, - ResponseFunctionToolCallParam, - ResponseInputItemParam, - ResponseInputParam, - ResponseOutputMessageParam, -) -from openai.types.responses.response_input_param import ( - Message as ResponseMessage, -) -from openai.types.responses.response_input_text_param import ResponseInputTextParam -from openai.types.responses.response_output_text_param import ResponseOutputTextParam -from pydantic import BaseModel - -from speech_to_speech.utils.utils import _generate_id - -logger = logging.getLogger(__name__) - -AUDIO_INPUT_HISTORY_PLACEHOLDER = "[User audio input]" - - -class ChatItemError(Exception): - """Raised when a conversation item fails validation in :meth:`Chat.add_item`.""" - - -class CompactionResult(BaseModel): - """Output of a :data:`CompactFn` summarization run.""" - - user_summary: str - assistant_summary: str - - -def _ensure_id(value: str | None, prefix: str) -> str: - if value is None: - return _generate_id(prefix) - if not value.startswith(f"{prefix}_"): - raise ChatItemError(f"ID must start with '{prefix}_', got {value!r}") - return value - - -SupportedItem = Union[ - RealtimeConversationItemSystemMessage, - RealtimeConversationItemUserMessage, - RealtimeConversationItemAssistantMessage, - RealtimeConversationItemFunctionCall, - RealtimeConversationItemFunctionCallOutput, -] - - -CompactFn = Callable[[ResponseInputParam], CompactionResult] - - -class Chat: - """Manages conversation history with bounded size to avoid OOM issues. - - The buffer stores ``ConversationItem`` objects (user messages, assistant - messages, function calls, function call outputs). System messages are - stored separately in ``init_chat_message`` and never placed in the buffer. - - History bounding is decided per ``add_item`` call via the ``compactor`` - argument: - - - ``compactor=None``: when the user-turn count exceeds ``size`` the oldest - complete turn is evicted in place. Synchronous, lossy, no LLM involvement. - - ``compactor=``: when ``size`` is exceeded, ``fn`` is invoked in a - background thread to summarize older turns into a single user/assistant - pair (with pending function calls preserved). Single-flight: while a - compaction is running, additional triggers are silently bypassed. - """ - - def __init__(self, size: int) -> None: - self.size = size - self.init_chat_message: RealtimeConversationItemSystemMessage | None = None - # ``size`` is the number of user turns to keep. When exceeded the - # oldest complete turn (everything up to the next user message) - # is evicted -- or, with a compactor, summarized in the background. - self.buffer: list[SupportedItem] = [] - self._pending_tool_calls: dict[str, RealtimeConversationItemFunctionCall] = {} - self._user_turn_count: int = 0 - - # All state mutations and serializations go through _lock. Public methods - # acquire it once; internal callers that already hold it use the - # ``_locked`` helpers, so no reentry is needed (regular Lock is safe). - self._lock = threading.Lock() - self._compact_in_flight: bool = False - self._compact_thread: threading.Thread | None = None - self._shutdown = threading.Event() - self._gen_counter = 0 - - # ── Internal mutators (caller holds _lock) ───────────────── - - def _evict_oldest_turn(self) -> None: - """Remove items from the front until the next user message boundary.""" - if not self.buffer: - return - first = self.buffer.pop(0) - if isinstance(first, RealtimeConversationItemUserMessage): - self._user_turn_count -= 1 - while self.buffer and not isinstance(self.buffer[0], RealtimeConversationItemUserMessage): - self.buffer.pop(0) - - def _has_call_id_in_buffer(self, call_id: str) -> bool: - for entry in self.buffer: - if isinstance(entry, RealtimeConversationItemFunctionCall) and entry.call_id == call_id: - return True - return False - - def _mark_call_completed( - self, call_id: str, status: Literal["completed", "incomplete", "in_progress"] | None = None - ) -> None: - """Set ``status`` to ``"completed"`` on the matching function_call.""" - for entry in self.buffer: - if isinstance(entry, RealtimeConversationItemFunctionCall) and entry.call_id == call_id: - entry.status = "completed" if status is None else status - return - - def append_tool_output(self, call_id: str, output_item: RealtimeConversationItemFunctionCallOutput) -> None: - """Append a ``function_call_output``, re-injecting its ``function_call`` if evicted. - - Also marks the paired ``function_call`` as ``"completed"`` if its - status was ``None``. - - Raises :class:`ChatItemError` if *call_id* is unknown. - """ - with self._lock: - self._append_tool_output_locked(call_id, output_item) - - def _append_tool_output_locked(self, call_id: str, output_item: RealtimeConversationItemFunctionCallOutput) -> None: - """Body of :meth:`append_tool_output`. Caller must hold ``_lock``.""" - if self._has_call_id_in_buffer(call_id): - self._pending_tool_calls.pop(call_id, None) - self._mark_call_completed(call_id, output_item.status) - self.buffer.append(output_item) - return - - if call_id in self._pending_tool_calls: - logger.info("Re-injecting evicted function_call for call_id=%s", call_id) - fc = self._pending_tool_calls.pop(call_id) - fc.status = "completed" if output_item.status is None else output_item.status - self.buffer.append(fc) - self.buffer.append(output_item) - return - - raise ChatItemError(f"No function_call with call_id '{call_id}' found in conversation history.") - - def init_chat(self, message: RealtimeConversationItemSystemMessage) -> None: - with self._lock: - self.init_chat_message = message - - def add_item(self, item: SupportedItem) -> SupportedItem: - """Validate and route a conversation item into the chat buffer. - - Does not enforce the soft size limit — call :meth:`trim_if_needed` - explicitly after each successful generation to evict or compact old - turns. A hard upper bound at ``2 * size`` is enforced inline as a - runaway-client safety net: if the user-turn count exceeds it, the - oldest complete turn is evicted (lossy, no compaction). - - Raises :class:`ChatItemError` if the item fails validation. - """ - with self._lock: - if isinstance(item, RealtimeConversationItemSystemMessage): - item.id = _ensure_id(item.id, "sys") - self.init_chat_message = item - logger.debug("Set system message via conversation item") - - elif isinstance(item, RealtimeConversationItemUserMessage): - item.id = _ensure_id(item.id, "msg") - item.content = [ - p - for p in item.content - if (p.type == "input_text" and p.text) - or (p.type == "input_image" and p.image_url) - or (p.type == "input_audio" and p.audio) - ] - if not item.content: - raise ChatItemError( - "Message has no supported content. Supported modalities: input_text, input_image, input_audio." - ) - self.buffer.append(item) - self._user_turn_count += 1 - logger.debug("Added user message to chat (%d parts)", len(item.content)) - - elif isinstance(item, RealtimeConversationItemAssistantMessage): - item.id = _ensure_id(item.id, "msg") - item.content = [p for p in item.content if p.type == "output_text" and p.text] - if not item.content: - return item - self.buffer.append(item) - logger.debug("Added assistant message to chat (%d parts)", len(item.content)) - - elif isinstance(item, RealtimeConversationItemFunctionCall): - item.id = _ensure_id(item.id, "fc") - item.call_id = _ensure_id(item.call_id, "call") - self._pending_tool_calls[item.call_id] = item - logger.debug("Added function_call to chat (call_id=%s)", item.call_id) - - elif isinstance(item, RealtimeConversationItemFunctionCallOutput): - item.id = _ensure_id(item.id, "fco") - self._append_tool_output_locked(item.call_id, item) - logger.debug("Added function_call_output to chat (call_id=%s)", item.call_id) - - else: - raise ChatItemError(f"Unsupported item type: {getattr(item, 'type', None)}") - - if self.size > 0 and self._user_turn_count > 2 * self.size: - logger.warning( - "Chat buffer exceeded hard cap (%d > 2 * size=%d); evicting oldest turn", - self._user_turn_count, - self.size, - ) - while self._user_turn_count > 2 * self.size: - self._evict_oldest_turn() - - return item - - def trim_if_needed(self, compactor: CompactFn | None = None) -> None: - """Enforce the size limit after a generation completes. Fires when - ``user_turn_count > size``. - - - ``compactor=None``: synchronous eviction of the oldest complete turn. - - ``compactor=``: launch a background compaction (single-flight). - - Call once after each successful generation, not inside :meth:`add_item`. - """ - with self._lock: - if self._user_turn_count <= self.size: - return - if compactor is not None: - self._maybe_trigger_compaction(compactor) - else: - while self._user_turn_count > self.size: - self._evict_oldest_turn() - - def replace_user_message_text(self, item_id: str, text: str) -> bool: - """Replace the text content of an existing user message. - - Used by speculative turn revisions: the conversation turn remains the - same, but the STT transcript is superseded by a transcription of a - longer raw-audio buffer. - """ - - with self._lock: - for item in self.buffer: - if not isinstance(item, RealtimeConversationItemUserMessage) or item.id != item_id: - continue - item.content = [UserContent(type="input_text", text=text)] - logger.debug("Replaced speculative user message %s", item_id) - return True - return False - - def remove_user_message(self, item_id: str) -> bool: - """Remove an existing user message from the bounded chat buffer.""" - - with self._lock: - for index, item in enumerate(self.buffer): - if not isinstance(item, RealtimeConversationItemUserMessage) or item.id != item_id: - continue - del self.buffer[index] - self._user_turn_count -= 1 - logger.debug("Removed speculative user message %s", item_id) - return True - return False - - def rollback_generation( - self, - user_message_id: str, - *, - item_ids: set[str], - call_ids: set[str], - ) -> None: - """Remove only the provisional state written by one failed generation. - - A tool output may be appended by a fast client while generation is still - streaming, so rollback matches both item IDs and tool ``call_id`` values. - Unrelated messages injected concurrently for a later turn are preserved. - """ - - with self._lock: - kept: list[SupportedItem] = [] - for item in self.buffer: - remove = item.id == user_message_id or item.id in item_ids - if isinstance(item, (RealtimeConversationItemFunctionCall, RealtimeConversationItemFunctionCallOutput)): - remove = remove or item.call_id in call_ids - if not remove: - kept.append(item) - self.buffer = kept - for call_id in call_ids: - self._pending_tool_calls.pop(call_id, None) - self._user_turn_count = sum(isinstance(item, RealtimeConversationItemUserMessage) for item in self.buffer) - logger.debug("Rolled back failed generation for user message %s", user_message_id) - - def compact_audio_history(self, max_audio_turns: int) -> None: - """Retain only the newest bounded set of audio turns. - - Older audio parts are replaced with a textual placeholder so the user - role and its paired assistant response remain valid in serialized - history. The newest turns keep their audio semantics for subsequent - Chat Completions requests. - """ - - with self._lock: - remaining = max(0, max_audio_turns) - for item in reversed(self.buffer): - if not isinstance(item, RealtimeConversationItemUserMessage): - continue - if not any(part.type == "input_audio" for part in item.content): - continue - if remaining: - remaining -= 1 - continue - replacement_added = False - compacted: list[UserContent] = [] - for part in item.content: - if part.type != "input_audio": - compacted.append(part) - elif not replacement_added: - compacted.append(UserContent(type="input_text", text=AUDIO_INPUT_HISTORY_PLACEHOLDER)) - replacement_added = True - item.content = compacted - - def to_responses_api_chat(self, items: list[SupportedItem] | None = None) -> ResponseInputParam: - """Serialize the chat (system prompt + buffer) for the OpenAI Responses API. - - If *items* is provided, serialize that slice instead of the live buffer - (used by the compaction snapshot). - """ - with self._lock: - return self._to_responses_api_chat_locked(items if items is not None else self.buffer) - - def _to_responses_api_chat_locked(self, items: list[SupportedItem]) -> ResponseInputParam: - """Body of :meth:`to_responses_api_chat`. Caller must hold ``_lock``.""" - buffer_items = list(items) - result: list[ResponseInputItemParam] = [] - if self.init_chat_message: - result.append( - ResponseMessage( - content=[ - ResponseInputTextParam(text=p.text or "A helpful AI assistant.", type="input_text") - for p in self.init_chat_message.content - ], - role="system", - type="message", - ) - ) - for item in buffer_items: - assert item.id is not None and item.id != "", f"item.id is {item.id}" - if isinstance(item, RealtimeConversationItemUserMessage): - content: ResponseInputMessageContentListParam = [] - audio_placeholder_added = False - for user_part in item.content: - if user_part.type == "input_text" and user_part.text is not None: - content.append(ResponseInputTextParam(text=user_part.text or "", type="input_text")) - elif user_part.type == "input_image" and user_part.image_url is not None: - img = ResponseInputImageParam(type="input_image", detail=user_part.detail or "auto") - if user_part.image_url is not None: - img["image_url"] = user_part.image_url - content.append(img) - elif user_part.type == "input_audio" and not audio_placeholder_added: - content.append(ResponseInputTextParam(text=AUDIO_INPUT_HISTORY_PLACEHOLDER, type="input_text")) - audio_placeholder_added = True - if content: - result.append(ResponseMessage(content=content, role="user", type="message")) - elif isinstance(item, RealtimeConversationItemAssistantMessage): - assistant_content: list[ResponseOutputTextParam] = [] - for assistant_part in item.content: - if assistant_part.type == "output_text" and assistant_part.text is not None: - assistant_content.append( - ResponseOutputTextParam(text=assistant_part.text, type="output_text", annotations=[]) - ) - if assistant_content: - result.append( - ResponseOutputMessageParam( - id=item.id, - content=assistant_content, - role="assistant", - status=item.status or "completed", - type="message", - ) - ) - elif isinstance(item, RealtimeConversationItemFunctionCall) and item.call_id is not None: - assert item.call_id is not None and item.call_id != "" - function_call = ResponseFunctionToolCallParam( - arguments=item.arguments, - call_id=item.call_id, - name=item.name, - type="function_call", - id=item.id, - ) - if item.id is not None: - function_call["id"] = item.id - if item.status is not None: - function_call["status"] = item.status - result.append(function_call) - elif isinstance(item, RealtimeConversationItemFunctionCallOutput): - function_call_output = FunctionCallOutput( - call_id=item.call_id, - output=item.output, - type="function_call_output", - ) - if item.id is not None: - function_call_output["id"] = item.id - if item.status is not None: - function_call_output["status"] = item.status - result.append(function_call_output) - return result - - def to_transformers_chat(self) -> list[dict[str, Any]]: - """Serialize the full chat for HuggingFace transformers ``apply_chat_template``. - - User messages with only text produce a plain string ``content`` value. - User messages containing images keep ``content`` as a list of dicts so - VLM pipelines can process them. - """ - with self._lock: - messages: list[TransformersChatMessage] = [] - if self.init_chat_message: - text = " ".join(p.text for p in self.init_chat_message.content if p.text) - messages.append(TransformersSystemMessage(content=text)) - for item in self.buffer: - if isinstance(item, RealtimeConversationItemUserMessage): - has_media = any(p.type in {"input_image", "input_audio"} for p in item.content) - if has_media: - messages.append( - TransformersUserMessage(content=[p.model_dump(exclude_none=True) for p in item.content]) - ) - else: - text = " ".join(p.text for p in item.content if p.type == "input_text" and p.text) - messages.append(TransformersUserMessage(content=text)) - elif isinstance(item, RealtimeConversationItemAssistantMessage): - text = " ".join(p.text for p in item.content if p.text) - messages.append(TransformersAssistantMessage(content=text)) - elif isinstance(item, RealtimeConversationItemFunctionCall): - assert item.call_id is not None and item.call_id != "" - args: Any = item.arguments - try: - args = json.loads(args) if isinstance(args, str) else args - except (json.JSONDecodeError, TypeError): - args = {} - messages.append( - TransformersFunctionCallMessage( - tool_calls=[ - TransformersToolCall( - id=item.call_id, - function=TransformersToolCallFunction(name=item.name, arguments=args), - ) - ] - ) - ) - elif isinstance(item, RealtimeConversationItemFunctionCallOutput): - name = "" - for prev in reversed(messages): - if isinstance(prev, TransformersFunctionCallMessage): - for tc in prev.tool_calls: - if tc.id == item.call_id: - name = tc.function.name - break - if name: - break - messages.append( - TransformersToolMessage( - tool_call_id=item.call_id, - name=name, - content=item.output, - ) - ) - return [m.model_dump() for m in messages] - - def copy(self) -> Chat: - """Return a shallow snapshot safe for concurrent read access.""" - with self._lock: - clone = Chat(self.size) - clone.init_chat_message = self.init_chat_message - clone.buffer = list(self.buffer) - clone._pending_tool_calls = dict(self._pending_tool_calls) - clone._user_turn_count = self._user_turn_count - return clone - - def reset(self) -> None: - """Clear all conversation state. Cancels any in-flight compaction splice.""" - with self._lock: - self._gen_counter += 1 - self._compact_in_flight = False - self.buffer = [] - self.init_chat_message = None - self._pending_tool_calls = {} - self._user_turn_count = 0 - - def close(self) -> None: - """Permanently shut down the chat. In-flight compaction splice is suppressed. - - The compaction worker (a daemon thread) is not joined: it may be blocked - in an LLM call. Process exit reaps it. - """ - self._shutdown.set() - with self._lock: - self._gen_counter += 1 - self._compact_in_flight = False - - def image_message_ids(self) -> set[str]: - """IDs of user messages currently carrying ``input_image`` content.""" - with self._lock: - return { - item.id - for item in self.buffer - if isinstance(item, RealtimeConversationItemUserMessage) - and item.id is not None - and any(p.type == "input_image" for p in item.content) - } - - def strip_images(self, only_ids: set[str] | None = None) -> None: - """Remove image content parts from user messages in the buffer. - - Called after appending the assistant response so images don't persist - across turns. With *only_ids*, strip only those message IDs — the images - the just-completed response actually consumed (captured before the - request was sent). This leaves intact an image a fast client injected - mid-generation for the *next* turn, which the current response never saw. - Without *only_ids*, every image is stripped. - """ - with self._lock: - for item in self.buffer: - if isinstance(item, RealtimeConversationItemUserMessage): - if only_ids is not None and item.id not in only_ids: - continue - item.content = [p for p in item.content if p.type != "input_image"] - - # ── Compaction internals ────────────────────────────────── - - def _snapshot_for_compaction( - self, - ) -> tuple[ResponseInputParam, set[str], int]: - """Compute the snapshot of items eligible for compaction. - - Caller must hold ``_lock``. Returns - ``(serialized_snapshot, marker_ids, n_turns)``. ``marker_ids`` - identifies the buffer items that may be removed when the splice runs. - Always leaves the most recent user turn untouched (it may be in-flight). - Returns an empty result if there are fewer than 2 compactable turns. - """ - n_turns = max(0, self._user_turn_count - 1) - if n_turns < 2: - return [], set(), n_turns - - # Slice up to (but not including) the (n_turns + 1)-th user message. - user_seen = 0 - end_idx = len(self.buffer) - for i, entry in enumerate(self.buffer): - if isinstance(entry, RealtimeConversationItemUserMessage): - user_seen += 1 - if user_seen == n_turns + 1: - end_idx = i - break - - items_to_compact = self.buffer[:end_idx] - marker_ids = {entry.id for entry in items_to_compact if entry.id is not None} - snapshot = self._to_responses_api_chat_locked(items=items_to_compact) - # Strip media parts so the summarizer doesn't have to handle them. - for raw in snapshot: - if not isinstance(raw, dict) or raw.get("role") != "user": - continue - msg: dict[str, Any] = raw # type: ignore[assignment] - content = msg.get("content") - if isinstance(content, list): - msg["content"] = [ - c for c in content if not (isinstance(c, dict) and c.get("type") in {"input_image", "input_audio"}) - ] - return snapshot, marker_ids, n_turns - - def _maybe_trigger_compaction(self, compactor: CompactFn) -> None: - """Start a background compaction worker. Bypass silently if one is running. - - Caller must hold ``_lock``. - """ - if self._shutdown.is_set() or self._compact_in_flight: - return - snapshot, marker_ids, n_turns = self._snapshot_for_compaction() - if n_turns < 2 or not marker_ids: - return - gen = self._gen_counter - self._compact_in_flight = True - thread = threading.Thread( - target=self._compact_worker, - args=(compactor, snapshot, marker_ids, gen), - daemon=True, - name="chat-compact", - ) - self._compact_thread = thread - logger.info( - "Chat compaction triggered: compacting %d turn(s) (%d item(s)), buffer size=%d", - n_turns, - len(marker_ids), - len(self.buffer), - ) - thread.start() - - def _compact_worker( - self, - compactor: CompactFn, - snapshot: ResponseInputParam, - marker_ids: set[str], - gen: int, - ) -> None: - """Worker thread entry point.""" - try: - if self._shutdown.is_set() or self._gen_counter != gen: - return - try: - result = compactor(snapshot) - except Exception: - logger.exception("Chat compaction failed; chat unchanged") - return - if not isinstance(result, CompactionResult): - logger.error("Compactor must return a CompactionResult, got %r", type(result).__name__) - return - if self._shutdown.is_set() or self._gen_counter != gen: - return - self._apply_compaction(result, marker_ids, gen) - finally: - # Don't clobber the flag if reset/close has advanced the gen. - with self._lock: - if self._gen_counter == gen: - self._compact_in_flight = False - - def _apply_compaction( - self, - result: CompactionResult, - marker_ids: set[str], - gen: int, - ) -> None: - """Splice the summary in front of items not consumed by compaction. - - FC/FCO pairing is left entirely to :meth:`add_item` / :meth:`append_tool_output`. - Compaction only drops items; it never inserts an FC into the buffer. - Pending FCs (no FCO yet) stay in ``_pending_tool_calls`` and will be - appended adjacent to their FCO when it arrives. - """ - with self._lock: - if self._shutdown.is_set() or self._gen_counter != gen: - return - # Keep FC if its FCO is outside the compacted range -- otherwise - # the FCO in `remaining` would be orphaned. - fco_call_ids_in_range = { - x.call_id - for x in self.buffer - if isinstance(x, RealtimeConversationItemFunctionCallOutput) and x.id in marker_ids - } - fc_ids_to_keep = { - x.id - for x in self.buffer - if x.id in marker_ids - and isinstance(x, RealtimeConversationItemFunctionCall) - and x.call_id not in fco_call_ids_in_range - } - drop_ids = marker_ids - fc_ids_to_keep - remaining = [x for x in self.buffer if x.id not in drop_ids] - - user_msg = make_user_message(result.user_summary) - user_msg.id = _generate_id("msg") - asst_msg = make_assistant_message(result.assistant_summary) - asst_msg.id = _generate_id("msg") - - self.buffer = [user_msg, asst_msg, *remaining] - self._user_turn_count = sum(1 for x in self.buffer if isinstance(x, RealtimeConversationItemUserMessage)) - logger.info( - "Chat compaction applied: buffer now %d item(s), %d user turn(s)", - len(self.buffer), - self._user_turn_count, - ) - - -# --------------------------------------------------------------------------- -# Transformers chat message models -# --------------------------------------------------------------------------- - - -class TransformersToolCallFunction(BaseModel): - name: str - arguments: dict[str, Any] - - -class TransformersToolCall(BaseModel): - type: Literal["function"] = "function" - id: str - function: TransformersToolCallFunction - - -class TransformersSystemMessage(BaseModel): - role: Literal["system"] = "system" - content: str - - -class TransformersUserMessage(BaseModel): - role: Literal["user"] = "user" - content: str | list[dict[str, Any]] - - -class TransformersAssistantMessage(BaseModel): - role: Literal["assistant"] = "assistant" - content: str - - -class TransformersFunctionCallMessage(BaseModel): - role: Literal["assistant"] = "assistant" - tool_calls: list[TransformersToolCall] - - -class TransformersToolMessage(BaseModel): - role: Literal["tool"] = "tool" - tool_call_id: str - name: str - content: str - - -TransformersChatMessage = Union[ - TransformersSystemMessage, - TransformersUserMessage, - TransformersAssistantMessage, - TransformersFunctionCallMessage, - TransformersToolMessage, -] - - -# --------------------------------------------------------------------------- -# Factory helpers -- hide verbose constructors behind simple calls -# --------------------------------------------------------------------------- - - -def make_user_message(text: str) -> RealtimeConversationItemUserMessage: - return RealtimeConversationItemUserMessage( - type="message", - role="user", - content=[UserContent(type="input_text", text=text)], - ) - - -def make_user_audio_message(audio_b64: str) -> RealtimeConversationItemUserMessage: - return RealtimeConversationItemUserMessage( - type="message", - role="user", - content=[UserContent(type="input_audio", audio=audio_b64)], - ) - - -def make_assistant_message(text: str) -> RealtimeConversationItemAssistantMessage: - return RealtimeConversationItemAssistantMessage( - type="message", - role="assistant", - content=[AssistantContent(type="output_text", text=text)], - ) - - -def make_system_message(text: str) -> RealtimeConversationItemSystemMessage: - return RealtimeConversationItemSystemMessage( - type="message", - role="system", - content=[SystemContent(type="input_text", text=text)], - ) - - -def add_supported_item(chat: Chat, item: ConversationItem) -> None: - """Narrow a protocol ``ConversationItem`` to a :data:`SupportedItem` and add it to *chat*. - - Raises :class:`ChatItemError` on validation failure or unsupported type. Shared - by the conversation handler (in-band item injection) and the language-model - handlers (seeding an out-of-band response's throwaway chat from ``response.input``). - """ - # call_id on function_call items must be client-supplied: it is referenced later by - # function_call_output items, so we cannot silently generate one here. - if isinstance(item, RealtimeConversationItemFunctionCall) and ( - item.call_id is None or not item.call_id.startswith("call_") - ): - raise ChatItemError("function_call item is missing a call_id. The call_id should start with 'call_'.") - - if isinstance( - item, - ( - RealtimeConversationItemSystemMessage, - RealtimeConversationItemUserMessage, - RealtimeConversationItemAssistantMessage, - RealtimeConversationItemFunctionCall, - RealtimeConversationItemFunctionCallOutput, - ), - ): - chat.add_item(item) - return - - raise ChatItemError(f"Unsupported item type: {getattr(item, 'type', None)}") - - -def build_active_chat(original_chat: Chat, response: RealtimeResponseCreateParams | None) -> Chat: - """Build the chat an *out-of-band* response generates against (caller ensures out-of-band). - - Mirrors the OpenAI realtime semantics for ``input``: - - - ``input is None`` -> a read-only **copy of the default conversation** (the - out-of-band response reads history but never commits back). - - ``input == []`` -> a **fresh, empty chat** (context cleared; only the - system prompt, added later by the handler, will be present). - - ``input == [...]`` -> a **fresh chat seeded** with those items. - - Raises :class:`ChatItemError` if an ``input`` item fails validation. - """ - if response is not None and response.input is not None: - fresh = Chat(original_chat.size) - for item in response.input: - add_supported_item(fresh, item) - return fresh - return original_chat.copy() diff --git a/runtime/src/speech_to_speech/LLM/chat_completions_language_model.py b/runtime/src/speech_to_speech/LLM/chat_completions_language_model.py deleted file mode 100644 index d425e7145b0a952e0353a9222d2441ba52b99fd1..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/LLM/chat_completions_language_model.py +++ /dev/null @@ -1,360 +0,0 @@ -from __future__ import annotations - -import json -import logging -import time -from collections.abc import Iterator -from typing import Any, cast - -from openai import Stream -from openai.types.chat import ( - ChatCompletionChunk, - ChatCompletionContentPartImageParam, - ChatCompletionContentPartParam, - ChatCompletionContentPartTextParam, - ChatCompletionNamedToolChoiceParam, - ChatCompletionToolChoiceOptionParam, - ChatCompletionToolParam, -) -from openai.types.chat.chat_completion_content_part_image_param import ImageURL -from openai.types.chat.chat_completion_named_tool_choice_param import Function as NamedToolChoiceFunction -from openai.types.realtime.realtime_conversation_item_assistant_message import ( - Content as AssistantContent, -) -from openai.types.responses import ResponseFunctionToolCall -from openai.types.shared_params import FunctionDefinition - -from speech_to_speech.LLM.base_openai_compatible_language_model import ( - WARMUP_MAX_RETRIES, - AssistantMessage, - BaseOpenAICompatibleHandler, - ProviderEvent, - TextDelta, - ToolCall, - Usage, -) -from speech_to_speech.LLM.chat import Chat -from speech_to_speech.LLM.compaction_prompt import CompactGenerateFn -from speech_to_speech.utils.utils import _generate_id - -logger = logging.getLogger(__name__) - - -def _to_chat_tools(req_tools: Any) -> list[ChatCompletionToolParam] | None: - """Convert Responses-API function tools to Chat-Completions tool format. - - Responses tools are flat ``{type:"function", name, description, parameters}``; - Chat Completions nests them under a ``function`` key. Items already in the - nested form (or non-function tools) are passed through untouched. - """ - if not req_tools: - return None - chat_tools: list[ChatCompletionToolParam] = [] - for t in req_tools: - d = t if isinstance(t, dict) else t.model_dump(exclude_none=True) - if d.get("type") == "function" and "function" not in d: - fn = FunctionDefinition(name=d["name"]) - if d.get("description") is not None: - fn["description"] = d["description"] - if d.get("parameters") is not None: - fn["parameters"] = d["parameters"] - chat_tools.append(ChatCompletionToolParam(type="function", function=fn)) - else: - chat_tools.append(cast("ChatCompletionToolParam", d)) - return chat_tools - - -def _to_chat_tool_choice(tool_choice: Any) -> ChatCompletionToolChoiceOptionParam: - """Convert a Responses-API tool_choice to Chat-Completions form. - - The string forms ("auto"/"required"/"none") are identical across both APIs; - only the forced-function object differs (flat ``name`` vs nested ``function``). - """ - if isinstance(tool_choice, dict) and tool_choice.get("type") == "function" and "name" in tool_choice: - return ChatCompletionNamedToolChoiceParam( - type="function", function=NamedToolChoiceFunction(name=tool_choice["name"]) - ) - if tool_choice is not None and not isinstance(tool_choice, (str, dict)): - d = tool_choice.model_dump(exclude_none=True) - if d.get("type") == "function" and "name" in d: - return ChatCompletionNamedToolChoiceParam(type="function", function=NamedToolChoiceFunction(name=d["name"])) - return cast("ChatCompletionToolChoiceOptionParam", d) - return cast("ChatCompletionToolChoiceOptionParam", tool_choice) - - -def _build_chat_optional_kwargs(req_tools: Any, req_tool_choice: Any) -> dict[str, Any]: - """Build Chat Completions tool arguments from the Responses-style session config.""" - optional_kwargs: dict[str, Any] = {} - chat_tools = _to_chat_tools(req_tools) - if chat_tools is not None: - optional_kwargs["tools"] = chat_tools - if req_tool_choice is not None: - optional_kwargs["tool_choice"] = _to_chat_tool_choice(req_tool_choice) - return optional_kwargs - - -def _to_chat_content_part( - part: dict[str, Any], - audio_content_type: str = "input_audio", -) -> ChatCompletionContentPartParam: - """Convert one Realtime/transformers content part to Chat Completions shape.""" - ptype = part.get("type") - if ptype == "input_text": - return ChatCompletionContentPartTextParam(type="text", text=part.get("text") or "") - if ptype == "input_image": - raw_url: Any = part.get("image_url") - if isinstance(raw_url, dict): - image_url = cast("ImageURL", raw_url) - else: - image_url = ImageURL(url=raw_url) - detail = part.get("detail") - if detail is not None: - image_url["detail"] = detail - return ChatCompletionContentPartImageParam(type="image_url", image_url=image_url) - if ptype == "input_audio": - audio_b64 = part.get("audio") or "" - if audio_content_type == "audio_url": - return cast( - "ChatCompletionContentPartParam", - { - "type": "audio_url", - "audio_url": { - "url": f"data:audio/wav;base64,{audio_b64}", - }, - }, - ) - return cast( - "ChatCompletionContentPartParam", - { - "type": "input_audio", - "input_audio": { - "data": audio_b64, - "format": "wav", - }, - }, - ) - return cast("ChatCompletionContentPartParam", part) - - -def _chat_messages( - chat: Chat, - audio_content_type: str = "input_audio", -) -> list[dict[str, Any]]: - """Serialise chat history, including media and tool messages, for Chat Completions.""" - messages = chat.to_transformers_chat() - for message in messages: - for tool_call in message.get("tool_calls") or []: - fn = tool_call.get("function") - if fn is not None and not isinstance(fn.get("arguments"), str): - fn["arguments"] = json.dumps(fn.get("arguments") or {}, ensure_ascii=False) - content = message.get("content") - if isinstance(content, list): - message["content"] = [ - _to_chat_content_part(part, audio_content_type=audio_content_type) for part in content - ] - if message.get("role") == "tool": - message.pop("name", None) - return messages - - -def _request_chat_completions( - *, - client: Any, - model_name: str, - messages: list[dict[str, Any]], - stream: bool, - extra_body: dict[str, Any] | None, - timeout: Any, - optional_kwargs: dict[str, Any], -) -> Any: - """Issue a Chat Completions request with consistent streaming usage accounting.""" - create_kwargs = dict(optional_kwargs) - if stream: - create_kwargs["stream_options"] = {"include_usage": True} - return client.chat.completions.create( - model=model_name, - messages=messages, - stream=stream, - extra_body=extra_body, - timeout=timeout, - **create_kwargs, - ) - - -def _tool_calls_from_accum(tool_accum: dict[int, dict[str, str]]) -> Iterator[ToolCall]: - """Turn accumulated Chat Completions tool deltas into normalized tool-call events.""" - for index in sorted(tool_accum): - entry = tool_accum[index] - if not entry["name"]: - continue - yield ToolCall( - item=ResponseFunctionToolCall( - type="function_call", - name=entry["name"], - arguments=entry["args"] or "{}", - call_id=_generate_id("call"), - id=_generate_id("fc"), - status="completed", - ) - ) - - -def _iter_chat_stream_events(api_response: Stream[ChatCompletionChunk]) -> Iterator[ProviderEvent]: - """Normalize a streaming Chat Completions response.""" - tool_accum: dict[int, dict[str, str]] = {} - usage: Usage | None = None - raw_text = "" - for chunk in api_response: - if chunk.usage is not None: - usage = Usage( - input_tokens=chunk.usage.prompt_tokens or 0, - output_tokens=chunk.usage.completion_tokens or 0, - ) - if not chunk.choices: - continue - delta = chunk.choices[0].delta - if delta.tool_calls: - for tool_call in delta.tool_calls: - entry = tool_accum.setdefault(tool_call.index, {"name": "", "args": "", "id": ""}) - if tool_call.id: - entry["id"] = tool_call.id - if tool_call.function is not None: - if tool_call.function.name: - entry["name"] = tool_call.function.name - if tool_call.function.arguments: - entry["args"] += tool_call.function.arguments - text_piece = delta.content or getattr(delta, "refusal", None) - if text_piece: - raw_text += text_piece - yield TextDelta(text=text_piece) - - if raw_text.strip(): - yield AssistantMessage(content=[AssistantContent(type="output_text", text=raw_text)]) - yield from _tool_calls_from_accum(tool_accum) - if usage is not None: - yield usage - - -def _iter_chat_response_events(api_response: Any) -> Iterator[ProviderEvent]: - """Normalize a non-streaming Chat Completions response.""" - usage = api_response.usage - if usage: - yield Usage(input_tokens=usage.prompt_tokens or 0, output_tokens=usage.completion_tokens or 0) - message = api_response.choices[0].message if api_response.choices else None - if message is None: - return - raw_content = message.content or getattr(message, "refusal", None) - if raw_content: - yield AssistantMessage(content=[AssistantContent(type="output_text", text=raw_content)]) - yield TextDelta(text=raw_content) - tool_accum: dict[int, dict[str, str]] = {} - for tool_call in message.tool_calls or []: - tool_accum[len(tool_accum)] = { - "name": tool_call.function.name or "", - "args": tool_call.function.arguments or "", - "id": tool_call.id or "", - } - yield from _tool_calls_from_accum(tool_accum) - - -class ChatCompletionsApiModelHandler(BaseOpenAICompatibleHandler): - """LLM handler that talks to an OpenAI-compatible ``/v1/chat/completions`` server. - - Functionally mirrors :class:`ResponsesApiModelHandler` but uses the mature - Chat Completions streaming tool-call protocol (``choices[].delta.tool_calls``) - instead of ``/v1/responses``. This is the robust path for vLLM + Qwen tool - calling. The conversation is serialised with :meth:`Chat.to_transformers_chat`, - which already emits OpenAI chat messages including ``tool_calls``/``tool`` roles. - """ - - def warmup(self) -> None: - logger.info(f"Warming up {self.__class__.__name__}") - start = time.time() - self.client.with_options(max_retries=WARMUP_MAX_RETRIES).chat.completions.create( - model=self.model_name, - messages=[ - {"role": "system", "content": "You are a helpful assistant"}, - {"role": "user", "content": "Hello"}, - ], - extra_body=self._extra_body, - timeout=self.request_timeout, - ) - end = time.time() - logger.info(f"{self.__class__.__name__}: warmed up! time: {(end - start):.3f} s") - - def _build_compaction_generate_fn(self) -> CompactGenerateFn: - """Return a generate fn that calls Chat Completions for compaction.""" - client = self.client - model_name = self.model_name - timeout = self.request_timeout - extra_body = self._extra_body - - def generate(system: str, user: str) -> str: - response = client.chat.completions.create( - model=model_name, - messages=[ - {"role": "system", "content": system}, - {"role": "user", "content": user}, - ], - extra_body=extra_body, - timeout=timeout, - ) - return response.choices[0].message.content or "" - - return generate - - @staticmethod - def _to_chat_content_part( - part: dict[str, Any], - audio_content_type: str = "input_audio", - ) -> ChatCompletionContentPartParam: - return _to_chat_content_part(part, audio_content_type=audio_content_type) - - @classmethod - def _chat_messages( - cls, - chat: Chat, - audio_content_type: str = "input_audio", - ) -> list[dict[str, Any]]: - return _chat_messages(chat, audio_content_type=audio_content_type) - - # ── base hooks ──────────────────────────────────────────────────────────-- - - def _serialize(self, active_chat: Chat) -> list[dict[str, Any]]: - return self._chat_messages(active_chat, audio_content_type=self.audio_content_type) - - def _build_optional_kwargs(self, req_tools: Any, req_tool_choice: Any) -> dict[str, Any]: - return _build_chat_optional_kwargs(req_tools, req_tool_choice) - - def _request( - self, - api_input: list[dict[str, Any]], - optional_kwargs: dict[str, Any], - runtime_config: Any, - ) -> Any: - extra_body = self._extra_body - session_reasoning_effort = getattr(runtime_config.session, "reasoning_effort", None) - if session_reasoning_effort: - extra_body = {"reasoning_effort": session_reasoning_effort} - return _request_chat_completions( - client=self.client, - model_name=self.model_name, - messages=api_input, - stream=self.stream, - extra_body=extra_body, - timeout=self.request_timeout, - optional_kwargs=optional_kwargs, - ) - - def _iter_stream_events(self, api_response: Stream[ChatCompletionChunk]) -> Iterator[ProviderEvent]: - yield from _iter_chat_stream_events(api_response) - - def _iter_response_events(self, api_response: Any) -> Iterator[ProviderEvent]: - yield from _iter_chat_response_events(api_response) - - @staticmethod - def _tool_calls_from_accum(tool_accum: dict[int, dict[str, str]]) -> Iterator[ToolCall]: - yield from _tool_calls_from_accum(tool_accum) - - def on_session_end(self) -> None: - logger.debug("Chat Completions API language model session state reset") diff --git a/runtime/src/speech_to_speech/LLM/compaction_prompt.py b/runtime/src/speech_to_speech/LLM/compaction_prompt.py deleted file mode 100644 index c1fc041fb3d2917d12c91c4a143762ab1fb4d735..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/LLM/compaction_prompt.py +++ /dev/null @@ -1,181 +0,0 @@ -"""Prompt template and factory for the conversation compaction (history summarization) function. - -Compaction reduces an unbounded conversation history to a tight user/assistant -summary pair, letting the pipeline continue indefinitely without running out of -context window. The factory :func:`build_compactor` returns a :data:`CompactFn` -compatible with :meth:`~speech_to_speech.LLM.chat.Chat.trim_if_needed`. - -:data:`CompactGenerateFn` is the backend-agnostic generation interface: -``(system_prompt: str, user_prompt: str) -> response_text: str``. -Each handler wraps its own model into this interface and passes it to -:func:`build_compactor`. -""" - -from __future__ import annotations - -import json -import logging -import re -from collections.abc import Callable -from typing import Any - -from speech_to_speech.LLM.chat import CompactFn, CompactionResult - -logger = logging.getLogger(__name__) - -# Callable[[system_prompt, user_prompt], response_text] -CompactGenerateFn = Callable[[str, str], str] - -# --------------------------------------------------------------------------- -# Prompts -# --------------------------------------------------------------------------- - -COMPACTION_SYSTEM_PROMPT = """\ -You are a conversation memory compressor for a real-time voice AI assistant. - -Your task: read a multi-turn conversation and produce a dense summary so the -assistant can continue naturally, as if it remembers everything that was said. - -Output a single JSON object with exactly two string fields: - "user_summary" — 1–5 sentences capturing what the user has been asking - about, any preferences or constraints they have stated, - and where the conversation stands from their perspective. - "assistant_summary" — 1–5 sentences capturing what the assistant has - explained, decided, or done (including tool calls and - their results), plus any open questions or commitments. - -Rules: -- Be information-dense: preserve names, numbers, file paths, error messages, and - other specifics that would be needed to continue the conversation correctly. -- Omit small-talk and filler that carries no forward context. -- Write in third person, past tense - (e.g. "The user asked about…", "The assistant provided…"). -- Emit only the JSON object — no markdown, no code fences, no extra keys.\ -""" - -COMPACTION_USER_TEMPLATE = """\ -Summarize the following conversation. Return only the JSON object. - ---- CONVERSATION START --- -{conversation} ---- CONVERSATION END ---\ -""" - -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - - -def _render_transcript(snapshot: list[Any]) -> str: - """Render a ResponseInputParam snapshot as a readable plain-text transcript.""" - lines: list[str] = [] - for item in snapshot: - if not isinstance(item, dict): - continue - item_type = item.get("type", "message") - role: str = item.get("role", "") - - if role == "system": - continue - - if item_type == "function_call": - name = item.get("name", "") - args = item.get("arguments", "") - lines.append(f"[Tool call: {name}({args})]") - continue - - if item_type == "function_call_output": - out = item.get("output", "") - lines.append(f"[Tool result: {out}]") - continue - - # Regular user / assistant message - content = item.get("content", "") - if isinstance(content, list): - text = " ".join( - c.get("text", "") - for c in content - if isinstance(c, dict) and c.get("type") in ("input_text", "output_text") - ).strip() - elif isinstance(content, str): - text = content.strip() - else: - continue - - if text: - label = role.capitalize() if role else "Unknown" - lines.append(f"{label}: {text}") - - return "\n\n".join(lines) - - -_JSON_BLOCK_RE = re.compile(r"```(?:json)?\s*(\{.*?\})\s*```", re.DOTALL) - - -def _extract_json(text: str) -> dict[str, Any]: - """Extract the first JSON object from *text*, stripping markdown code fences.""" - try: - return json.loads(text) - except json.JSONDecodeError: - pass - - m = _JSON_BLOCK_RE.search(text) - if m: - return json.loads(m.group(1)) - - start = text.find("{") - end = text.rfind("}") - if start != -1 and end > start: - return json.loads(text[start : end + 1]) - - raise ValueError(f"No JSON object found in compaction response: {text!r}") - - -# --------------------------------------------------------------------------- -# Factory -# --------------------------------------------------------------------------- - - -def build_compactor(generate_fn: CompactGenerateFn) -> CompactFn: - """Return a :data:`CompactFn` that summarizes history using *generate_fn*. - - *generate_fn* is the only model-specific dependency: it receives - ``(system_prompt, user_prompt)`` and returns the model's text response. - Both :class:`~speech_to_speech.LLM.openai_api_language_model.OpenApiModelHandler` - and :class:`~speech_to_speech.LLM.language_model.BaseLanguageModelHandler` - subclasses expose a ``_build_compaction_generate_fn()`` method that wraps - their respective backend into this interface. - - The returned callable is safe to call from a background thread. - - Args: - generate_fn: Backend-agnostic text generation callable. - - Returns: - A callable ``(snapshot: ResponseInputParam) -> CompactionResult``. - """ - - def compact(snapshot: list[Any]) -> CompactionResult: - transcript = _render_transcript(snapshot) - if not transcript.strip(): - logger.warning("Compaction called with an empty transcript; returning empty summaries") - return CompactionResult(user_summary="", assistant_summary="") - - user_content = COMPACTION_USER_TEMPLATE.format(conversation=transcript) - raw_text = generate_fn(COMPACTION_SYSTEM_PROMPT, user_content) - - data = _extract_json(raw_text) - user_summary = str(data.get("user_summary", "")).strip() - assistant_summary = str(data.get("assistant_summary", "")).strip() - - if not user_summary or not assistant_summary: - raise ValueError(f"Compaction response missing required fields. Got: {data!r}") - - logger.debug( - "Compaction complete. user=%d chars assistant=%d chars", - len(user_summary), - len(assistant_summary), - ) - return CompactionResult(user_summary=user_summary, assistant_summary=assistant_summary) - - return compact diff --git a/runtime/src/speech_to_speech/LLM/language_model.py b/runtime/src/speech_to_speech/LLM/language_model.py deleted file mode 100644 index 4f414cd9ee8565c7c19a1505de145e0c632611b9..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/LLM/language_model.py +++ /dev/null @@ -1,1011 +0,0 @@ -from __future__ import annotations - -import logging -from abc import ABC, abstractmethod -from collections.abc import Iterable, Iterator, Sized -from queue import Empty -from threading import Lock, Thread -from typing import Any, Literal, Optional, Protocol, runtime_checkable - -import torch -from nltk import sent_tokenize -from openai.types.realtime.realtime_conversation_item_function_call import ( - RealtimeConversationItemFunctionCall, -) -from openai.types.realtime.realtime_response_create_params import RealtimeResponseCreateParams -from openai.types.responses import ResponseFunctionToolCall -from pydantic import BaseModel, ConfigDict, Field -from transformers import ( - AutoModelForCausalLM, - AutoModelForImageTextToText, - AutoProcessor, - AutoTokenizer, - Pipeline, - PreTrainedModel, - PreTrainedTokenizerBase, - StoppingCriteria, - StoppingCriteriaList, - TextIteratorStreamer, - pipeline, -) - -from speech_to_speech.api.openai_realtime.runtime_config import RuntimeConfig -from speech_to_speech.baseHandler import BaseHandler -from speech_to_speech.LLM.chat import ( - Chat, - ChatItemError, - build_active_chat, - make_assistant_message, - make_system_message, - make_user_message, -) -from speech_to_speech.LLM.compaction_prompt import CompactGenerateFn, build_compactor -from speech_to_speech.LLM.text_prompt import build_text_system_prompt -from speech_to_speech.LLM.tool_call.function_call import extract_function_calls_from_text -from speech_to_speech.LLM.tool_call.function_tool import FunctionTool -from speech_to_speech.LLM.tool_call.tool_prompt import END_CODE, ENTER_CODE, build_block_regex, build_tool_system_prompt -from speech_to_speech.LLM.utils import image_url_to_pil, remove_unspeechable, resolve_auto_language -from speech_to_speech.LLM.voice_prompt import build_voice_system_prompt -from speech_to_speech.pipeline.cancel_scope import CancelScope -from speech_to_speech.pipeline.handler_types import LLMIn, LLMOut -from speech_to_speech.pipeline.messages import ( - EndOfResponse, - GenerateResponseRequest, - LLMResponseChunk, - TokenUsage, -) -from speech_to_speech.pipeline.speculative_turns import SpeculativeTurnTracker -from speech_to_speech.utils.utils import is_out_of_band, response_wants_audio - -try: - import mlx.core as mx - from mlx_lm import generate as mlx_generate - from mlx_lm import load as mlx_load - from mlx_lm import stream_generate as mlx_stream_generate - - from speech_to_speech.utils.mlx_lock import MLXLockContext - - HAS_MLX = True -except ImportError: - HAS_MLX = False - -try: - from mlx_vlm import load as mlx_vlm_load - from mlx_vlm import stream_generate as mlx_vlm_stream_generate - - HAS_MLX_VLM = True -except ImportError: - HAS_MLX_VLM = False - -logger = logging.getLogger(__name__) - - -@runtime_checkable -class _Tokenizer(Protocol): - """Minimal interface the base class needs from any tokenizer.""" - - def encode(self, text: str, /, **kwargs: Any) -> list[int] | Sized: ... - - -@runtime_checkable -class _Processor(Protocol): - """Minimal interface the VLM handler needs from any processor.""" - - tokenizer: _Tokenizer - - def apply_chat_template(self, conversation: Any, /, **kwargs: Any) -> Any: ... - def __call__(self, **kwargs: Any) -> Any: ... - - -class _CancelCriteria(StoppingCriteria): - """Stopping criteria that can be signalled externally to abort generation.""" - - def __init__(self) -> None: - self._cancelled = False - - def cancel(self) -> None: - self._cancelled = True - - def reset(self) -> None: - self._cancelled = False - - def __call__(self, input_ids: Any, scores: Any, **kwargs: Any) -> bool: - return self._cancelled - - -class StreamContext(BaseModel): - """Mutable accumulator passed through ``_stream_tokens`` so the caller - can read back generation state after the iterator is exhausted.""" - - model_config = ConfigDict(arbitrary_types_allowed=True) - - cancelled: bool = False - stopped: bool = False - raw_generated_text: str = "" - generated_text: str = "" - printable_text: str = "" - tools: list[ResponseFunctionToolCall] = Field(default_factory=list) - function_tools: list[FunctionTool] = Field(default_factory=list) - block_regex: Optional[str] = None - enter_code: Optional[str] = None - end_code: Optional[str] = None - input_tokens: int = 0 - sentence_batch: list[str] = Field(default_factory=list) - turn_id: str | None = None - turn_revision: int | None = None - speech_stopped_at_s: float | None = None - cancel_generation: int | None = None - - @property - def interrupted(self) -> bool: - """True when generation ended early for any reason.""" - return self.cancelled or self.stopped - - -class BaseLanguageModelHandler(BaseHandler[LLMIn, LLMOut], ABC): - """Abstract base for text-only and vision language model handlers. - - Holds shared pipeline logic (streaming, tool extraction, chat management) - and delegates model-specific behaviour to two abstract hooks: - ``_load_model`` and ``_generate``. - """ - - _cancel_criteria: _CancelCriteria - streamer: Iterable[str] - tokenizer: _Tokenizer - - def setup( - self, - model_name: str = "Qwen/Qwen3-4B-Instruct-2507", - device: str = "cuda", - torch_dtype: str = "float16", - gen_kwargs: dict[str, Any] = {}, - user_role: str = "user", - cancel_scope: CancelScope | None = None, - speculative_turns: SpeculativeTurnTracker | None = None, - backend: Literal["transformers", "mlx"] = "transformers", - enable_thinking: bool = False, - stream_batch_sentences: int = 3, - enable_lang_prompt: bool = False, - compact_history: bool = False, - **_kwargs: Any, - ) -> None: - self.backend = backend - self.cancel_scope = cancel_scope - self.speculative_turns = speculative_turns - self.device = device - self.model_name = model_name - self.enable_thinking = enable_thinking - self.stream_batch_sentences = max(1, stream_batch_sentences) - self.enable_lang_prompt = enable_lang_prompt - - self._load_model(model_name, device, torch_dtype, gen_kwargs) - - self.user_role = user_role - # Serializes transformers pipe/model.generate calls between the speech - # response path and the background compaction worker. MLX paths use - # MLXLockContext instead. - self._transformers_lock = Lock() - self.compactor = build_compactor(self._build_compaction_generate_fn()) if compact_history else None - - def _turn_is_latest(self, turn_id: str | None, turn_revision: int | None) -> bool: - return self.speculative_turns is None or self.speculative_turns.is_latest(turn_id, turn_revision) - - def _turn_output_allowed(self, turn_id: str | None, turn_revision: int | None) -> bool: - if self.speculative_turns is None: - return True - return self.speculative_turns.is_latest_after_reopen_grace(turn_id, turn_revision) - - @abstractmethod - def _load_model( - self, - model_name: str, - device: str, - torch_dtype: str, - gen_kwargs: dict[str, Any], - ) -> None: - """Load the model, tokenizer, and any backend-specific resources.""" - - def _build_compaction_generate_fn(self) -> CompactGenerateFn: - """Return a ``(system, user) -> text`` callable for compaction. - - Subclasses must override this when ``compact_history=True`` is used. - """ - raise NotImplementedError( - f"{type(self).__name__} does not implement _build_compaction_generate_fn. " - "Override it or pass compact_history=False." - ) - - @abstractmethod - def _generate( - self, - chat: Chat, - language_code: Optional[str], - gen: int | None, - ctx: StreamContext, - runtime_config: RuntimeConfig | None = None, - response: RealtimeResponseCreateParams | None = None, - ) -> Iterator[LLMResponseChunk]: - """Run model generation, stream sentence chunks, set ``ctx.input_tokens``.""" - - # ------------------------------------------------------------------ - # Generation lifecycle helpers - # ------------------------------------------------------------------ - - def _finish_mlx_generation(self, token_iter: Any) -> None: - """Close the MLX generator to release resources immediately.""" - token_iter.close() - - def _finish_transformers_generation(self, thread: Thread) -> None: - """Signal the stopping criteria, drain the streamer, and join the thread.""" - self._cancel_criteria.cancel() - try: - for _ in self.streamer: - pass - except Empty: - pass - thread.join(timeout=5.0) - - # ------------------------------------------------------------------ - # Shared helpers - # ------------------------------------------------------------------ - - def _apply_instructions( - self, - chat: Chat, - instructions: Optional[str], - raw_tools: list[Any] | None, - tool_choice: Optional[str], - ctx: StreamContext | None = None, - wants_audio: bool = True, - ) -> None: - if not instructions: - return - - raw_tools = raw_tools or [] - - function_tools = [FunctionTool(**t.model_dump()) for t in raw_tools if t.type == "function"] - - build_system_prompt = build_voice_system_prompt if wants_audio else build_text_system_prompt - - if function_tools and tool_choice != "none": - tool_section = build_tool_system_prompt(function_tools, text_only=not wants_audio) - full_instructions = build_system_prompt(instructions, tool_section=tool_section) - block_regex = build_block_regex() - enter_code = ENTER_CODE - end_code = END_CODE - else: - full_instructions = build_system_prompt(instructions) - block_regex = None - enter_code = None - end_code = None - - chat.add_item(make_system_message(full_instructions)) - - if ctx is not None: - ctx.function_tools = function_tools - ctx.block_regex = block_regex - ctx.enter_code = enter_code - ctx.end_code = end_code - - def _process_printable_text( - self, - printable_text: str, - language_code: Optional[str], - tools: list[ResponseFunctionToolCall], - ctx: StreamContext, - runtime_config: RuntimeConfig | None = None, - response: RealtimeResponseCreateParams | None = None, - ) -> tuple[list[LLMResponseChunk], list[ResponseFunctionToolCall], str]: - """Extract complete code blocks and return complete sentences to yield. - - Returns ``(chunks_to_yield, updated_tools, remaining_printable_text)``. - Each element in *chunks_to_yield* is an :class:`LLMResponseChunk` - ready to be yielded to the downstream pipeline. - - Sentences are accumulated in ``ctx.sentence_batch`` and only yielded - when the batch reaches ``self.stream_batch_sentences``. - """ - chunks: list[LLMResponseChunk] = [] - - if ctx.enter_code and ctx.enter_code in printable_text: - idx = printable_text.index(ctx.enter_code) - before = printable_text[:idx] - code_and_after = printable_text[idx:] - if before.strip(): - for s in sent_tokenize(before): - ctx.sentence_batch.append(s) - if ctx.sentence_batch: - chunks.append( - LLMResponseChunk( - text=" ".join(ctx.sentence_batch), - language_code=language_code, - runtime_config=runtime_config, - response=response, - turn_id=ctx.turn_id, - turn_revision=ctx.turn_revision, - speech_stopped_at_s=ctx.speech_stopped_at_s, - cancel_generation=ctx.cancel_generation, - ) - ) - ctx.sentence_batch = [] - if ctx.block_regex and ctx.end_code and ctx.end_code in code_and_after: - stripped, func_calls = extract_function_calls_from_text( - code_and_after, - ctx.block_regex, - ) - parsed_tools: list[ResponseFunctionToolCall] = [] - for fc in func_calls: - if tools: - logger.warning( - "Skipping extra tool call '%s'; only one tool call is allowed per response", - fc.function_name, - ) - continue - try: - tool_call = fc.to_realtime_function_tool_call(ctx.function_tools) - except ValueError as e: - logger.warning("Skipping invalid tool call: %s", e) - continue - tools.append(tool_call) - parsed_tools.append(tool_call) - if parsed_tools: - chunks.append( - LLMResponseChunk( - text="", - language_code=language_code, - tools=parsed_tools, - runtime_config=runtime_config, - response=response, - turn_id=ctx.turn_id, - turn_revision=ctx.turn_revision, - speech_stopped_at_s=ctx.speech_stopped_at_s, - cancel_generation=ctx.cancel_generation, - ) - ) - printable_text = stripped - else: - printable_text = code_and_after - return chunks, tools, printable_text - - if printable_text and not response_wants_audio(response) and ctx.enter_code is None: - # Text-only with no tool block: stream the raw text immediately. No TTS - # means no need to split into sentences, and sent_tokenize would collapse - # newlines / markdown. Streaming (vs buffering to the end) keeps the - # response interruptible by a new speech turn. - chunks.append( - LLMResponseChunk( - text=printable_text, - language_code=language_code, - runtime_config=runtime_config, - response=response, - turn_id=ctx.turn_id, - turn_revision=ctx.turn_revision, - speech_stopped_at_s=ctx.speech_stopped_at_s, - cancel_generation=ctx.cancel_generation, - ) - ) - return chunks, tools, "" - - if printable_text: - sentences = sent_tokenize(printable_text) - if len(sentences) > 1: - for s in sentences[:-1]: - ctx.sentence_batch.append(s) - if len(ctx.sentence_batch) >= self.stream_batch_sentences: - chunks.append( - LLMResponseChunk( - text=" ".join(ctx.sentence_batch), - language_code=language_code, - runtime_config=runtime_config, - response=response, - turn_id=ctx.turn_id, - turn_revision=ctx.turn_revision, - speech_stopped_at_s=ctx.speech_stopped_at_s, - cancel_generation=ctx.cancel_generation, - ) - ) - ctx.sentence_batch = [] - printable_text = sentences[-1] - - return chunks, tools, printable_text - - def _check_stop(self, gen: int | None, ctx: StreamContext) -> bool: - """Check whether generation should be aborted and mark the reason on *ctx*.""" - if gen is not None and self.cancel_scope is not None and self.cancel_scope.is_stale(gen): - ctx.cancelled = True - logger.info("LLM generation cancelled (interruption)") - return True - if not self._turn_is_latest(ctx.turn_id, ctx.turn_revision): - ctx.cancelled = True - logger.info("LLM generation cancelled (stale speculative turn)") - return True - if self.stop_event.is_set(): - ctx.stopped = True - logger.info("LLM generation stopped (shutdown)") - return True - return False - - def _stream_tokens( - self, - token_iter: Iterator[Any], - gen: int | None, - language_code: Optional[str], - ctx: StreamContext, - runtime_config: RuntimeConfig | None = None, - response: RealtimeResponseCreateParams | None = None, - ) -> Iterator[LLMResponseChunk]: - """Consume *token_iter*, accumulate text in *ctx*, yield sentence chunks.""" - # Text-only output keeps every character; only audio strips TTS-unfriendly - # symbols via remove_unspeechable. - wants_audio = response_wants_audio(response) - while True: - try: - token = next(token_iter) - except StopIteration: - break - except Empty: - if self._check_stop(gen, ctx): - break - continue - - if self._check_stop(gen, ctx): - break - - raw_text: str = token.text if hasattr(token, "text") else token - ctx.raw_generated_text += raw_text - clean = raw_text if not wants_audio else remove_unspeechable(raw_text) - ctx.generated_text += clean - ctx.printable_text += clean - chunks, ctx.tools, ctx.printable_text = self._process_printable_text( - ctx.printable_text, - language_code, - ctx.tools, - ctx, - runtime_config, - response, - ) - if chunks and not self._turn_output_allowed(ctx.turn_id, ctx.turn_revision): - ctx.cancelled = True - logger.info("LLM generation cancelled (stale speculative turn)") - break - yield from chunks - - if ctx.sentence_batch and not ctx.interrupted: - if ctx.printable_text.strip(): - ctx.sentence_batch.append(ctx.printable_text.strip()) - ctx.printable_text = "" - if not self._turn_output_allowed(ctx.turn_id, ctx.turn_revision): - ctx.cancelled = True - logger.info("LLM generation cancelled (stale speculative turn)") - return - yield LLMResponseChunk( - text=" ".join(ctx.sentence_batch), - language_code=language_code, - runtime_config=runtime_config, - response=response, - turn_id=ctx.turn_id, - turn_revision=ctx.turn_revision, - speech_stopped_at_s=ctx.speech_stopped_at_s, - cancel_generation=ctx.cancel_generation, - ) - ctx.sentence_batch = [] - - # ------------------------------------------------------------------ - # Main pipeline entry point - # ------------------------------------------------------------------ - - def process(self, request: LLMIn) -> Iterator[LLMOut]: - ctx = StreamContext() - - if not isinstance(request, GenerateResponseRequest): - raise TypeError(f"Unexpected request type: {type(request)}") - - ctx.turn_id = request.turn_id - ctx.turn_revision = request.turn_revision - ctx.speech_stopped_at_s = request.speech_stopped_at_s - if not self._turn_is_latest(ctx.turn_id, ctx.turn_revision): - logger.info("Skipping stale LLM request for turn=%s rev=%s", ctx.turn_id, ctx.turn_revision) - yield EndOfResponse(turn_id=ctx.turn_id, turn_revision=ctx.turn_revision) - return - - runtime_config = request.runtime_config - response = request.response - original_chat = runtime_config.chat - out_of_band = is_out_of_band(response) - if out_of_band: - try: - active_chat = build_active_chat(original_chat, response) - except ChatItemError as exc: - logger.info("Out-of-band response rejected: %s", exc) - yield EndOfResponse(turn_id=ctx.turn_id, turn_revision=ctx.turn_revision, error=str(exc)) - return - else: - active_chat = original_chat.copy() - language_code = request.language_code - instructions = ( - response.instructions if response and response.instructions else runtime_config.session.instructions - ) - tools = response.tools if response and response.tools else runtime_config.session.tools - tool_choice = response.tool_choice if response and response.tool_choice else runtime_config.session.tool_choice - self._apply_instructions( - active_chat, - instructions, - tools, - str(tool_choice) if tool_choice else None, - ctx, - response_wants_audio(response), - ) - language_code, lang_name = resolve_auto_language(language_code) - if lang_name and self.enable_lang_prompt: - active_chat.add_item(make_user_message(f"Please reply to my message in {lang_name}.")) - - gen = self.cancel_scope.generation if self.cancel_scope else None - ctx.cancel_generation = gen - # Images the model sees this turn; only these are stripped on write-back, - # so an image a fast client injects mid-generation for the next turn - # survives (it is not in this serialized snapshot). - consumed_image_ids = active_chat.image_message_ids() - - try: - yield from self._generate(active_chat, language_code, gen, ctx, runtime_config, response) - - if ctx.stopped: - return - - turn_output_allowed = not ctx.cancelled and self._turn_output_allowed(ctx.turn_id, ctx.turn_revision) - # Out-of-band responses still emit output, but never write back to the default - # conversation (their context was a throwaway chat). - commit_allowed = turn_output_allowed and not out_of_band - if commit_allowed: - original_chat.add_item(make_assistant_message(ctx.generated_text)) - if commit_allowed and ctx.tools: - for t in ctx.tools: - original_chat.add_item( - RealtimeConversationItemFunctionCall( - type="function_call", - id=t.id, - call_id=t.call_id, - name=t.name, - arguments=t.arguments, - status=t.status, - ) - ) - if commit_allowed: - original_chat.strip_images(consumed_image_ids) - original_chat.trim_if_needed(self.compactor) - logger.debug("Clean text: %s", ctx.generated_text) - logger.info(f"Tools: {ctx.tools}") - - if turn_output_allowed and ctx.printable_text.strip(): - yield LLMResponseChunk( - text=ctx.printable_text.strip(), - language_code=language_code, - runtime_config=runtime_config, - response=response, - turn_id=ctx.turn_id, - turn_revision=ctx.turn_revision, - speech_stopped_at_s=ctx.speech_stopped_at_s, - cancel_generation=ctx.cancel_generation, - ) - - output_tokens = len(self.tokenizer.encode(ctx.raw_generated_text)) if ctx.raw_generated_text else 0 - if turn_output_allowed and (ctx.input_tokens or output_tokens): - yield TokenUsage( - input_tokens=ctx.input_tokens, - output_tokens=output_tokens, - turn_id=ctx.turn_id, - turn_revision=ctx.turn_revision, - ) - except Exception as exc: - # Any generation failure must still terminate the response. Without this - # the exception would escape process() and no EndOfResponse would be - # emitted, leaving st.in_response stuck and locking every later response. - logger.exception("LLM generation failed; ending the current response") - yield EndOfResponse( - turn_id=ctx.turn_id, - turn_revision=ctx.turn_revision, - cancel_generation=ctx.cancel_generation, - error=f"Language model generation failed: {exc}", - ) - return - yield EndOfResponse( - turn_id=ctx.turn_id, - turn_revision=ctx.turn_revision, - cancel_generation=ctx.cancel_generation, - ) - - def on_session_end(self) -> None: - logger.debug("Language model session state reset") - - -# ====================================================================== -# Text-only LLM -# ====================================================================== - - -class LanguageModelHandler(BaseLanguageModelHandler): - """Text-only language model handler (transformers or MLX backend).""" - - model: PreTrainedModel - tokenizer: PreTrainedTokenizerBase - gen_kwargs: dict - torch_dtype: torch.dtype - pipe: Pipeline - streamer: TextIteratorStreamer - - def setup(self, **kwargs: Any) -> None: # type: ignore[override] - super().setup(**kwargs) - logger.info(f"LLM Backend: {self.backend}") - self.warmup() - - def _load_model(self, model_name: str, device: str, torch_dtype: str, gen_kwargs: dict[str, Any]) -> None: - - logger.info("LLM Language Model Handler setup") - - if self.backend == "mlx": - if not HAS_MLX: - raise ImportError("mlx_lm is required for the 'mlx' backend. Install with: pip install mlx-lm") - self.model, self.tokenizer = mlx_load(model_name) # type: ignore[assignment, misc] - self.gen_kwargs = gen_kwargs - else: - self.torch_dtype = getattr(torch, torch_dtype) - self.tokenizer = AutoTokenizer.from_pretrained(model_name) # type: ignore[assignment] - self.model = AutoModelForCausalLM.from_pretrained( - model_name, torch_dtype=torch_dtype, trust_remote_code=True - ).to(device) # type: ignore[arg-type] - self.pipe = pipeline( # type: ignore[call-overload] - "text-generation", model=self.model, tokenizer=self.tokenizer, device=device - ) - self.streamer = TextIteratorStreamer( - self.tokenizer, # type: ignore[arg-type] - skip_prompt=True, - skip_special_tokens=True, - timeout=1.0, - ) - self._cancel_criteria = _CancelCriteria() - self.gen_kwargs = { - "streamer": self.streamer, - "return_full_text": False, - "stopping_criteria": StoppingCriteriaList([self._cancel_criteria]), - **gen_kwargs, - } - - def _generate( - self, - chat: Chat, - language_code: Optional[str], - gen: int | None, - ctx: StreamContext, - runtime_config: RuntimeConfig | None = None, - response: RealtimeResponseCreateParams | None = None, - ) -> Iterator[LLMResponseChunk]: - chat_messages = chat.to_transformers_chat() - chat_input = self.tokenizer.apply_chat_template(chat_messages, tokenize=True) - ctx.input_tokens += len(chat_input if isinstance(chat_input, list) else chat_input["input_ids"]) # type: ignore[index] - logger.debug("Prompt token count: %d", ctx.input_tokens) - - chat_prompt = self.tokenizer.apply_chat_template( - chat_messages, tokenize=False, add_generation_prompt=True, enable_thinking=False - ) - - if self.backend == "mlx": - with MLXLockContext(handler_name="MLX-LLM", timeout=10.0): - token_iter = mlx_stream_generate( - self.model, # type: ignore[arg-type] - self.tokenizer, # type: ignore[arg-type] - chat_prompt, # type: ignore[arg-type] - max_tokens=self.gen_kwargs["max_new_tokens"], - ) - yield from self._stream_tokens(token_iter, gen, language_code, ctx, runtime_config, response) - self._finish_mlx_generation(token_iter) - try: - mx.clear_cache() - except Exception: - pass - torch.mps.empty_cache() - else: - self._cancel_criteria.reset() - lock = self._transformers_lock - - def _locked_pipe() -> None: - with lock: - self.pipe(chat_prompt, **self.gen_kwargs) - - thread = Thread(target=_locked_pipe) - thread.start() - yield from self._stream_tokens(self.streamer, gen, language_code, ctx, runtime_config, response) - self._finish_transformers_generation(thread) - if self.device == "mps": - torch.mps.empty_cache() - - def _build_compaction_generate_fn(self) -> CompactGenerateFn: - if self.backend == "mlx": - model = self.model - tokenizer = self.tokenizer - max_tokens = 1024 - - def generate_mlx(system: str, user: str) -> str: - messages = [{"role": "system", "content": system}, {"role": "user", "content": user}] - prompt = tokenizer.apply_chat_template( # type: ignore[union-attr] - messages, tokenize=False, add_generation_prompt=True - ) - with MLXLockContext(handler_name="MLX-compact", timeout=10.0): - return mlx_generate(model, tokenizer, prompt=prompt, max_tokens=max_tokens, verbose=False) # type: ignore[arg-type] - - return generate_mlx - else: - pipe = self.pipe - tokenizer = self.tokenizer - gen_kwargs = { - k: v - for k, v in self.gen_kwargs.items() - if k not in ("streamer", "stopping_criteria", "max_new_tokens", "return_full_text") - } - max_new_tokens = 1024 - lock = self._transformers_lock - - def generate_transformers(system: str, user: str) -> str: - messages = [{"role": "system", "content": system}, {"role": "user", "content": user}] - prompt = tokenizer.apply_chat_template( # type: ignore[union-attr] - messages, tokenize=False, add_generation_prompt=True, enable_thinking=False - ) - with lock: - result = pipe(prompt, return_full_text=False, max_new_tokens=max_new_tokens, **gen_kwargs) - return result[0]["generated_text"] # type: ignore[index] - - return generate_transformers - - def warmup(self) -> None: - logger.info(f"Warming up {self.__class__.__name__}") - - dummy_input_text = "Repeat the word 'home'." - dummy_chat = [{"role": self.user_role, "content": dummy_input_text}] - n_steps = 2 - - if self.backend == "mlx": - for _ in range(n_steps): - prompt = self.tokenizer.apply_chat_template(dummy_chat, tokenize=False) - mlx_generate( - self.model, # type: ignore[arg-type] - self.tokenizer, # type: ignore[arg-type] - prompt=prompt, # type: ignore[arg-type] - max_tokens=self.gen_kwargs["max_new_tokens"], - verbose=False, - ) - else: - warmup_gen_kwargs = { - "min_new_tokens": self.gen_kwargs["min_new_tokens"], - "max_new_tokens": self.gen_kwargs["max_new_tokens"], - **self.gen_kwargs, - } - - if self.device == "cuda": - start_event = torch.cuda.Event(enable_timing=True) - end_event = torch.cuda.Event(enable_timing=True) - torch.cuda.synchronize() - start_event.record() - - for _ in range(n_steps): - thread = Thread(target=self.pipe, args=(dummy_chat,), kwargs=warmup_gen_kwargs) - thread.start() - for _ in self.streamer: - pass - - if self.device == "cuda": - end_event.record() - torch.cuda.synchronize() - logger.info( - f"{self.__class__.__name__}: warmed up! time: {start_event.elapsed_time(end_event) * 1e-3:.3f} s" - ) - - -# ====================================================================== -# Vision Language Model (VLM) -# ====================================================================== - - -class VisionLanguageModelHandler(BaseLanguageModelHandler): - """Vision language model handler (transformers or MLX-VLM backend).""" - - model: PreTrainedModel - tokenizer: PreTrainedTokenizerBase - processor: _Processor - gen_kwargs: dict - torch_dtype: torch.dtype - streamer: TextIteratorStreamer - - def setup(self, **kwargs: Any) -> None: # type: ignore[override] - super().setup(**kwargs) - logger.info(f"VLM Backend: {self.backend}") - - def _load_model(self, model_name: str, device: str, torch_dtype: str, gen_kwargs: dict[str, Any]) -> None: - - logger.info("VLM Language Model Handler setup") - - if self.backend == "mlx": - if not HAS_MLX_VLM: - raise ImportError("mlx-vlm is required for MLX VLM models. Install with: pip install mlx-vlm") - self.model, self.processor = mlx_vlm_load(model_name) # type: ignore[assignment] - self.tokenizer = self.processor.tokenizer # type: ignore[assignment] - self.gen_kwargs = gen_kwargs - else: - self.torch_dtype = getattr(torch, torch_dtype) - self.processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True) # type: ignore[assignment] - self.tokenizer = self.processor.tokenizer # type: ignore[assignment] - self.model = AutoModelForImageTextToText.from_pretrained( - model_name, torch_dtype=self.torch_dtype, trust_remote_code=True - ).to(device) # type: ignore[arg-type] - self.streamer = TextIteratorStreamer( - self.tokenizer, # type: ignore[arg-type] - skip_prompt=True, - skip_special_tokens=True, - timeout=1.0, - ) - self._cancel_criteria = _CancelCriteria() - self.gen_kwargs = gen_kwargs - - def _generate( - self, - chat: Chat, - language_code: Optional[str], - gen: int | None, - ctx: StreamContext, - runtime_config: RuntimeConfig | None = None, - response: RealtimeResponseCreateParams | None = None, - ) -> Iterator[LLMResponseChunk]: - prepared = chat.to_transformers_chat() - if self.backend == "mlx": - images, formatted_prompt = self._prepare_mlx_vlm_inputs(prepared) - ctx.input_tokens += len(self.tokenizer.encode(formatted_prompt)) - logger.debug("MLX VLM prompt token count: %d", ctx.input_tokens) - - with MLXLockContext(handler_name="MLX-VLM", timeout=10.0): - token_iter = mlx_vlm_stream_generate( # type: ignore[arg-type] - self.model, - self.processor, - formatted_prompt, - images or None, - max_tokens=self.gen_kwargs.get("max_new_tokens", 1024), - enable_thinking=self.enable_thinking, - ) - yield from self._stream_tokens(token_iter, gen, language_code, ctx, runtime_config, response) - self._finish_mlx_generation(token_iter) - try: - mx.clear_cache() - except Exception: - pass - torch.mps.empty_cache() - else: - inputs, input_tokens = self._prepare_vlm_inputs(prepared) - ctx.input_tokens += input_tokens - logger.debug("VLM prompt token count: %d", ctx.input_tokens) - - self._cancel_criteria.reset() - generate_kwargs = { - **inputs, - "max_new_tokens": self.gen_kwargs.get("max_new_tokens", 1024), - "streamer": self.streamer, - "stopping_criteria": StoppingCriteriaList([self._cancel_criteria]), - } - lock = self._transformers_lock - - def _locked_generate() -> None: - with lock: - self.model.generate(**generate_kwargs) # type: ignore[union-attr,operator] - - thread = Thread(target=_locked_generate) - thread.start() - yield from self._stream_tokens(self.streamer, gen, language_code, ctx, runtime_config, response) - self._finish_transformers_generation(thread) - if self.device == "mps": - torch.mps.empty_cache() - - def _build_compaction_generate_fn(self) -> CompactGenerateFn: - if self.backend == "mlx": - model = self.model - processor = self.processor - tokenizer = self.tokenizer - max_tokens = 1024 - - def generate_mlx(system: str, user: str) -> str: - messages = [{"role": "system", "content": system}, {"role": "user", "content": user}] - formatted_prompt = processor.apply_chat_template( # type: ignore[union-attr] - messages, tokenize=False, add_generation_prompt=True - ) - with MLXLockContext(handler_name="MLX-VLM-compact", timeout=10.0): - token_iter = mlx_vlm_stream_generate( # type: ignore[arg-type] - model, processor, formatted_prompt, None, max_tokens=max_tokens - ) - return "".join(t.text if hasattr(t, "text") else str(t) for t in token_iter) - - return generate_mlx - else: - model = self.model - processor = self.processor - tokenizer = self.tokenizer - device = self.device - max_new_tokens = 1024 - lock = self._transformers_lock - - def generate_transformers(system: str, user: str) -> str: - messages = [{"role": "system", "content": system}, {"role": "user", "content": user}] - prompt = processor.apply_chat_template( # type: ignore[union-attr] - messages, tokenize=False, add_generation_prompt=True - ) - inputs = processor(text=[prompt], return_tensors="pt", padding=True).to(device) # type: ignore[operator] - input_len = inputs["input_ids"].shape[1] - with lock: - output_ids = model.generate(**inputs, max_new_tokens=max_new_tokens) # type: ignore[union-attr,operator] - return str(tokenizer.decode(output_ids[0][input_len:], skip_special_tokens=True)) # type: ignore[union-attr] - - return generate_transformers - - def _prepare_vlm_inputs(self, chat_messages: list[dict[str, Any]]) -> tuple[Any, int]: - """Build processor inputs for transformers VLM generation. - - Converts Realtime-style content (``input_text``/``input_image``) to the - transformers chat-template format (``text``/``image``) and decodes image - URLs to PIL images. - """ - converted_messages = [] - images = [] - for msg in chat_messages: - content = msg.get("content") - if isinstance(content, list): - parts = [] - for p in content: - if p.get("type") == "input_text": - parts.append({"type": "text", "text": p["text"]}) - elif p.get("type") == "input_image": - pil_img = image_url_to_pil(p["image_url"]) - images.append(pil_img) - parts.append({"type": "image", "image": pil_img}) - converted_messages.append({"role": msg["role"], "content": parts}) - else: - converted_messages.append(msg) - - text_prompt = self.processor.apply_chat_template( - converted_messages, - tokenize=False, - add_generation_prompt=True, - ) - inputs = self.processor( - text=[text_prompt], - images=images if images else None, - return_tensors="pt", - padding=True, - ).to(self.device) - return inputs, len(inputs["input_ids"][0]) - - def _prepare_mlx_vlm_inputs(self, chat_messages: list[dict[str, Any]]) -> tuple[list[Any], str]: - """Build formatted prompt and extract PIL images for MLX VLM generation. - - Converts Realtime-style content (``input_text``/``input_image``) to the - HF chat-template format and uses the processor's own ``apply_chat_template`` - for full multi-turn support. - - Returns ``(images, formatted_prompt)`` ready for ``mlx_vlm.stream_generate``. - """ - images = [] - converted_messages = [] - for msg in chat_messages: - content = msg.get("content") - if isinstance(content, list): - parts = [] - for p in content: - if p.get("type") == "input_text": - parts.append({"type": "text", "text": p["text"]}) - elif p.get("type") == "input_image": - images.append(image_url_to_pil(p["image_url"])) - parts.append({"type": "image"}) - converted_messages.append({"role": msg["role"], "content": parts}) - else: - converted_messages.append(msg) - - formatted_prompt = self.processor.apply_chat_template( - converted_messages, - tokenize=False, - add_generation_prompt=True, - ) - return images, formatted_prompt diff --git a/runtime/src/speech_to_speech/LLM/lm_output_processor.py b/runtime/src/speech_to_speech/LLM/lm_output_processor.py deleted file mode 100644 index fcfb6a7bb82087a94519c802f47f52aab1100e03..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/LLM/lm_output_processor.py +++ /dev/null @@ -1,148 +0,0 @@ -""" -LLM Output Processor - -Intercepts LLM output to: -1. Extract tool calls and send them via text_output_queue -2. Forward clean text to TTS pipeline -""" - -from __future__ import annotations - -import logging -from collections.abc import Iterator -from queue import Queue - -from speech_to_speech.baseHandler import BaseHandler -from speech_to_speech.pipeline.events import AssistantTextEvent, ResponseFailedEvent, TokenUsageEvent -from speech_to_speech.pipeline.handler_types import LLMOut, TTSIn -from speech_to_speech.pipeline.messages import EndOfResponse, LLMResponseChunk, TokenUsage, TTSInput -from speech_to_speech.pipeline.queue_types import TextEventItem -from speech_to_speech.pipeline.speculative_turns import SpeculativeTurnTracker -from speech_to_speech.utils.utils import response_wants_audio - -logger = logging.getLogger(__name__) - - -class LMOutputProcessor(BaseHandler[LLMOut, TTSIn]): - """ - Processes LLM output to extract tool calls and forward clean text to TTS. - - Input: :class:`LLMResponseChunk`, :class:`TokenUsage`, or :class:`EndOfResponse` from LLM - Output: :class:`TTSInput` or :class:`EndOfResponse` to TTS - Side effect: Sends :class:`AssistantTextEvent` / :class:`TokenUsageEvent` to text_output_queue - """ - - def setup( - self, - text_output_queue: Queue[TextEventItem] | None = None, - speculative_turns: SpeculativeTurnTracker | None = None, - ) -> None: - """ - Initialize the processor. - - Args: - text_output_queue: Queue to send text messages and tool calls - """ - self.text_output_queue = text_output_queue - self.speculative_turns = speculative_turns - - def _turn_output_allowed(self, turn_id: str | None, turn_revision: int | None) -> bool: - if self.speculative_turns is None: - return True - return self.speculative_turns.is_latest_after_reopen_grace(turn_id, turn_revision) - - def process(self, lm_output: LLMOut) -> Iterator[TTSIn]: - """ - Process LLM output: send text/tools to WebSocket, forward clean text to TTS. - - Yields: - :class:`TTSInput` or :class:`EndOfResponse` for TTS - """ - if isinstance(lm_output, TokenUsage): - if not self._turn_output_allowed( - lm_output.turn_id, - lm_output.turn_revision, - ): - logger.debug( - "Dropping stale token usage for turn=%s rev=%s", lm_output.turn_id, lm_output.turn_revision - ) - return - if self.text_output_queue is not None: - self.text_output_queue.put( - TokenUsageEvent( - input_tokens=lm_output.input_tokens or 0, - output_tokens=lm_output.output_tokens or 0, - turn_id=lm_output.turn_id, - turn_revision=lm_output.turn_revision, - ) - ) - return - - if isinstance(lm_output, EndOfResponse): - if not self._turn_output_allowed( - lm_output.turn_id, - lm_output.turn_revision, - ): - logger.debug( - "Dropping stale end-of-response for turn=%s rev=%s", - lm_output.turn_id, - lm_output.turn_revision, - ) - return - # A failed generation (e.g. invalid out-of-band input) closes the response as - # "failed" via the text side-channel, then falls through to emit the normal - # EndOfResponse so the audio path still re-enables listening / releases the slot. - if lm_output.error and self.text_output_queue is not None: - self.text_output_queue.put( - ResponseFailedEvent( - message=lm_output.error, - turn_id=lm_output.turn_id, - turn_revision=lm_output.turn_revision, - ) - ) - yield EndOfResponse( - turn_id=lm_output.turn_id, - turn_revision=lm_output.turn_revision, - cancel_generation=lm_output.cancel_generation, - ) - return - - if not isinstance(lm_output, LLMResponseChunk): - logger.warning("LMOutputProcessor received unexpected type: %s", type(lm_output)) - return - - if not self._turn_output_allowed( - lm_output.turn_id, - lm_output.turn_revision, - ): - logger.debug("Dropping stale LLM chunk for turn=%s rev=%s", lm_output.turn_id, lm_output.turn_revision) - return - - logger.debug(f"LM processor: text='{lm_output.text}', tools={lm_output.tools}") - - if self.text_output_queue is not None: - event = AssistantTextEvent( - text=lm_output.text, - turn_id=lm_output.turn_id, - turn_revision=lm_output.turn_revision, - cancel_generation=lm_output.cancel_generation, - ) - if lm_output.tools: - event.tools = lm_output.tools - logger.info(f"Sending to clients: text='{lm_output.text}', tools={[t.name for t in lm_output.tools]}") - else: - logger.debug(f"Sending to clients: text='{lm_output.text}' (no tools)") - self.text_output_queue.put(event) - - if lm_output.text and response_wants_audio(lm_output.response): - logger.debug(f"Forwarding to TTS: '{lm_output.text}'") - yield TTSInput( - text=lm_output.text, - language_code=lm_output.language_code, - runtime_config=lm_output.runtime_config, - response=lm_output.response, - turn_id=lm_output.turn_id, - turn_revision=lm_output.turn_revision, - speech_stopped_at_s=lm_output.speech_stopped_at_s, - cancel_generation=lm_output.cancel_generation, - ) diff --git a/runtime/src/speech_to_speech/LLM/responses_api_language_model.py b/runtime/src/speech_to_speech/LLM/responses_api_language_model.py deleted file mode 100644 index 53ee60b22768ab22e4383ce9127af321c4ae379b..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/LLM/responses_api_language_model.py +++ /dev/null @@ -1,201 +0,0 @@ -from __future__ import annotations - -import logging -import time -from collections.abc import Iterator -from typing import Any - -from openai import Stream -from openai.types.realtime.realtime_conversation_item_assistant_message import ( - Content as AssistantContent, -) -from openai.types.responses import ( - ResponseCompletedEvent, - ResponseFunctionToolCall, - ResponseOutputItemDoneEvent, - ResponseOutputMessage, - ResponseTextDeltaEvent, -) - -from speech_to_speech.LLM.base_openai_compatible_language_model import ( - WARMUP_MAX_RETRIES, - AssistantMessage, - BaseOpenAICompatibleHandler, - ProviderEvent, - TextDelta, - ToolCall, - Usage, -) -from speech_to_speech.LLM.chat import Chat -from speech_to_speech.LLM.chat_completions_language_model import ( - _build_chat_optional_kwargs, - _chat_messages, - _iter_chat_response_events, - _iter_chat_stream_events, - _request_chat_completions, -) -from speech_to_speech.LLM.compaction_prompt import CompactGenerateFn -from speech_to_speech.utils.utils import _generate_id - -logger = logging.getLogger(__name__) - - -class ResponsesApiModelHandler(BaseOpenAICompatibleHandler): - """LLM handler that talks to an OpenAI ``/v1/responses`` server.""" - - def warmup(self) -> None: - logger.info(f"Warming up {self.__class__.__name__}") - start = time.time() - self.client.with_options(max_retries=WARMUP_MAX_RETRIES).responses.create( - model=self.model_name, - input=[ - { - "type": "message", - "role": "system", - "content": [{"type": "input_text", "text": "You are a helpful assistant"}], - }, - {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "Hello"}]}, - ], - timeout=self.request_timeout, - ) - end = time.time() - logger.info(f"{self.__class__.__name__}: warmed up! time: {(end - start):.3f} s") - - def _build_compaction_generate_fn(self) -> CompactGenerateFn: - """Return a generate fn that calls the Responses API for compaction.""" - client = self.client - model_name = self.model_name - timeout = self.request_timeout - - def generate(system: str, user: str) -> str: - response = client.responses.create( - model=model_name, - input=[ - { - "type": "message", - "role": "system", - "content": [{"type": "input_text", "text": system}], - }, - { - "type": "message", - "role": "user", - "content": [{"type": "input_text", "text": user}], - }, - ], - timeout=timeout, - ) - return response.output_text - - return generate - - def _build_audio_optional_kwargs( - self, - response: Any, - req_tools: Any, - req_tool_choice: Any, - ) -> dict[str, Any]: - kwargs = _build_chat_optional_kwargs(req_tools, req_tool_choice) - max_tokens = getattr(response, "max_output_tokens", None) if response is not None else None - kwargs.setdefault("max_tokens", max_tokens or self.audio_max_tokens) - kwargs.setdefault("temperature", self.audio_temperature) - return kwargs - - def _serialize_audio(self, active_chat: Chat) -> list[dict[str, Any]]: - return _chat_messages(active_chat, audio_content_type=self.audio_content_type) - - def _request_audio( - self, - api_input: list[dict[str, Any]], - optional_kwargs: dict[str, Any], - runtime_config: Any, - ) -> Any: - extra_body = self._extra_body - session_reasoning_effort = getattr(runtime_config.session, "reasoning_effort", None) - if session_reasoning_effort: - extra_body = {"reasoning_effort": session_reasoning_effort} - return _request_chat_completions( - client=self.client, - model_name=self.model_name, - messages=api_input, - stream=self.stream, - extra_body=extra_body, - timeout=self.request_timeout, - optional_kwargs=optional_kwargs, - ) - - def _iter_audio_events(self, api_response: Any) -> Iterator[ProviderEvent]: - if self.stream: - yield from _iter_chat_stream_events(api_response) - else: - yield from _iter_chat_response_events(api_response) - - # ── base hooks ──────────────────────────────────────────────────────────-- - - def _serialize(self, active_chat: Chat) -> Any: - return active_chat.to_responses_api_chat() - - def _build_optional_kwargs(self, req_tools: Any, req_tool_choice: Any) -> dict[str, Any]: - optional_kwargs: dict[str, Any] = {} - if req_tools is not None: - optional_kwargs["tools"] = req_tools - if req_tool_choice is not None: - optional_kwargs["tool_choice"] = req_tool_choice - return optional_kwargs - - def _request(self, api_input: Any, optional_kwargs: dict[str, Any], runtime_config: Any) -> Any: - extra_body = self._extra_body - session_reasoning_effort = getattr(runtime_config.session, "reasoning_effort", None) - if session_reasoning_effort: - extra_body = {"reasoning_effort": session_reasoning_effort} - return self.client.responses.create( - model=self.model_name, - input=api_input, - stream=self.stream, - extra_body=extra_body, - timeout=self.request_timeout, - **optional_kwargs, - ) - - @staticmethod - def _assistant_content(content: Any) -> list[AssistantContent]: - return [ - AssistantContent(type="output_text", text=c.text if c.type == "output_text" else c.refusal) for c in content - ] - - def _iter_stream_events(self, api_response: Stream) -> Iterator[ProviderEvent]: - for raw_event in api_response: - if isinstance(raw_event, ResponseTextDeltaEvent): - yield TextDelta(text=raw_event.delta) - elif isinstance(raw_event, ResponseOutputItemDoneEvent): - item = raw_event.item - if isinstance(item, ResponseFunctionToolCall): - item.call_id = _generate_id("call") - item.id = _generate_id("fc") - yield ToolCall(item=item) - elif isinstance(item, ResponseOutputMessage): - yield AssistantMessage(content=self._assistant_content(item.content)) - elif isinstance(raw_event, ResponseCompletedEvent): - usage = getattr(raw_event.response, "usage", None) - if usage: - yield Usage(input_tokens=usage.input_tokens or 0, output_tokens=usage.output_tokens or 0) - - def _iter_response_events(self, api_response: Any) -> Iterator[ProviderEvent]: - usage = api_response.usage - if usage: - yield Usage(input_tokens=usage.input_tokens or 0, output_tokens=usage.output_tokens or 0) - for message in api_response.output: - if isinstance(message, ResponseFunctionToolCall): - message.call_id = _generate_id("call") - message.id = _generate_id("fc") - yield ToolCall(item=message) - elif isinstance(message, ResponseOutputMessage): - yield AssistantMessage(content=self._assistant_content(message.content)) - # Text-only keeps every character; the base applies remove_unspeechable - # for audio. Only output_text parts are spoken (refusals are stored). - raw = "".join(c.text for c in message.content if c.type == "output_text") - yield TextDelta(text=raw) - else: - logger.warning(f"Not supported message type: {message.type}") - - def on_session_end(self) -> None: - logger.debug("OpenAI API language model session state reset") diff --git a/runtime/src/speech_to_speech/LLM/text_prompt.py b/runtime/src/speech_to_speech/LLM/text_prompt.py deleted file mode 100644 index 3ded9b8687e197f6346035aa1c7f28307cea2dab..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/LLM/text_prompt.py +++ /dev/null @@ -1,44 +0,0 @@ -"""Text-channel system prompt: lead + session prompt + tail (strongest constraints last).""" - -TEXT_SYSTEM_PROMPT_LEAD = """\ -You are a helpful assistant in a text conversation. -""" - -TEXT_SYSTEM_PROMPT_TAIL = """\ -## Text Rules -- Write clearly and directly. Match length to the request: concise for simple questions, fuller when the task genuinely needs it. -- Use markdown when it helps (lists, code blocks, tables, emphasis); don't over-format simple answers. -- This is a written channel: no spoken-style filler and no action/emote text like *laughs*. -- Use tools when they help fulfill the request. No preamble sentence is required before a tool call. -- For slow or external tools, just call the tool and use the result; you don't need to announce it. -- If unsure whether a tool is needed, just answer directly. -""" - -# Skeleton for the assembled system message (placeholders filled in build_text_system_prompt). -_TEXT_SYSTEM_PROMPT_FULL = """\ -{lead} - -Session Prompt: -{session_prompt}{optional_tools} - -{tail} -""" - - -def build_text_system_prompt(session_prompt: str, *, tool_section: str = "") -> str: - """Context → session prompt → optional tool block → strongest text rules last.""" - tools = tool_section.strip() - optional_tools = f"\n\n{tools}" if tools else "" - return _TEXT_SYSTEM_PROMPT_FULL.format( - lead=TEXT_SYSTEM_PROMPT_LEAD.rstrip(), - session_prompt=session_prompt.strip(), - optional_tools=optional_tools, - tail=TEXT_SYSTEM_PROMPT_TAIL.rstrip(), - ) - - -# Full text instructions without a separate session block (legacy / rare direct use). -TEXT_SYSTEM_PROMPT = "{lead}\n\n{tail}".format( - lead=TEXT_SYSTEM_PROMPT_LEAD.rstrip(), - tail=TEXT_SYSTEM_PROMPT_TAIL.rstrip(), -) diff --git a/runtime/src/speech_to_speech/LLM/tool_call/function_call.py b/runtime/src/speech_to_speech/LLM/tool_call/function_call.py deleted file mode 100644 index 1871db749ecf8f1cfb4927f6116f33c7c2ad290d..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/LLM/tool_call/function_call.py +++ /dev/null @@ -1,321 +0,0 @@ -#!/usr/bin/env python3 -""" -Function parser for extracting function names, parameter names, and values from string function calls. - -Uses Python's ``tokenize`` and ``ast`` modules so that nested parentheses, -strings containing ')' characters, tuples, dicts, etc. are handled correctly. -""" - -import ast -import io -import json -import logging -import re -import tokenize -from collections import OrderedDict -from typing import Any, Dict, List, Tuple - -from openai.types.responses import ResponseFunctionToolCall -from pydantic import BaseModel - -from speech_to_speech.LLM.tool_call.function_tool import FunctionTool -from speech_to_speech.utils.utils import _generate_id - -logger = logging.getLogger(__name__) - -_POSITIONAL_RE = re.compile(r"^__arg_\d+__$") -_LENIENT_CALL_RE = re.compile( - r"\b[A-Za-z_]\w*(?:\.[A-Za-z_]\w*)*\s*" - r"\((?:[^()\"']+|\"(?:\\.|[^\"])*\"|'(?:\\.|[^'])*')*\)" -) - - -# ── AST / tokenize helpers ─────────────────────────────────────────── - - -def _split_top_level_calls(source: str) -> List[str]: - """Split *source* into individual ``name(...)`` expression strings. - - Uses the tokenizer to walk tokens and track parenthesis depth so that - nested parens, strings with ')' chars, etc. are handled correctly. - """ - tokens = list(tokenize.generate_tokens(io.StringIO(source).readline)) - calls: List[str] = [] - i = 0 - - while i < len(tokens): - tok = tokens[i] - if tok.type != tokenize.NAME: - i += 1 - continue - - start = i - j = i + 1 - - # Walk past dotted attribute access (e.g. ``mobile.click``) - while j + 1 < len(tokens) and tokens[j].string == "." and tokens[j + 1].type == tokenize.NAME: - j += 2 - - if j >= len(tokens) or tokens[j].string != "(": - i += 1 - continue - - # Track balanced parens - depth = 0 - end = None - k = j - while k < len(tokens): - t = tokens[k] - if t.type == tokenize.OP and t.string == "(": - depth += 1 - elif t.type == tokenize.OP and t.string == ")": - depth -= 1 - if depth == 0: - end = k - break - k += 1 - - if end is None: - i += 1 - continue - - calls.append(tokenize.untokenize(tokens[start : end + 1]).strip()) - i = end + 1 - - return calls - - -def _split_simple_calls_with_regex(source: str) -> List[str]: - """Extract complete simple ``name(args)`` spans from malformed model output. - - This fallback can recover well-formed siblings before a tokenizer error, - but not the incomplete call that caused the tokenizer error. - """ - return [match.group(0).strip() for match in _LENIENT_CALL_RE.finditer(source)] - - -def _parse_function_exprs( - expressions: List[str], - pattern_to_match: list[str], - *, - skip_invalid: bool = False, -) -> List["FunctionToolCall"]: - results: List[FunctionToolCall] = [] - for expr in expressions: - try: - call = _parse_call_expr(expr) - except Exception: - if skip_invalid: - continue - raise - if pattern_to_match and all(pattern not in call.function_name for pattern in pattern_to_match): - continue - results.append(call) - return results - - -def _extract_function_name(node: ast.expr) -> str: - """Return the dotted function name from a Call node's ``func`` attribute.""" - if isinstance(node, ast.Name): - return node.id - if isinstance(node, ast.Attribute): - base = _extract_function_name(node.value) - return f"{base}.{node.attr}" if base else node.attr - raise ValueError(f"Unsupported function target: {ast.dump(node)}") - - -def _literal_from_ast(node: ast.AST) -> Any: - """Convert an AST node to a Python literal value.""" - if isinstance(node, ast.Constant): - return node.value - - if isinstance(node, ast.Name): - return node.id - - if isinstance(node, ast.List): - return [_literal_from_ast(elt) for elt in node.elts] - - if isinstance(node, ast.Tuple): - return [_literal_from_ast(elt) for elt in node.elts] - - if isinstance(node, ast.Dict): - return { - _literal_from_ast(key): _literal_from_ast(value) - for key, value in zip(node.keys, node.values) - if key is not None - } - - if isinstance(node, ast.UnaryOp) and isinstance(node.op, (ast.USub, ast.UAdd)): - value = _literal_from_ast(node.operand) - if not isinstance(value, (int, float)): - raise ValueError(f"Unsupported unary literal: {ast.dump(node)}") - return -value if isinstance(node.op, ast.USub) else value - - raise ValueError(f"Unsupported literal: {ast.dump(node)}") - - -def _parse_call_expr(expr: str) -> "FunctionToolCall": - """Parse a single ``name(args...)`` expression string into a FunctionToolCall.""" - parsed = ast.parse(expr, mode="eval").body - if not isinstance(parsed, ast.Call): - raise ValueError(f"Expression is not a function call: {expr!r}") - - parameters: "OrderedDict[str, Any]" = OrderedDict() - - for idx, arg in enumerate(parsed.args): - parameters[f"__arg_{idx}__"] = _literal_from_ast(arg) - - for kw in parsed.keywords: - if kw.arg is None: - raise ValueError("**kwargs are not supported") - parameters[kw.arg] = _literal_from_ast(kw.value) - - return FunctionToolCall( - function_name=_extract_function_name(parsed.func), - parameters=parameters, - original_string=expr, - ) - - -# ── Data model ─────────────────────────────────────────────────────── - - -class FunctionToolCall(BaseModel): - """Represents a parsed function call with its parameters.""" - - function_name: str - parameters: Dict[str, Any] - original_string: str - description: str = "" - - def to_realtime_function_tool_call( - self, - function_tools: list[FunctionTool] | None = None, - ) -> ResponseFunctionToolCall: - positional = {k for k in self.parameters if _POSITIONAL_RE.match(k)} - if positional: - logger.warning( - "Dropping positional arguments for '%s': %s", - self.function_name, - positional, - ) - arguments = {k: v for k, v in self.parameters.items() if not _POSITIONAL_RE.match(k)} - - if function_tools is not None: - tool = next( - (t for t in function_tools if t.name == self.function_name), - None, - ) - if tool is None: - available = [t.name for t in function_tools] - raise ValueError(f"Function '{self.function_name}' not found in available tools: {available}") - - schema = tool.parameters if isinstance(tool.parameters, dict) else {} - properties = schema.get("properties", {}) - required = set(schema.get("required", [])) - - undeclared = {k for k in arguments if k not in properties} - if undeclared: - logger.warning( - "Dropping undeclared parameters for '%s': %s", - self.function_name, - undeclared, - ) - arguments = {k: v for k, v in arguments.items() if k in properties} - - missing = required - set(arguments.keys()) - if missing: - raise ValueError(f"Missing required parameters for '{self.function_name}': {missing}") - - return ResponseFunctionToolCall( - name=self.function_name, - arguments=json.dumps(arguments), - call_id=_generate_id("call"), - type="function_call", - id=_generate_id("fc"), - status="in_progress", - ) - - -# ── Public API ─────────────────────────────────────────────────────── - - -def parse_function_call(function_string: str, pattern_to_match: list[str] = []) -> List[FunctionToolCall]: - """Parse a function call string and extract all function calls found. - - Args: - function_string: String representation of function calls. - pattern_to_match: If non-empty, only calls whose function name - contains at least one of these substrings are returned. - - Returns: - List of FunctionToolCall objects with parsed information. - """ - function_string = function_string.strip() - if not function_string: - return [] - - try: - expressions = _split_top_level_calls(function_string) - except tokenize.TokenError: - return _parse_function_exprs( - _split_simple_calls_with_regex(function_string), - pattern_to_match, - skip_invalid=True, - ) - - return _parse_function_exprs(expressions, pattern_to_match) - - -def parse_multiple_functions(function_strings: List[str]) -> List[FunctionToolCall]: - """Parse multiple function call strings. - - Args: - function_strings: List of function call strings. - - Returns: - List of FunctionToolCall objects. - """ - results: List[FunctionToolCall] = [] - for func_str in function_strings: - try: - results.extend(parse_function_call(func_str)) - except Exception: - continue - return results - - -def extract_function_calls_from_text(text: str, block_regex: str = ".*") -> Tuple[str, List[FunctionToolCall]]: - """Extract function calls from delimited code blocks inside *text*. - - The LLM is prompted to wrap tool calls inside code-block delimiters - (e.g. ``func(x=1)``). This function finds those blocks, - parses the function calls within them, and returns the remaining text - (with blocks stripped) alongside the parsed calls. - - Args: - text: Full model output potentially containing code blocks. - block_regex: Regex matching the code-block delimiters **and** their - content (e.g. ``r".*?"``). Only text **inside** - matched blocks is scanned for function calls. - - Returns: - ``(outside_text, function_calls)`` -- the text with blocks stripped - and the parsed function calls found inside the blocks. - """ - if not block_regex: - return text, [] - - matches = list(re.finditer(block_regex, text, flags=re.DOTALL)) - if not matches: - return text, [] - - outside = re.sub(block_regex, "", text, flags=re.DOTALL) - inside = " ".join(match.group(0) for match in matches).strip() - if not inside: - return outside, [] - - try: - return outside, parse_function_call(inside) - except Exception: - return outside, [] diff --git a/runtime/src/speech_to_speech/LLM/tool_call/function_tool.py b/runtime/src/speech_to_speech/LLM/tool_call/function_tool.py deleted file mode 100644 index 56c8519e7be4998ee67cc67eff4f8f193ad7c674..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/LLM/tool_call/function_tool.py +++ /dev/null @@ -1,34 +0,0 @@ -import textwrap - -from openai.types.realtime import RealtimeFunctionTool - -from speech_to_speech.LLM.tool_call.signature_from_schema import signature_from_schema - - -class FunctionTool(RealtimeFunctionTool): - def to_code_prompt(self, include_args_doc: bool = True) -> str: - """Generate a code-style prompt string for this function tool. - - Args: - include_args_doc: If True, include argument descriptions in the docstring. - ⚠️ This lets the model see each argument's purpose but significantly increases - token usage (e.g. 906 tokens without vs 3434 with for the default Reachy Mini - tool profile). Enable depending on the model's capabilities and context limit. - """ - signature = signature_from_schema(self.parameters) - - tool_doc = self.description or "" - - if isinstance(self.parameters, dict) and include_args_doc: - props = self.parameters.get("properties", {}) - if props: - arg_lines = [] - for arg_name, arg_schema in props.items(): - desc = arg_schema.get("description", "") if isinstance(arg_schema, dict) else "" - arg_lines.append(f"{arg_name}: {desc}") - args_doc = f"Args:\n{textwrap.indent(chr(10).join(arg_lines), ' ')}" - tool_doc += f"\n\n{args_doc}" - - tool_doc = f'"""{tool_doc}\n"""' - - return f"def {self.name}{signature}:\n{textwrap.indent(tool_doc, ' ')}" diff --git a/runtime/src/speech_to_speech/LLM/tool_call/signature_from_schema.py b/runtime/src/speech_to_speech/LLM/tool_call/signature_from_schema.py deleted file mode 100644 index dea6d2f3d92b57471d3e760add71011f629c8d5e..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/LLM/tool_call/signature_from_schema.py +++ /dev/null @@ -1,108 +0,0 @@ -import inspect -from typing import Any, Literal, Union - -JSON_TYPE_TO_PYTHON_TYPE = { - "string": str, - "number": float, - "boolean": bool, - "integer": int, - "object": dict, - "array": list, - "null": type(None), -} - - -def _dedupe_types(types: list[Any]) -> list[Any]: - seen = [] - for t in types: - if t not in seen: - seen.append(t) - return seen - - -def _annotation_from_spec(spec: dict[str, Any]) -> Any: - if not spec or not isinstance(spec, dict): - return Any - - # const → Literal[value] - if "const" in spec: - return Literal[spec["const"]] - - # enum → Literal[val1, val2, ...] - if "enum" in spec: - values = spec["enum"] - if not values: - return Any - return Literal[tuple(values)] - - # anyOf / oneOf → Union[Type1, Type2, ...] - for key in ("anyOf", "oneOf"): - if key in spec: - variants = [_annotation_from_spec(s) for s in spec[key]] - unique = _dedupe_types(variants) - if len(unique) == 0: - return Any - if len(unique) == 1: - return unique[0] - return Union[tuple(unique)] - - # allOf → merge sub-schemas then resolve - if "allOf" in spec: - merged = {} - for sub in spec["allOf"]: - merged.update(sub) - return _annotation_from_spec(merged) - - json_type = spec.get("type") - - if json_type is None: - return Any - - # type as list, e.g. ["string", "null"] → Union[str, None] (Optional) - if isinstance(json_type, list): - types = [JSON_TYPE_TO_PYTHON_TYPE.get(t, Any) for t in json_type] - unique = _dedupe_types(types) - if len(unique) == 0: - return Any - if len(unique) == 1: - return unique[0] - return Union[tuple(unique)] - - # array with items → list[ItemType] - if json_type == "array" and "items" in spec: - item_type = _annotation_from_spec(spec["items"]) - return list[item_type] # type: ignore[valid-type] - - return JSON_TYPE_TO_PYTHON_TYPE.get(json_type, Any) - - -def signature_from_schema(schema: object | None) -> inspect.Signature: - if not schema or not isinstance(schema, dict): - return inspect.Signature() - - props = schema.get("properties", {}) - required = set(schema.get("required", [])) - params = [] - - for name, spec in props.items(): - annotation = _annotation_from_spec(spec) - - has_schema_default = "default" in spec if isinstance(spec, dict) else False - - if name in required and not has_schema_default: - default = inspect.Parameter.empty - elif has_schema_default: - default = spec["default"] - else: - default = None - - params.append( - inspect.Parameter( - name, - inspect.Parameter.POSITIONAL_OR_KEYWORD, - default=default, - annotation=annotation, - ) - ) - - return inspect.Signature(params) diff --git a/runtime/src/speech_to_speech/LLM/tool_call/tool_prompt.py b/runtime/src/speech_to_speech/LLM/tool_call/tool_prompt.py deleted file mode 100644 index 5f5cff3e57d4aaeeca1771785d3ef835bcc92784..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/LLM/tool_call/tool_prompt.py +++ /dev/null @@ -1,111 +0,0 @@ -""" -Optional system-prompt builder that instructs a local LLM to output tool calls -inside delimited code blocks (e.g. ``func(arg='val')``). - -The prompt is rendered from a Jinja2 template and relies on -``FunctionTool.to_code_prompt()`` to expose each tool's Python-style signature. -""" - -import re - -from jinja2 import Template - -from speech_to_speech.LLM.tool_call.function_tool import FunctionTool - -# --------------------------------------------------------------------------- -# Default delimiters -# --------------------------------------------------------------------------- -ENTER_CODE = "" -END_CODE = "" - -# --------------------------------------------------------------------------- -# Jinja2 template -# --------------------------------------------------------------------------- -# ``enter_code`` / ``end_code`` are the block delimiters the model must wrap -# every tool call in. ``tools`` is a list of FunctionTool instances whose -# ``.to_code_prompt()`` is called inside the template. -# --------------------------------------------------------------------------- - -TOOL_PROMPT_TEMPLATE = Template( - """\ -Available tools: - -{% for tool in tools %} -{{ tool.to_code_prompt() }} - -{% endfor %} -To call a tool, put exactly one named-argument function call inside {{ enter_code }}...{{ end_code }}: -{{ enter_code }}function_name(required_arg='value'){{ end_code }} - -Rules: -- You may say one brief natural sentence before the tool call; for slow information tools, briefly say that you will check. -- For expression/background tools, always speak first. For requested expressions, use a short pattern like "Sure, here's my best ."; otherwise use a fitting empathetic sentence. -- Do not mention tags, functions, or tools. Keep prose outside tags brief, and do not claim tool results before a tool result is available. -- Use named arguments only; quote strings. Omit optional args instead of placeholder values like "random", "none", "", or null. -- Only one tool call may appear in a response.\ -""", - keep_trailing_newline=True, -) - -# Text-channel variant: same call format and structural rules, without the -# voice-specific "speak first" prose. -TEXT_TOOL_PROMPT_TEMPLATE = Template( - """\ -Available tools: - -{% for tool in tools %} -{{ tool.to_code_prompt() }} - -{% endfor %} -To call a tool, put exactly one named-argument function call inside {{ enter_code }}...{{ end_code }}: -{{ enter_code }}function_name(required_arg='value'){{ end_code }} - -Rules: -- Call a tool directly when it helps fulfill the request; no preamble sentence is required. -- Do not mention tags, functions, or tools in your prose, and do not claim tool results before a tool result is available. -- Use named arguments only; quote strings. Omit optional args instead of placeholder values like "random", "none", "", or null. -- Only one tool call may appear in a response.\ -""", - keep_trailing_newline=True, -) - - -# --------------------------------------------------------------------------- -# Public helpers -# --------------------------------------------------------------------------- - - -def build_tool_system_prompt( - tools: list[FunctionTool], - enter_code: str = ENTER_CODE, - end_code: str = END_CODE, - *, - text_only: bool = False, -) -> str: - """Render the tool-calling system-prompt section. - - Returns an empty string when *tools* is empty so it can be - unconditionally appended to a base system prompt. When *text_only* is set, - the written-channel variant is used (no voice "speak first" prose). - """ - if not tools: - return "" - - template = TEXT_TOOL_PROMPT_TEMPLATE if text_only else TOOL_PROMPT_TEMPLATE - return template.render( - tools=tools, - enter_code=enter_code, - end_code=end_code, - ) - - -def build_block_regex( - enter_code: str = ENTER_CODE, - end_code: str = END_CODE, -) -> str: - """Build a regex that matches a single code block (non-greedy). - - >>> build_block_regex("", "") - '.*?' - """ - return f"{re.escape(enter_code)}.*?{re.escape(end_code)}" diff --git a/runtime/src/speech_to_speech/LLM/utils.py b/runtime/src/speech_to_speech/LLM/utils.py deleted file mode 100644 index b07dac6acec2c98e707f0aa5a1c6e5bd89d416e2..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/LLM/utils.py +++ /dev/null @@ -1,75 +0,0 @@ -import base64 -import io -import re -from typing import Optional - -import requests # type: ignore[import-untyped] -from PIL import Image - -SMART_PUNCT_TRANSLATION = str.maketrans( - { - "\u2018": "'", - "\u2019": "'", - "\u201c": '"', - "\u201d": '"', - } -) - -SPEECHABLE_PATTERN = re.compile( - r"[^\w\s.,!?;:'\"\-()\/\\@#%&*+=$€£¥₹₽¢\[\]{}<>~`^|…—–\n\r\t]", - flags=re.UNICODE, -) - - -def remove_unspeechable(text: str) -> str: - """Keep only speechable characters: letters, digits, punctuation, whitespace. - support unicode characters (english, arabic, chinese, japanese, korean, etc.) - """ - text = text.translate(SMART_PUNCT_TRANSLATION) - return SPEECHABLE_PATTERN.sub("", text) - - -WHISPER_LANGUAGE_TO_LLM_LANGUAGE = { - "en": "english", - "fr": "french", - "es": "spanish", - "zh": "chinese", - "ja": "japanese", - "ko": "korean", - "hi": "hindi", - "de": "german", - "pt": "portuguese", - "pl": "polish", - "it": "italian", - "nl": "dutch", -} - - -def resolve_auto_language(language_code: Optional[str]) -> tuple[Optional[str], Optional[str]]: - """Strip the ``-auto`` suffix and resolve the human-readable language name. - - Returns ``(clean_code, language_name)``. ``language_name`` is non-None - when the code (with or without ``-auto``) maps to a known language. - """ - if not language_code: - return language_code, None - if language_code.endswith("-auto"): - language_code = language_code[:-5] - if language_code not in WHISPER_LANGUAGE_TO_LLM_LANGUAGE: - return language_code, None - return language_code, WHISPER_LANGUAGE_TO_LLM_LANGUAGE.get(language_code) - - -def image_url_to_pil(image_url: str) -> Image.Image: - """Convert an image URL or base64 data URI to a PIL Image. - - Accepts: - - 'data:image/...;base64,' data URIs - - 'https://...`` or ``http://...' URLs (fetched with a 10s timeout) - """ - if image_url.startswith("data:"): - _, b64_data = image_url.split(",", 1) - return Image.open(io.BytesIO(base64.b64decode(b64_data))) - resp = requests.get(image_url, timeout=10) - resp.raise_for_status() - return Image.open(io.BytesIO(resp.content)) diff --git a/runtime/src/speech_to_speech/LLM/voice_prompt.py b/runtime/src/speech_to_speech/LLM/voice_prompt.py deleted file mode 100644 index 29c2f105c86f7ea982467c2cea97d4413109b79e..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/LLM/voice_prompt.py +++ /dev/null @@ -1,48 +0,0 @@ -"""Voice-channel system prompt: lead + session prompt + tail (strongest constraints last).""" - -VOICE_SYSTEM_PROMPT_LEAD = """\ -You are in a spoken conversation. The user speaks and hears you. -The session prompt defines persona, facts, goals, and tool descriptions. These channel rules only control spoken output and tool-use behavior. -""" - -VOICE_SYSTEM_PROMPT_TAIL = """\ -## Voice Rules -- Keep replies brief by default: usually one spoken sentence, two if needed. Go longer only when asked. -- Speak naturally. No markdown, bullets, headings, visual formatting, or action/emote text like *laughs*. -- Treat transcripts as noisy. Correct likely mishearings only if asked or meaning depends on it. -- Speech is the default. Use at most one tool when it helps fulfill the request or clearly fits the moment. -- Before a tool call, use a brief natural utterance unless the user asked for silence or tool-only output. For slow information tools, briefly say that you will check. -- For expression/background tools, speak first. If asked to show an expression, use a short pattern like "Sure, here's my best ." Otherwise use a fitting empathetic sentence. Never mention tools. -- After completed expression/background/physical-action tools, do not add a second spoken comment unless the result has user-facing information. -- Use motion, dance, emotion, and similar tools sparingly when they add empathy, celebration, playfulness, or a requested physical action. -- If unsure whether a tool is needed, just speak. -""" - -# Skeleton for the assembled system message (placeholders filled in build_voice_system_prompt). -_VOICE_SYSTEM_PROMPT_FULL = """\ -{lead} - -Session Prompt: -{session_prompt}{optional_tools} - -{tail} -""" - - -def build_voice_system_prompt(session_prompt: str, *, tool_section: str = "") -> str: - """Context → session prompt → optional tool block → strongest voice rules last.""" - tools = tool_section.strip() - optional_tools = f"\n\n{tools}" if tools else "" - return _VOICE_SYSTEM_PROMPT_FULL.format( - lead=VOICE_SYSTEM_PROMPT_LEAD.rstrip(), - session_prompt=session_prompt.strip(), - optional_tools=optional_tools, - tail=VOICE_SYSTEM_PROMPT_TAIL.rstrip(), - ) - - -# Full voice instructions without a separate session block (legacy / rare direct use). -VOICE_SYSTEM_PROMPT = "{lead}\n\n{tail}".format( - lead=VOICE_SYSTEM_PROMPT_LEAD.rstrip(), - tail=VOICE_SYSTEM_PROMPT_TAIL.rstrip(), -) diff --git a/runtime/src/speech_to_speech/STT/README.md b/runtime/src/speech_to_speech/STT/README.md deleted file mode 100644 index 73e759fc12aaa5fb98590c5685668d0abd846c20..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/STT/README.md +++ /dev/null @@ -1,163 +0,0 @@ -# STT Summary - -This document summarizes the Speech-to-Text (STT) implementations in the `STT/` folder, including language support, language abbreviations, and usage in `s2s_pipeline.py`. - -## Available STT Modes (`--stt`) - -- `whisper` → `STT/whisper_stt_handler.py` -- `whisper-mlx` → `STT/lightning_whisper_mlx_handler.py` -- `mlx-audio-whisper` → `STT/mlx_audio_whisper_handler.py` -- `faster-whisper` → `STT/faster_whisper_handler.py` -- `parakeet-tdt` → `STT/parakeet_tdt_handler.py` -- `paraformer` → `STT/paraformer_handler.py` - -## Language Support by Handler - -### 1) Whisper (`--stt whisper`) - -- Handler: `WhisperSTTHandler` -- Language input flag: `--language` (from shared Whisper args) -- Supports fixed language (e.g. `en`) or `auto` -- Internal supported language list: - - `en`, `fr`, `es`, `zh`, `ja`, `ko`, `hi`, `de`, `pt`, `pl`, `it`, `nl` -- Behavior: - - Detects language from token output - - If detected language is outside the supported list, it falls back to the previous language - -### 2) Lightning Whisper MLX (`--stt whisper-mlx`) - -- Handler: `LightningWhisperSTTHandler` -- Uses same shared `--language` argument as Whisper -- Internal supported language list: - - `en`, `fr`, `es`, `zh`, `ja`, `ko`, `hi`, `de`, `pt`, `pl`, `it`, `nl` -- Behavior: - - If `--language auto`, model auto-detects each utterance - - If detected language is unsupported, falls back to last supported language - -### 3) MLX Audio Whisper (`--stt mlx-audio-whisper`) - -- Handler: `MLXAudioWhisperSTTHandler` -- Model flag: `--mlx_audio_whisper_model_name` -- Language still comes from shared `--language` flag (wired in pipeline) -- Internal supported language list: - - `en`, `fr`, `es`, `zh`, `ja`, `ko`, `hi`, `de`, `pt`, `pl`, `it`, `nl` -- Behavior: - - Uses fixed language unless `--language auto` - - Falls back to last known supported language when needed - -### 4) Faster-Whisper (`--stt faster-whisper`) - -- Handler: `FasterWhisperSTTHandler` -- Language flag: `--faster_whisper_stt_gen_language` -- Default language: `en` -- Note: - - This handler passes generation kwargs directly to `faster_whisper.WhisperModel.transcribe(...)` - - Effective language coverage depends on selected Faster-Whisper/OpenAI Whisper model - -### 5) Parakeet TDT (`--stt parakeet-tdt`) - -- Handler: `ParakeetTDTSTTHandler` -- Language flag: `--parakeet_tdt_language` (optional) -- Supports auto language detection when language not specified -- Declared supported language list (25 European languages): - - `en`, `de`, `fr`, `es`, `it`, `pt`, `nl`, `pl`, `ru`, `uk`, `cs`, `sk`, `hu`, `ro`, `bg`, `hr`, `sl`, `sr`, `da`, `no`, `sv`, `fi`, `et`, `lv`, `lt` -- Backend behavior: - - On macOS/MPS: MLX (`mlx-community/parakeet-tdt-0.6b-v3`) - - On CUDA/CPU: nano-parakeet (`nvidia/parakeet-tdt-0.6b-v3`) - -### 6) Paraformer (`--stt paraformer`) - -- Handler: `ParaformerSTTHandler` -- Model flag: `--paraformer_stt_model_name` -- Default model: `paraformer-zh` -- No dedicated language flag in current args class -- Practical support: - - Depends on selected FunASR model checkpoint - - Default setup is Chinese-oriented (`zh`) - -## Language Abbreviations (ISO-style codes seen in STT handlers) - -| Code | Language | -|---|---| -| `en` | English | -| `fr` | French | -| `es` | Spanish | -| `zh` | Chinese | -| `ja` | Japanese | -| `ko` | Korean | -| `hi` | Hindi | -| `de` | German | -| `pt` | Portuguese | -| `pl` | Polish | -| `it` | Italian | -| `nl` | Dutch | -| `ru` | Russian | -| `uk` | Ukrainian | -| `cs` | Czech | -| `sk` | Slovak | -| `hu` | Hungarian | -| `ro` | Romanian | -| `bg` | Bulgarian | -| `hr` | Croatian | -| `sl` | Slovenian | -| `sr` | Serbian | -| `da` | Danish | -| `no` | Norwegian | -| `sv` | Swedish | -| `fi` | Finnish | -| `et` | Estonian | -| `lv` | Latvian | -| `lt` | Lithuanian | -| `auto` | Per-utterance automatic language detection | - -## Usage Examples - -### Whisper (Transformers) - -```bash -python s2s_pipeline.py --stt whisper --language en -python s2s_pipeline.py --stt whisper --language auto -``` - -### Whisper MLX (LightningWhisperMLX) - -```bash -python s2s_pipeline.py --stt whisper-mlx --language auto --device mps -``` - -### MLX Audio Whisper - -```bash -python s2s_pipeline.py --stt mlx-audio-whisper \ - --mlx_audio_whisper_model_name mlx-community/whisper-large-v3-turbo \ - --language auto -``` - -### Faster-Whisper - -```bash -python s2s_pipeline.py --stt faster-whisper \ - --faster_whisper_stt_model_name large-v3 \ - --faster_whisper_stt_gen_language en -``` - -### Parakeet TDT - -```bash -python s2s_pipeline.py --stt parakeet-tdt --parakeet_tdt_device auto -python s2s_pipeline.py --stt parakeet-tdt --parakeet_tdt_language de -``` - -With live transcription (MLX or CUDA/nano-parakeet backend): - -```bash -python s2s_pipeline.py --stt parakeet-tdt \ - --enable_live_transcription \ - --live_transcription_update_interval 0.25 -``` - -### Paraformer - -```bash -python s2s_pipeline.py --stt paraformer --paraformer_stt_model_name paraformer-zh -``` diff --git a/runtime/src/speech_to_speech/STT/__init__.py b/runtime/src/speech_to_speech/STT/__init__.py deleted file mode 100644 index 8b137891791fe96927ad78e64b0aad7bded08bdc..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/STT/__init__.py +++ /dev/null @@ -1 +0,0 @@ - diff --git a/runtime/src/speech_to_speech/STT/base_stt_handler.py b/runtime/src/speech_to_speech/STT/base_stt_handler.py deleted file mode 100644 index 58524b74c1dc77514c4de6266dedc56a4f79bfe6..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/STT/base_stt_handler.py +++ /dev/null @@ -1,212 +0,0 @@ -from __future__ import annotations - -import logging -from collections import Counter, OrderedDict -from time import perf_counter -from typing import Any - -from speech_to_speech.baseHandler import BaseHandler -from speech_to_speech.pipeline.handler_types import STTIn, STTOut -from speech_to_speech.pipeline.messages import PartialTranscription, Transcription, VADAudio -from speech_to_speech.pipeline.speculative_turns import SpeculativeTurnTracker - -logger = logging.getLogger(__name__) - - -class BaseSTTHandler(BaseHandler[STTIn, STTOut]): - """Base STT handler with speculative-turn stale input filtering.""" - - _MAX_COMPLETED_FINAL_REVISIONS = 2048 - - speculative_turns: SpeculativeTurnTracker | None = None - final_revision_settle_s: float = 0.0 - - def should_process_input(self, item: STTIn) -> bool: - mode = getattr(item, "mode", None) - turn_id = getattr(item, "turn_id", None) - turn_revision = getattr(item, "turn_revision", None) - if self._is_completed_final_revision(item): - queued_drops = self._drop_stale_queued_inputs() - self._log_stale_turn_item(item, "input-after-final", queued_drops=queued_drops) - return False - if mode == "progressive" and self._has_queued_final_for_revision(item): - self._log_stale_turn_item(item, "progressive-before-final") - return False - - wait_for_stability = mode == "final" - gate_start = perf_counter() - is_latest = self._is_latest_turn_item( - item, - wait_for_pending_reopen=True, - wait_for_stability=wait_for_stability, - ) - gate_wait_s = perf_counter() - gate_start - if gate_wait_s >= 0.05: - logger.info( - "%s: STT input gate waited %.3fs for turn=%s rev=%s mode=%s latest=%s age=%.3fs queue=%s", - self.__class__.__name__, - gate_wait_s, - turn_id, - turn_revision, - mode, - is_latest, - self._item_age_s(item), - self._safe_qsize(), - ) - - if not is_latest: - queued_drops = self._drop_stale_queued_inputs() - self._log_stale_turn_item(item, "input", queued_drops=queued_drops) - return False - return True - - def should_emit_output(self, output: STTOut) -> bool: - if isinstance(output, PartialTranscription) and self._is_completed_final_revision(output): - self._log_stale_turn_item(output, "output-after-final") - return False - - if not self._is_latest_turn_item(output, wait_for_pending_reopen=True, wait_for_stability=False): - self._log_stale_turn_item(output, "output") - return False - return True - - def before_emit_output(self, output: STTOut) -> None: - if isinstance(output, Transcription): - self._mark_completed_final_revision(output) - - def _is_latest_turn_item( - self, - item: object, - *, - wait_for_pending_reopen: bool, - wait_for_stability: bool, - ) -> bool: - if self.speculative_turns is None: - return True - turn_id = getattr(item, "turn_id", None) - turn_revision = getattr(item, "turn_revision", None) - if turn_id is None or turn_revision is None: - return True - - if wait_for_stability: - is_latest = self.speculative_turns.is_latest_after_stability_window( - turn_id, - turn_revision, - self.final_revision_settle_s, - ) - elif wait_for_pending_reopen: - is_latest = self.speculative_turns.is_latest_after_pending_reopen(turn_id, turn_revision) - else: - is_latest = self.speculative_turns.is_latest(turn_id, turn_revision) - return is_latest - - def _drop_stale_queued_inputs(self) -> int: - if self.speculative_turns is None or not hasattr(self.queue_in, "mutex") or not hasattr(self.queue_in, "queue"): - return 0 - - dropped = 0 - with self.queue_in.mutex: - kept: list[Any] = [] - while self.queue_in.queue: - queued_item = self.queue_in.queue.popleft() - if isinstance(queued_item, VADAudio) and ( - self._is_completed_final_revision(queued_item) - or (queued_item.mode == "progressive" and self._has_queued_final_for_revision_locked(queued_item)) - or not self._is_latest_turn_item( - queued_item, - wait_for_pending_reopen=False, - wait_for_stability=False, - ) - ): - dropped += 1 - else: - kept.append(queued_item) - self.queue_in.queue.extend(kept) - if dropped: - self.queue_in.not_full.notify_all() - return dropped - - def _log_stale_turn_item(self, item: object, stage: str, *, queued_drops: int = 0) -> None: - turn_id = getattr(item, "turn_id", None) - turn_revision = getattr(item, "turn_revision", None) - if turn_id is None or turn_revision is None: - return - - if not hasattr(self, "_stale_drop_counts"): - self._stale_drop_counts: Counter[tuple[str, str, int]] = Counter() - key = (stage, turn_id, turn_revision) - self._stale_drop_counts[key] += 1 - - message = "%s: dropping stale STT %s for turn=%s rev=%s age=%.3fs" - args: tuple[object, ...] = ( - self.__class__.__name__, - stage, - turn_id, - turn_revision, - self._item_age_s(item), - ) - if queued_drops: - message += " (+%d queued)" - args = (*args, queued_drops) - - if self._stale_drop_counts[key] == 1: - logger.info(message, *args) - else: - logger.debug(message, *args) - - def _item_age_s(self, item: object) -> float: - created_at_s = getattr(item, "created_at_s", None) - if not isinstance(created_at_s, float): - return 0.0 - return max(0.0, perf_counter() - created_at_s) - - def _safe_qsize(self) -> int | str: - try: - return self.queue_in.qsize() - except NotImplementedError: - return "unknown" - - def _has_queued_final_for_revision(self, item: object) -> bool: - if not hasattr(self.queue_in, "mutex") or not hasattr(self.queue_in, "queue"): - return False - with self.queue_in.mutex: - return self._has_queued_final_for_revision_locked(item) - - def _has_queued_final_for_revision_locked(self, item: object) -> bool: - key = self._revision_key(item) - if key is None: - return False - return any( - isinstance(queued_item, VADAudio) and queued_item.mode == "final" and self._revision_key(queued_item) == key - for queued_item in self.queue_in.queue - ) - - def _revision_key(self, item: object) -> tuple[str, int] | None: - turn_id = getattr(item, "turn_id", None) - turn_revision = getattr(item, "turn_revision", None) - if not isinstance(turn_id, str) or not isinstance(turn_revision, int): - return None - return (turn_id, turn_revision) - - def _completed_final_revisions(self) -> OrderedDict[tuple[str, int], None]: - if not hasattr(self, "_completed_final_revision_keys"): - self._completed_final_revision_keys: OrderedDict[tuple[str, int], None] = OrderedDict() - return self._completed_final_revision_keys - - def _is_completed_final_revision(self, item: object) -> bool: - key = self._revision_key(item) - return key is not None and key in self._completed_final_revisions() - - def _mark_completed_final_revision(self, output: Transcription) -> None: - key = self._revision_key(output) - if key is None: - return - completed = self._completed_final_revisions() - completed[key] = None - completed.move_to_end(key) - while len(completed) > self._MAX_COMPLETED_FINAL_REVISIONS: - completed.popitem(last=False) - - def on_session_end(self) -> None: - if hasattr(self, "_completed_final_revision_keys"): - self._completed_final_revision_keys.clear() diff --git a/runtime/src/speech_to_speech/STT/faster_whisper_handler.py b/runtime/src/speech_to_speech/STT/faster_whisper_handler.py deleted file mode 100644 index 291e2149f0289dae8b3cd9f70898596488edb7c0..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/STT/faster_whisper_handler.py +++ /dev/null @@ -1,68 +0,0 @@ -from __future__ import annotations - -import logging -import os -from typing import Any, Iterator - -from faster_whisper import WhisperModel -from rich.console import Console - -from speech_to_speech.pipeline.handler_types import STTIn, STTOut -from speech_to_speech.pipeline.messages import Transcription -from speech_to_speech.STT.base_stt_handler import BaseSTTHandler - -console = Console() - -logger = logging.getLogger(__name__) - - -class FasterWhisperSTTHandler(BaseSTTHandler): - """ - Handles the Speech To Text generation using a Whisper model. - """ - - def setup( - self, - model_name: str = "tiny.en", - device: str = "auto", - compute_type: str = "auto", - gen_kwargs: dict[str, Any] = {}, - ) -> None: - self.gen_kwargs = self.adapt_gen_kwargs(gen_kwargs) - - os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" - self.model = WhisperModel(model_name, device=device, compute_type=compute_type) - - def process(self, vad_audio: STTIn) -> Iterator[STTOut]: - logger.debug("infering faster whisper...") - - segments, info = self.model.transcribe(vad_audio.audio, **self.gen_kwargs) - output_text = [] - - for segment in segments: - logger.debug("[%.2fs -> %.2fs] %s" % (segment.start, segment.end, segment.text)) - output_text.append(segment.text) - - pred_text = " ".join(output_text).strip() - - logger.debug("finished whisper inference") - if pred_text: - console.print(f"[yellow]USER: {pred_text}") - - yield Transcription( - text=pred_text, - turn_id=vad_audio.turn_id, - turn_revision=vad_audio.turn_revision, - speech_stopped_at_s=vad_audio.created_at_s, - ) - else: - logger.debug("no text detected. skipping...") - - def cleanup(self) -> None: - print("Stopping FasterWhisperSTTHandler") - del self.model - - def adapt_gen_kwargs(self, gen_kwargs: dict[str, Any]) -> dict[str, Any]: - gen_kwargs["without_timestamps"] = not gen_kwargs.pop("return_timestamps", True) - - return gen_kwargs diff --git a/runtime/src/speech_to_speech/STT/lightning_whisper_mlx_handler.py b/runtime/src/speech_to_speech/STT/lightning_whisper_mlx_handler.py deleted file mode 100644 index 296ec3190c684456060063d4dcad6a2c99d00f6d..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/STT/lightning_whisper_mlx_handler.py +++ /dev/null @@ -1,103 +0,0 @@ -from __future__ import annotations - -import logging -from typing import Any, Iterator, Optional - -import numpy as np -import torch -from lightning_whisper_mlx import LightningWhisperMLX -from rich.console import Console - -from speech_to_speech.pipeline.handler_types import STTIn, STTOut -from speech_to_speech.pipeline.messages import Transcription -from speech_to_speech.STT.base_stt_handler import BaseSTTHandler - -logger = logging.getLogger(__name__) - -console = Console() - -SUPPORTED_LANGUAGES = [ - "en", - "fr", - "es", - "zh", - "ja", - "ko", - "hi", - "de", - "pt", - "pl", - "it", - "nl", -] - - -class LightningWhisperSTTHandler(BaseSTTHandler): - """ - Handles the Speech To Text generation using a Whisper model. - """ - - def setup( - self, - model_name: str = "distil-large-v3", - device: str = "mps", - torch_dtype: str = "float16", - compile_mode: Optional[str] = None, - language: Optional[str] = None, - gen_kwargs: dict[str, Any] = {}, - ) -> None: - if len(model_name.split("/")) > 1: - model_name = model_name.split("/")[-1] - self.device = device - self.model = LightningWhisperMLX(model=model_name, batch_size=6, quant=None) - self.start_language = language - self.last_language = language - - self.warmup() - - def warmup(self) -> None: - logger.info(f"Warming up {self.__class__.__name__}") - - # 2 warmup steps for no compile or compile mode with CUDA graphs capture - n_steps = 1 - dummy_input = np.array([0] * 512) - - for _ in range(n_steps): - _ = self.model.transcribe(dummy_input)["text"].strip() - - def process(self, vad_audio: STTIn) -> Iterator[STTOut]: - logger.debug("infering whisper...") - - audio = vad_audio.audio - if self.start_language != "auto": - transcription_dict = self.model.transcribe(audio, language=self.start_language) - else: - transcription_dict = self.model.transcribe(audio) - language_code = transcription_dict["language"] - if language_code not in SUPPORTED_LANGUAGES: - logger.warning(f"Whisper detected unsupported language: {language_code}") - if self.last_language in SUPPORTED_LANGUAGES: # reprocess with the last language - transcription_dict = self.model.transcribe(audio, language=self.last_language) - else: - transcription_dict = {"text": "", "language": "en"} - else: - self.last_language = language_code - - pred_text = transcription_dict["text"].strip() - language_code = transcription_dict["language"] - torch.mps.empty_cache() - - logger.debug("finished whisper inference") - console.print(f"[yellow]USER: {pred_text}") - logger.debug(f"Language Code Whisper: {language_code}") - - if self.start_language == "auto": - language_code += "-auto" - - yield Transcription( - text=pred_text, - language_code=language_code, - turn_id=vad_audio.turn_id, - turn_revision=vad_audio.turn_revision, - speech_stopped_at_s=vad_audio.created_at_s, - ) diff --git a/runtime/src/speech_to_speech/STT/mlx_audio_whisper_handler.py b/runtime/src/speech_to_speech/STT/mlx_audio_whisper_handler.py deleted file mode 100644 index 72be3b4bdb32340f61286e0f8032c3d851fc4cde..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/STT/mlx_audio_whisper_handler.py +++ /dev/null @@ -1,174 +0,0 @@ -from __future__ import annotations - -import logging -from typing import Any, Iterator, Optional - -import numpy as np -from rich.console import Console - -from speech_to_speech.pipeline.handler_types import STTIn, STTOut -from speech_to_speech.pipeline.messages import Transcription -from speech_to_speech.STT.base_stt_handler import BaseSTTHandler -from speech_to_speech.utils.mlx_lock import MLXLockContext - -logger = logging.getLogger(__name__) - -console = Console() - -DEFAULT_LANGUAGE = "en" - -SUPPORTED_LANGUAGES = [ - "en", - "fr", - "es", - "zh", - "ja", - "ko", - "hi", - "de", - "pt", - "pl", - "it", - "nl", -] - - -class MLXAudioWhisperSTTHandler(BaseSTTHandler): - """ - Handles the Speech To Text generation using MLX Audio's Whisper implementation. - Optimized for Apple Silicon using the MLX framework. - """ - - def setup( - self, - model_name: str = "mlx-community/whisper-large-v3-turbo", - language: Optional[str] = None, - gen_kwargs: dict[str, Any] = {}, - ) -> None: - from mlx_audio.stt.generate import load_model - from transformers import WhisperProcessor - - self.model_name = model_name - self.start_language = language - # "auto" is a request to detect, not a language code, so it must never leak into - # last_language -- it would fail every SUPPORTED_LANGUAGES check downstream. - self.last_language = language if language != "auto" else None - self.gen_kwargs = gen_kwargs - - # Load the model directly - logger.info(f"Loading model {model_name}...") - self.model = load_model(model_name) - - # Check if processor was loaded, if not, load it manually from original model - if self.model._processor is None: - logger.info("Processor not found in MLX model, loading from original Whisper model...") - # Map MLX model names to their original Whisper counterparts - processor_model_map = { - "mlx-community/whisper-large-v3-turbo": "openai/whisper-large-v3", - "mlx-community/whisper-large-v3": "openai/whisper-large-v3", - "mlx-community/whisper-medium": "openai/whisper-medium", - "mlx-community/whisper-small": "openai/whisper-small", - "mlx-community/whisper-base": "openai/whisper-base", - "mlx-community/whisper-tiny": "openai/whisper-tiny", - } - - # Get the appropriate processor model name - processor_model = processor_model_map.get(model_name, "openai/whisper-large-v3") - logger.info(f"Loading processor from {processor_model}...") - - try: - self.model._processor = WhisperProcessor.from_pretrained(processor_model) - logger.info("Processor loaded successfully") - except Exception as e: - logger.error(f"Failed to load processor: {e}") - raise - - logger.info(f"Model {model_name} loaded successfully") - - self.warmup() - - def warmup(self) -> None: - logger.info(f"Warming up {self.__class__.__name__}") - - # Warmup with a dummy input - dummy_audio = np.zeros(16000, dtype=np.float32) - - try: - # Pre-warm the model by running a transcription - with MLXLockContext(handler_name=self.__class__.__name__): - _ = self.model.generate(dummy_audio, verbose=False) - logger.info("Model warmed up and ready") - except Exception as e: - logger.warning(f"Warmup failed: {e}") - - def _forced_language(self) -> Optional[str]: - """The language explicitly requested by the user, if any.""" - if self.start_language and self.start_language != "auto": - return self.start_language - return None - - def _resolve_language(self, result: Any) -> str: - """Pick the language code to report, updating the sticky fallback on success.""" - forced = self._forced_language() - if forced is not None: - # generate() ran with this language, so it is authoritative. - self.last_language = forced - return forced - - detected = getattr(result, "language", None) - if isinstance(detected, str) and detected: - if detected in SUPPORTED_LANGUAGES: - self.last_language = detected - return detected - logger.warning("Detected unsupported language: %s", detected) - - last_language = self.last_language - if last_language is not None and last_language in SUPPORTED_LANGUAGES: - return last_language - return DEFAULT_LANGUAGE - - def process(self, vad_audio: STTIn) -> Iterator[STTOut]: - logger.debug("inferring mlx-audio whisper...") - - assert isinstance(vad_audio.audio, np.ndarray), "Audio must be a numpy array" - audio_input = vad_audio.audio.astype(np.float32) - - # Prepare generation kwargs - only pass valid parameters - gen_kwargs = {} - - # Add language if specified - forced_language = self._forced_language() - if forced_language is not None: - gen_kwargs["language"] = forced_language - - try: - # MLX models share a single Metal command queue, so concurrent inference from - # the STT/LLM/TTS threads aborts the process with - # "Completed handler provided after commit call". Every other MLX path in the - # pipeline serializes through this lock; this one must too. - with MLXLockContext(handler_name=self.__class__.__name__): - result = self.model.generate(audio_input, verbose=False, **gen_kwargs) - - # Extract text from result - pred_text = result.text.strip() if hasattr(result, "text") else str(result).strip() - language_code = self._resolve_language(result) - - except Exception as e: - logger.error(f"MLX Audio Whisper inference failed: {e}") - pred_text = "" - language_code = self.last_language if self.last_language else DEFAULT_LANGUAGE - - logger.debug("finished mlx-audio whisper inference") - console.print(f"[yellow]USER: {pred_text}") - logger.debug(f"Language Code: {language_code}") - - if self.start_language == "auto": - language_code += "-auto" - - yield Transcription( - text=pred_text, - language_code=language_code, - turn_id=vad_audio.turn_id, - turn_revision=vad_audio.turn_revision, - speech_stopped_at_s=vad_audio.created_at_s, - ) diff --git a/runtime/src/speech_to_speech/STT/paraformer_handler.py b/runtime/src/speech_to_speech/STT/paraformer_handler.py deleted file mode 100644 index 2d8c5f56da791ad426064c6ddf01653dbe9b9a3f..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/STT/paraformer_handler.py +++ /dev/null @@ -1,79 +0,0 @@ -from __future__ import annotations - -import logging -from typing import Any, Iterator - -import numpy as np -import torch -from rich.console import Console - -from speech_to_speech.pipeline.handler_types import STTIn, STTOut -from speech_to_speech.pipeline.messages import PartialTranscription, Transcription -from speech_to_speech.STT.base_stt_handler import BaseSTTHandler - -logging.basicConfig( - format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", -) -logger = logging.getLogger(__name__) - -console = Console() - - -class ParaformerSTTHandler(BaseSTTHandler): - """ - Handles the Speech To Text generation using a Paraformer model. - The default for this model is set to Chinese. - This model was contributed by @wuhongsheng. - """ - - def setup( - self, - model_name: str = "paraformer-zh", - device: str = "cuda", - gen_kwargs: dict[str, Any] = {}, - ) -> None: - print(model_name) - if len(model_name.split("/")) > 1: - model_name = model_name.split("/")[-1] - self.device = device - try: - from funasr import AutoModel - except ModuleNotFoundError as exc: - raise ModuleNotFoundError( - "Paraformer STT requires the optional 'paraformer' extra. " - "Install it with `pip install speech-to-speech[paraformer]`." - ) from exc - self.model = AutoModel(model=model_name, device=device) - self.warmup() - - def warmup(self) -> None: - logger.info(f"Warming up {self.__class__.__name__}") - - # 2 warmup steps for no compile or compile mode with CUDA graphs capture - n_steps = 1 - dummy_input = np.array([0] * 512, dtype=np.float32) - for _ in range(n_steps): - _ = self.model.generate(dummy_input)[0]["text"].strip().replace(" ", "") - - def process(self, vad_audio: STTIn) -> Iterator[STTOut]: - logger.debug("infering paraformer...") - - pred_text = self.model.generate(vad_audio.audio)[0]["text"].strip().replace(" ", "") - torch.mps.empty_cache() - - logger.debug("finished paraformer inference") - console.print(f"[yellow]USER: {pred_text}") - - if vad_audio.mode == "progressive": - yield PartialTranscription( - text=pred_text, - turn_id=vad_audio.turn_id, - turn_revision=vad_audio.turn_revision, - ) - else: - yield Transcription( - text=pred_text, - turn_id=vad_audio.turn_id, - turn_revision=vad_audio.turn_revision, - speech_stopped_at_s=vad_audio.created_at_s, - ) diff --git a/runtime/src/speech_to_speech/STT/parakeet_tdt_handler.py b/runtime/src/speech_to_speech/STT/parakeet_tdt_handler.py deleted file mode 100644 index 989e94c538640e3f211d0fadbbe851444f778693..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/STT/parakeet_tdt_handler.py +++ /dev/null @@ -1,648 +0,0 @@ -""" -Parakeet TDT Speech-to-Text Handler - -Supports NVIDIA Parakeet TDT model for high-quality multilingual ASR. -- On Apple Silicon (MPS): Uses mlx-audio with mlx-community/parakeet-tdt-0.6b-v3 -- On CUDA/CPU: Uses nano-parakeet (pure PyTorch) with nvidia/parakeet-tdt-0.6b-v3 - -Model supports 25 European languages with automatic language detection. -""" - -from __future__ import annotations - -import logging -from contextlib import contextmanager -from sys import platform -from threading import Lock -from time import perf_counter -from typing import Any, Iterator, Optional - -import numpy as np -from rich.console import Console -from rich.text import Text - -from speech_to_speech.pipeline.handler_types import STTIn, STTOut -from speech_to_speech.pipeline.messages import PartialTranscription, Transcription -from speech_to_speech.STT.base_stt_handler import BaseSTTHandler -from speech_to_speech.STT.smart_progressive_streaming import PartialTranscription as ProgressiveStreamPartial -from speech_to_speech.utils.mlx_lock import MLXLockContext - -try: - from lingua import Language, LanguageDetectorBuilder - - LINGUA_AVAILABLE = True -except ImportError: - LINGUA_AVAILABLE = False - -logger = logging.getLogger(__name__) -console = Console() - -# Parakeet TDT v3 supports 25 European languages -SUPPORTED_LANGUAGES = [ - "en", - "de", - "fr", - "es", - "it", - "pt", - "nl", - "pl", - "ru", - "uk", - "cs", - "sk", - "hu", - "ro", - "bg", - "hr", - "sl", - "sr", - "da", - "no", - "sv", - "fi", - "et", - "lv", - "lt", -] - -# Lingua uses "nb" (Bokmål) for Norwegian instead of "no" -_LINGUA_CODE_MAP = {"no": "nb"} - -if LINGUA_AVAILABLE: - _lingua_iso_to_code = { - lang.iso_code_639_1.name.lower(): lang for lang in Language.all() if lang.iso_code_639_1 is not None - } - _lingua_languages = [ - _lingua_iso_to_code[_LINGUA_CODE_MAP.get(code, code)] - for code in SUPPORTED_LANGUAGES - if _LINGUA_CODE_MAP.get(code, code) in _lingua_iso_to_code - ] - - def _build_lingua_detector(): - # Preloading can take multiple seconds on some hardware, including the - # deployed server. Pay that cost at startup instead of on the first user - # request, where it would look like slow STT. - return LanguageDetectorBuilder.from_languages(*_lingua_languages).with_preloaded_language_models().build() - - _lingua_detector = _build_lingua_detector() - - -class ParakeetTDTSTTHandler(BaseSTTHandler): - """ - Handles Speech-to-Text using NVIDIA Parakeet TDT model. - - On Apple Silicon (MPS): Uses mlx-audio with the MLX-converted model. - On CUDA/CPU: Uses nano-parakeet (pure PyTorch) for NeMo-free inference. - - Parakeet TDT 0.6B v3 is a 600M parameter multilingual ASR model - supporting 25 European languages with automatic language detection. - """ - - def setup( - self, - model_name: Optional[str] = None, - device: str = "auto", - compute_type: str = "float16", - language: Optional[str] = None, - gen_kwargs: dict[str, Any] = {}, - enable_live_transcription: bool = False, - live_transcription_update_interval: float = 0.5, - ) -> None: - """ - Initialize the Parakeet TDT model. - - Args: - model_name: Model identifier. Defaults are: - - MPS: "mlx-community/parakeet-tdt-0.6b-v3" - - CUDA/CPU: "nvidia/parakeet-tdt-0.6b-v3" - device: Device to use ("auto", "cuda", "mps", "cpu") - compute_type: Compute precision ("float16", "float32") - language: Target language code (optional, model auto-detects) - gen_kwargs: Additional generation kwargs - """ - self.gen_kwargs = gen_kwargs - self.start_language = language - self.last_language = language if language else "en" - self.enable_live_transcription = enable_live_transcription - self.live_transcription_update_interval = live_transcription_update_interval - self.compute_lock = Lock() - self.sample_rate = 16000 - - # Determine device - if device == "auto": - if platform == "darwin": - self.device = "mps" - else: - import torch - - self.device = "cuda" if torch.cuda.is_available() else "cpu" - else: - self.device = device - - # Set default model based on device - if model_name is None: - if self.device == "mps": - model_name = "mlx-community/parakeet-tdt-0.6b-v3" - else: - model_name = "nvidia/parakeet-tdt-0.6b-v3" - - self.model_name = model_name - self.compute_type = compute_type - - logger.info(f"Loading Parakeet TDT model: {model_name} on {self.device}") - - if self.device == "mps": - self._setup_mlx(model_name) - else: - self._setup_nano_parakeet(model_name) - - # Setup streaming handler if live transcription is enabled - self.streaming_handler = None - if self.enable_live_transcription: - from speech_to_speech.STT.smart_progressive_streaming import ( - SmartProgressiveStreamingHandler, - ) - - self.streaming_handler = SmartProgressiveStreamingHandler( - self.model, - emission_interval=self.live_transcription_update_interval, - max_window_size=15.0, - sentence_buffer=2.0, - ) - self.processing_final = False # Track if we're processing final audio - logger.info(f"Live transcription enabled for Parakeet TDT ({self.backend})") - self._live_transcription_active = False - self._live_turn_key: tuple[str | None, int | None] | None = None - - self.warmup() - - def _setup_mlx(self, model_name: str) -> None: - """Setup for Apple Silicon using mlx-audio.""" - try: - from mlx_audio.stt.generate import load_model - - self.backend = "mlx" - self.model = load_model(model_name) - logger.info("MLX Audio Parakeet model loaded successfully") - except ImportError as e: - raise ImportError( - "mlx-audio is required for Parakeet TDT on Apple Silicon. Install with: pip install mlx-audio" - ) from e - - def _setup_nano_parakeet(self, model_name: str) -> None: - """Setup for CUDA/CPU using nano-parakeet.""" - try: - import torch - from nano_parakeet import from_pretrained - - self.backend = "nano_parakeet" - - if self.device == "cuda" and not torch.cuda.is_available(): - logger.warning("CUDA requested but not available. Falling back to CPU for nano-parakeet.") - self.device = "cpu" - - self.model = from_pretrained(model_name=model_name, device=self.device) - - logger.info(f"nano-parakeet model loaded successfully on {self.device}") - except ImportError as e: - raise ImportError( - "nano-parakeet is required for Parakeet TDT on CUDA/CPU. Install with: pip install nano-parakeet" - ) from e - - def warmup(self) -> None: - """Warm up the model with a dummy input.""" - logger.info(f"Warming up {self.__class__.__name__}") - - # Create 1 second of silence at 16kHz - dummy_audio = np.zeros(16000, dtype=np.float32) - - try: - if self.backend == "mlx": - import mlx.core as mx - - # Convert to mx.array and call decode_chunk directly - audio_mx = mx.array(dummy_audio, dtype=mx.float32) - _ = self.model.decode_chunk(audio_mx, verbose=False) - elif self.backend == "nano_parakeet": - _ = self.model.transcribe(dummy_audio) - else: - _ = self.model.transcribe([dummy_audio], batch_size=1, verbose=False) - - logger.info("Model warmed up and ready") - except Exception as e: - logger.warning(f"Warmup failed: {e}") - - def process(self, vad_audio: STTIn) -> Iterator[STTOut]: - """ - Process audio and generate transcription. - - Yields: - :class:`PartialTranscription` or :class:`Transcription` - """ - process_start_s = perf_counter() - is_progressive = vad_audio.mode == "progressive" - audio_input = vad_audio.audio - - # Ensure audio is float32 numpy array - if not isinstance(audio_input, np.ndarray): - audio_input = np.array(audio_input, dtype=np.float32) - else: - audio_input = audio_input.astype(np.float32) - audio_duration_s = len(audio_input) / getattr(self, "sample_rate", 16000) - item_age_s = self._item_age_s(vad_audio) - - self._prepare_live_transcription_turn(vad_audio.turn_id, vad_audio.turn_revision) - - # Handle progressive updates: yield tagged partial for TranscriptionNotifier - if self.enable_live_transcription and is_progressive: - # Ignore progressive updates if we're already processing final audio - if self.processing_final: - logger.debug("Skipping stale progressive update (final audio already received)") - return - - # Try to acquire lock with short timeout - skip if busy - lock_scope_start_s = perf_counter() - with self._compute_lock_context(handler_name="ParakeetSTT-Progressive", timeout=0.01) as acquired: - if acquired: - try: - inference_start_s = perf_counter() - progressive_text = self._show_progressive_transcription(audio_input) - inference_s = perf_counter() - inference_start_s - if inference_s >= 0.25: - logger.info( - "Parakeet progressive STT timing turn=%s rev=%s audio=%.3fs age=%.3fs " - "lock_scope=%.3fs inference=%.3fs chars=%d", - vad_audio.turn_id, - vad_audio.turn_revision, - audio_duration_s, - item_age_s, - perf_counter() - lock_scope_start_s, - inference_s, - len(progressive_text), - ) - if progressive_text: - yield PartialTranscription( - text=progressive_text, - turn_id=vad_audio.turn_id, - turn_revision=vad_audio.turn_revision, - ) - return - except Exception as e: - logger.debug(f"Progressive transcription failed: {e}") - else: - logger.debug("Skipping progressive update (compute busy)") - return - - # Handle final transcription (send to LLM) - logger.info( - "Parakeet final STT start turn=%s rev=%s audio=%.3fs age=%.3fs", - vad_audio.turn_id, - vad_audio.turn_revision, - audio_duration_s, - item_age_s, - ) - inference_s = 0.0 - lock_scope_s = 0.0 - try: - if self.enable_live_transcription: - # Mark that we're processing final audio (ignore stale progressive updates) - self.processing_final = True - - # Acquire lock with longer timeout for final transcription - lock_scope_start_s = perf_counter() - with self._compute_lock_context(handler_name="ParakeetSTT-Final", timeout=5.0) as acquired: - lock_scope_s = perf_counter() - lock_scope_start_s - if not acquired: - logger.error("Failed to acquire compute lock for final transcription") - pred_text = "" - language_code = self.last_language - else: - inference_start_s = perf_counter() - if self.backend == "mlx": - pred_text, language_code = self._process_mlx_final(audio_input) - else: - pred_text, language_code = self._process_nano_parakeet(audio_input) - inference_s = perf_counter() - inference_start_s - lock_scope_s = perf_counter() - lock_scope_start_s - - # Validate and update language - if language_code and language_code in SUPPORTED_LANGUAGES: - self.last_language = language_code - else: - language_code = self.last_language - - except Exception as e: - logger.error(f"Parakeet TDT inference failed: {e}") - pred_text = "" - language_code = self.last_language - - total_s = perf_counter() - process_start_s - logger.info( - "Parakeet final STT done turn=%s rev=%s total=%.3fs lock_scope=%.3fs inference=%.3fs chars=%d", - vad_audio.turn_id, - vad_audio.turn_revision, - total_s, - lock_scope_s, - inference_s, - len(pred_text), - ) - logger.debug("Finished Parakeet TDT inference") - self._clear_live_transcription_line() - if pred_text.strip(): - console.print(f"[yellow]USER: {pred_text.strip()}") - if language_code: - console.print(f"[dim]Language: {language_code}[/dim]") - - # Reset per-utterance live transcription state only after final STT - # completes. The streaming handler carries fixed sentence timing within - # an utterance, and stale timing must not leak into the next turn. - if self.enable_live_transcription: - self.processing_final = False - self._reset_live_transcription_state(clear_turn=True) - - yield Transcription( - text=pred_text, - language_code=language_code, - turn_id=vad_audio.turn_id, - turn_revision=vad_audio.turn_revision, - speech_stopped_at_s=vad_audio.created_at_s, - ) - - @property - def timing_log_level(self) -> int: - return logging.INFO - - def should_log_timing(self, output: STTOut) -> bool: - return isinstance(output, Transcription) and self.last_time > self.min_time_to_debug - - def _detect_language_from_text(self, text: str) -> Optional[str]: - """ - Detect language from transcribed text using lingua-py. - - Args: - text: Transcribed text string - - Returns: - Detected language code or None if detection fails - """ - if not LINGUA_AVAILABLE: - logger.warning("lingua-py not available, cannot detect language from text") - return None - - # Skip very short utterances where language ID is still too noisy. - if not text or len(text.strip()) < 20: - return None - - detected = _lingua_detector.detect_language_of(text) - if detected is None: - return None - - code = detected.iso_code_639_1.name.lower() - # Map back lingua-specific codes to our supported codes - return {v: k for k, v in _LINGUA_CODE_MAP.items()}.get(code, code) - - @contextmanager - def _compute_lock_context(self, handler_name: str, timeout: float) -> Iterator[bool]: - if self.backend == "mlx": - with MLXLockContext(handler_name=handler_name, timeout=timeout) as acquired: - yield acquired - return - - lock_start_s = perf_counter() - acquired = self.compute_lock.acquire(timeout=timeout) - wait_s = perf_counter() - lock_start_s - hold_start_s: float | None = None - if acquired: - if wait_s >= 0.25: - logger.info("%s: compute lock acquired after %.2fs", handler_name, wait_s) - else: - logger.debug("%s: compute lock acquired after %.3fs", handler_name, wait_s) - hold_start_s = perf_counter() - else: - logger.warning("%s: Failed to acquire compute lock after %.3fs (timeout=%s)", handler_name, wait_s, timeout) - try: - yield acquired - finally: - if acquired: - assert hold_start_s is not None - self.compute_lock.release() - hold_s = perf_counter() - hold_start_s - if hold_s >= 0.25: - logger.info("%s: compute lock released after holding %.2fs", handler_name, hold_s) - else: - logger.debug("%s: compute lock released after holding %.3fs", handler_name, hold_s) - - def _show_progressive_transcription(self, audio_input: np.ndarray) -> str: - """Run progressive transcription, print to console, and return the text.""" - result = self.streaming_handler.transcribe_incremental(audio_input) - rich_text = Text() - if result.fixed_text: - rich_text.append("Live: ", style="dim") - rich_text.append(result.fixed_text, style="yellow") - if result.active_text: - rich_text.append(" ", style="dim") - - if result.active_text: - if not result.fixed_text: - rich_text.append("Live: ", style="dim") - rich_text.append(result.active_text, style="cyan dim") - - progressive_text = self._build_progressive_text(result) - if progressive_text: - self._print_live_transcription(rich_text, progressive_text) - - return progressive_text - - def _print_live_transcription(self, rich_text: Text, progressive_text: str) -> None: - is_terminal = bool(getattr(console, "is_terminal", False)) - if is_terminal: - self._write_live_control("\r\x1b[2K") - if rich_text: - console.print(self._truncate_live_transcription(rich_text), end="") - else: - fallback = Text("Live: ", style="dim") - fallback.append(progressive_text, style="cyan dim") - console.print(self._truncate_live_transcription(fallback), end="") - self._write_live_control("\r") - self._live_transcription_active = True - return - - if rich_text: - console.print(rich_text) - else: - console.print(f"[dim]Live: [/dim]{progressive_text}") - - def _clear_live_transcription_line(self) -> None: - if not getattr(self, "_live_transcription_active", False): - return - self._write_live_control("\r\x1b[2K") - self._live_transcription_active = False - - def _write_live_control(self, sequence: str) -> None: - file = getattr(console, "file", None) - if file is None: - return - file.write(sequence) - file.flush() - - def _truncate_live_transcription(self, text: Text) -> Text: - text = text.copy() - width = getattr(console, "width", 80) - try: - max_width = max(1, int(width) - 1) - except (TypeError, ValueError): - max_width = 79 - text.truncate(max_width, overflow="ellipsis") - return text - - def _prepare_live_transcription_turn(self, turn_id: str | None, turn_revision: int | None) -> None: - if not self.enable_live_transcription: - return - turn_key = (turn_id, turn_revision) - if getattr(self, "_live_turn_key", None) == turn_key: - return - self._reset_live_transcription_state(clear_turn=False) - self._live_turn_key = turn_key - - def _reset_live_transcription_state(self, clear_turn: bool) -> None: - self._clear_live_transcription_line() - streaming_handler = getattr(self, "streaming_handler", None) - if streaming_handler is not None: - streaming_handler.reset() - if clear_turn: - self._live_turn_key = None - - def _build_progressive_text(self, result: ProgressiveStreamPartial) -> str: - parts = [] - if result.fixed_text: - parts.append(result.fixed_text.strip()) - if result.active_text: - parts.append(result.active_text.strip()) - return " ".join(part for part in parts if part).strip() - - def _process_mlx_final(self, audio_input: np.ndarray) -> tuple[str, str]: - """Process final audio using MLX backend with streaming handler.""" - # If we have fixed sentences from progressive updates, only transcribe the new part - if ( - self.streaming_handler is not None - and hasattr(self.streaming_handler, "fixed_sentences") - and self.streaming_handler.fixed_sentences - ): - self._clear_live_transcription_line() - - # Get fixed text from previous progressive updates - fixed_text = " ".join(self.streaming_handler.fixed_sentences).strip() - - # Calculate where fixed part ends in audio - fixed_end_time = self.streaming_handler.fixed_end_time - sample_rate = 16000 - fixed_end_sample = int(fixed_end_time * sample_rate) - if fixed_end_sample > len(audio_input): - logger.warning( - "Ignoring stale progressive fixed text: fixed_end_sample=%d exceeds final audio samples=%d", - fixed_end_sample, - len(audio_input), - ) - pred_text, language_code = self._process_mlx(audio_input) - return pred_text, language_code - - # Only transcribe the part after fixed sentences - if fixed_end_sample < len(audio_input): - remaining_audio = audio_input[fixed_end_sample:] - - import mlx.core as mx - - audio_mx = mx.array(remaining_audio, dtype=mx.float32) - result = self.model.decode_chunk(audio_mx, verbose=False) - - if hasattr(result, "text"): - new_text = result.text.strip() - else: - new_text = str(result).strip() - - # Combine fixed + new - pred_text = f"{fixed_text} {new_text}".strip() if new_text else fixed_text - else: - # All audio already transcribed in progressive updates - pred_text = fixed_text - - # Reset streaming handler for next utterance - self.streaming_handler.reset() - else: - # No progressive updates, transcribe everything - pred_text, language_code = self._process_mlx(audio_input) - return pred_text, language_code - - # Determine language - if self.start_language and self.start_language != "auto": - language_code = self.start_language - else: - detected_lang = self._detect_language_from_text(pred_text) - if detected_lang: - language_code = detected_lang - else: - language_code = self.last_language - - return pred_text, language_code - - def _process_mlx(self, audio_input: np.ndarray) -> tuple[str, str]: - """Process audio using MLX backend.""" - import mlx.core as mx - - # Convert numpy array to mx.array - audio_mx = mx.array(audio_input, dtype=mx.float32) - - # Call decode_chunk directly with the audio array - result = self.model.decode_chunk(audio_mx, verbose=False) - - # Extract text from result - if hasattr(result, "text"): - pred_text = result.text.strip() - else: - pred_text = str(result).strip() - - # Determine language: - # 1. Use fixed language if specified by user - # 2. Try to detect from transcribed text using langdetect - # 3. Fall back to last known language - if self.start_language and self.start_language != "auto": - language_code = self.start_language - else: - # Detect language from transcribed text - detected_lang = self._detect_language_from_text(pred_text) - if detected_lang: - language_code = detected_lang - else: - language_code = self.last_language - - return pred_text, language_code - - def _process_nano_parakeet(self, audio_input: np.ndarray) -> tuple[str, str]: - """Process audio using nano-parakeet backend.""" - pred_text = self.model.transcribe(audio_input).strip() - - if self.start_language and self.start_language != "auto": - language_code = self.start_language - else: - detected_lang = self._detect_language_from_text(pred_text) - if detected_lang: - language_code = detected_lang - else: - language_code = self.last_language - - return pred_text, language_code - - def cleanup(self) -> None: - """Clean up model resources.""" - logger.info(f"Cleaning up {self.__class__.__name__}") - if hasattr(self, "model"): - del self.model - - def on_session_end(self) -> None: - super().on_session_end() - self.last_language = self.start_language if self.start_language else "en" - if self.enable_live_transcription: - self.processing_final = False - self._reset_live_transcription_state(clear_turn=True) - logger.debug("Parakeet TDT session state reset") diff --git a/runtime/src/speech_to_speech/STT/smart_progressive_streaming.py b/runtime/src/speech_to_speech/STT/smart_progressive_streaming.py deleted file mode 100644 index 61e197f210fba16b6fdbb2a326e57e79b7adfa1d..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/STT/smart_progressive_streaming.py +++ /dev/null @@ -1,343 +0,0 @@ -#!/usr/bin/env python3 -""" -Smart Progressive Streaming Handler - -Provides frequent partial transcriptions (every 500ms) with: -- Growing window up to 15s for accuracy -- Sentence-boundary-aware window sliding for audio > 15s -- Fixed sentences + active transcription -""" - -from dataclasses import dataclass -from types import SimpleNamespace -from typing import Any, Generator - -import numpy as np - - -@dataclass -class PartialTranscription: - """Result from progressive streaming.""" - - fixed_text: str # Sentences that won't change - active_text: str # Current partial transcription - timestamp: float # Current position in audio - is_final: bool # True if this is the last update - - -class SmartProgressiveStreamingHandler: - """ - Smart progressive streaming with sentence-aware window management. - - Strategy: - 1. Emit partial transcriptions every 500ms - 2. Use growing window (up to 15s) for better accuracy - 3. When audio > 15s, slide window using sentence boundaries: - - Keep completed sentences as "fixed" - - Only re-transcribe the "active" portion - """ - - def __init__( - self, - model: Any, - emission_interval: float = 0.5, - max_window_size: float = 15.0, - sentence_buffer: float = 2.0, - ) -> None: - """ - Args: - model: Parakeet model with sentence alignment - emission_interval: Emit partial transcription every N seconds (default 500ms) - max_window_size: Maximum window size before sliding (default 15s) - sentence_buffer: Keep last N seconds of sentences in active window (default 2s) - """ - self.model = model - self.emission_interval = emission_interval - self.max_window_size = max_window_size - self.sentence_buffer = sentence_buffer - self.sample_rate = 16000 - - # State for incremental streaming - self.reset() - - def reset(self) -> None: - """Reset state for new streaming session.""" - self.fixed_sentences: list[str] = [] - self.fixed_end_time: float = 0.0 - self.last_transcribed_length: int = 0 - - def _decode_window(self, audio_window: np.ndarray) -> Any: - """ - Decode an audio window and return an object with: - - text: full transcript for the window - - sentences: list of objects with .text and .end (seconds) - """ - if hasattr(self.model, "decode_chunk"): - import mlx.core as mx - - audio_mx = mx.array(audio_window, dtype=mx.float32) - return self.model.decode_chunk(audio_mx, verbose=False) - - # nano-parakeet: use timestamps from transcribe - result = self.model.transcribe(audio_window, timestamps=True) - segments = result.timestamp.get("segment", []) if hasattr(result, "timestamp") else [] - sentences = [SimpleNamespace(text=seg.get("segment", "").strip(), end=seg.get("end", 0.0)) for seg in segments] - return SimpleNamespace(text=getattr(result, "text", ""), sentences=sentences) - - def transcribe_incremental(self, audio: np.ndarray) -> PartialTranscription: - """ - Transcribe audio incrementally (for live streaming). - - Call this repeatedly with growing audio buffer. - Returns a single PartialTranscription for current state. - """ - # Skip if not enough new audio - current_length = len(audio) - if current_length < self.sample_rate * 0.5: # Need at least 500ms - return PartialTranscription( - fixed_text=" ".join(self.fixed_sentences), - active_text="", - timestamp=current_length / self.sample_rate, - is_final=False, - ) - - # Skip if no new audio since last transcription - if current_length == self.last_transcribed_length: - return PartialTranscription( - fixed_text=" ".join(self.fixed_sentences), - active_text="", - timestamp=current_length / self.sample_rate, - is_final=False, - ) - - self.last_transcribed_length = current_length - - # Extract window for transcription (from last fixed sentence to end) - window_start_samples = int(self.fixed_end_time * self.sample_rate) - audio_window = audio[window_start_samples:] - - # Transcribe current window - result = self._decode_window(audio_window) - - # Check if window exceeds max_window_size - window_duration = len(audio_window) / self.sample_rate - - if window_duration >= self.max_window_size and len(result.sentences) > 1: - # Window is too large - fix some sentences - cutoff_time = window_duration - self.sentence_buffer - - # Find sentences to fix - new_fixed_sentences = [] - new_fixed_end_time = self.fixed_end_time - - for sentence in result.sentences: - sentence_abs_time = self.fixed_end_time + sentence.end - - if sentence.end < cutoff_time: - # Fix this sentence - new_fixed_sentences.append(sentence.text.strip()) - new_fixed_end_time = sentence_abs_time - else: - break - - if new_fixed_sentences: - self.fixed_sentences.extend(new_fixed_sentences) - self.fixed_end_time = new_fixed_end_time - - # Re-transcribe from new fixed point - window_start_samples = int(self.fixed_end_time * self.sample_rate) - audio_window = audio[window_start_samples:] - result = self._decode_window(audio_window) - - # Build output - fixed_text = " ".join(self.fixed_sentences) - active_text = result.text.strip() - timestamp = len(audio) / self.sample_rate - - return PartialTranscription(fixed_text=fixed_text, active_text=active_text, timestamp=timestamp, is_final=False) - - def transcribe_progressive(self, audio: np.ndarray) -> Generator[PartialTranscription, None, None]: - """ - Transcribe audio with smart progressive emissions. - - Yields PartialTranscription with: - - fixed_text: Completed sentences (won't change) - - active_text: Current partial transcription - - timestamp: Current position - - is_final: True for last update - """ - position = 0 # Start of current window (in samples) - fixed_sentences = [] # List of completed sentence texts - fixed_end_time = 0.0 # End time of last fixed sentence - - while position < len(audio): - # Determine current window end - remaining = (len(audio) - position) / self.sample_rate - - if remaining <= self.emission_interval: - # Last chunk - process everything remaining - window_end = len(audio) - is_final = True - else: - # Regular chunk - window_end = min(len(audio), position + int(self.emission_interval * self.sample_rate)) - is_final = False - - # Extract window for transcription - # Window includes: fixed_end_time to current position - # (we re-transcribe from last fixed sentence to get better context) - window_start_samples = int(fixed_end_time * self.sample_rate) - audio_window = audio[window_start_samples:window_end] - - # Transcribe current window - result = self._decode_window(audio_window) - - # Check if window exceeds max_window_size - window_duration = (window_end - window_start_samples) / self.sample_rate - - if window_duration >= self.max_window_size and len(result.sentences) > 1: - # Window is too large - need to fix some sentences - # Strategy: Keep sentences up to (max_window - sentence_buffer) as fixed - cutoff_time = window_duration - self.sentence_buffer - - # Find sentences that are completed and before cutoff - new_fixed_sentences = [] - new_fixed_end_time = fixed_end_time - - for sentence in result.sentences: - sentence_abs_time = fixed_end_time + sentence.end - - if sentence.end < cutoff_time: - # This sentence is complete and old enough to fix - new_fixed_sentences.append(sentence.text.strip()) - new_fixed_end_time = sentence_abs_time - else: - # Stop here - keep remaining sentences as active - break - - if new_fixed_sentences: - fixed_sentences.extend(new_fixed_sentences) - fixed_end_time = new_fixed_end_time - - # Re-transcribe from new fixed_end_time - window_start_samples = int(fixed_end_time * self.sample_rate) - audio_window = audio[window_start_samples:window_end] - result = self._decode_window(audio_window) - - # Build output - fixed_text = " ".join(fixed_sentences) - active_text = result.text.strip() - timestamp = window_end / self.sample_rate - - yield PartialTranscription( - fixed_text=fixed_text, active_text=active_text, timestamp=timestamp, is_final=is_final - ) - - # Move to next position - position = window_end - - if is_final: - break - - -def demo_smart_progressive() -> None: - """Demonstrate smart progressive streaming.""" - import time - - import soundfile as sf - from mlx_audio.stt.generate import load_model - - print("=" * 80) - print("SMART PROGRESSIVE STREAMING DEMO") - print("=" * 80) - print("\nFeatures:") - print(" • Partial updates every 500ms") - print(" • Growing window up to 15s") - print(" • Sentence-aware window sliding for long audio") - print() - - # Load model - print("Loading model...") - model = load_model("mlx-community/parakeet-tdt-0.6b-v3") - - # Load audio - audio, sr = sf.read("reachy-voice-test.wav") - if len(audio.shape) > 1: - audio = audio.mean(axis=1) - if sr != 16000: - from scipy import signal - - audio = signal.resample(audio, int(len(audio) * 16000 / sr)) - audio = audio.astype(np.float32) - - # Test cases - test_cases = [ - (audio, "Short (5.8s)"), - (np.concatenate([audio] * 3), "Medium (17.5s) - exceeds 15s window"), - (np.concatenate([audio] * 5), "Long (29.1s) - multiple sentence fixings"), - ] - - for test_audio, label in test_cases: - duration = len(test_audio) / 16000 - - print(f"\n{'=' * 80}") - print(f"TEST: {label} - Duration: {duration:.1f}s") - print("=" * 80) - - handler = SmartProgressiveStreamingHandler( - model, - emission_interval=0.5, # 500ms updates - max_window_size=15.0, # Max 15s window - sentence_buffer=2.0, # Keep 2s of sentences in active window - ) - - print("\nProgressive transcriptions (showing every 4th update for readability):") - print() - - update_count = 0 - start_time = time.perf_counter() - - for result in handler.transcribe_progressive(test_audio): - update_count += 1 - - # Show every other update (every 1s) to keep output manageable - if update_count % 2 == 0 or result.is_final: - elapsed = time.perf_counter() - start_time - marker = "FINAL" if result.is_final else f"Update {update_count}" - - print(f"[{result.timestamp:5.1f}s | {elapsed:.3f}s elapsed] {marker}:") - - if result.fixed_text: - print(f" Fixed: {result.fixed_text[:70]}...") - print(f" Active: {result.active_text[:70]}...") - print() - - total_time = time.perf_counter() - start_time - print(f"Total updates: {update_count}") - print(f"Total time: {total_time:.3f}s ({total_time / duration:.3f}s per audio second)") - - print("\n" + "=" * 80) - print("KEY INSIGHTS") - print("=" * 80) - print(""" -1. Frequent updates (500ms): - - User sees transcription building in real-time - - Low latency feel - -2. Growing window (up to 15s): - - Better accuracy with more context - - No arbitrary chunking for short audio - -3. Sentence-aware sliding (for audio > 15s): - - Completed sentences become "fixed" - - Only active portion re-transcribed - - Maintains accuracy while managing window size - -4. Performance: - - Processing happens DURING user speaking - - Final result available immediately when user stops - """) - - -if __name__ == "__main__": - demo_smart_progressive() diff --git a/runtime/src/speech_to_speech/STT/transcription_notifier.py b/runtime/src/speech_to_speech/STT/transcription_notifier.py deleted file mode 100644 index dd75f9ae9699d18e38f7ff54ac8ae910a97c1229..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/STT/transcription_notifier.py +++ /dev/null @@ -1,103 +0,0 @@ -from __future__ import annotations - -import logging -from queue import Queue -from threading import Event -from typing import Iterator, Union - -from speech_to_speech.api.openai_realtime.runtime_config import RuntimeConfig -from speech_to_speech.baseHandler import BaseHandler -from speech_to_speech.LLM.chat import make_user_message -from speech_to_speech.pipeline.events import PartialTranscriptionEvent, TranscriptionCompletedEvent -from speech_to_speech.pipeline.handler_types import LLMIn, STTOut -from speech_to_speech.pipeline.messages import GenerateResponseRequest, PartialTranscription, Transcription -from speech_to_speech.pipeline.queue_types import TextEventItem - -logger = logging.getLogger(__name__) - - -class TranscriptionNotifier(BaseHandler[STTOut, Union[STTOut, LLMIn]]): - """Sits between STT and LLM. - - For **realtime mode** (no ``runtime_config``): emits transcription events - on ``text_output_queue`` for protocol translation but yields nothing -- the - ``RealtimeService`` builds ``GenerateResponseRequest`` directly. - - For **legacy mode** (``runtime_config`` provided): appends the user - message to ``runtime_config.chat`` and yields a - ``GenerateResponseRequest`` so the LLM handler receives a uniform input - type regardless of pipeline mode. - """ - - def setup( - self, - text_output_queue: Queue[TextEventItem] | None = None, - runtime_config: RuntimeConfig | None = None, - should_listen: Event | None = None, - ) -> None: - self.text_output_queue = text_output_queue - self.runtime_config = runtime_config - self.should_listen = should_listen - - def process(self, transcription: STTOut) -> Iterator[Union[STTOut, LLMIn]]: - if isinstance(transcription, PartialTranscription): - if self.text_output_queue and transcription.text: - self.text_output_queue.put( - PartialTranscriptionEvent( - delta=str(transcription.text), - turn_id=transcription.turn_id, - turn_revision=transcription.turn_revision, - ) - ) - logger.debug("Partial transcription: %s", str(transcription.text)[:80]) - return - - if isinstance(transcription, Transcription): - text = transcription.text - language_code = transcription.language_code - turn_id = transcription.turn_id - turn_revision = transcription.turn_revision - speech_stopped_at_s = transcription.speech_stopped_at_s - else: - text = transcription - language_code = None - turn_id = None - turn_revision = None - speech_stopped_at_s = None - - transcript = str(text) - # Always close the client-visible transcription item. Empty final STT - # results should not trigger the LLM, but clients may already have - # received partial deltas and still need a completed event. - if self.text_output_queue is not None: - self.text_output_queue.put( - TranscriptionCompletedEvent( - transcript=transcript, - language_code=language_code, - turn_id=turn_id, - turn_revision=turn_revision, - speech_stopped_at_s=speech_stopped_at_s, - ) - ) - - if not transcript: - logger.debug("Transcription completed with empty transcript") - if self.should_listen is not None: - self.should_listen.set() - logger.debug("Empty transcription completed; listening re-enabled") - return - - if language_code: - logger.info("Transcription completed (language=%s): %s", language_code, transcript) - else: - logger.info("Transcription completed: %s", transcript) - - if self.runtime_config is not None: - self.runtime_config.chat.add_item(make_user_message(transcript)) - yield GenerateResponseRequest( - runtime_config=self.runtime_config, - language_code=language_code, - turn_id=turn_id, - turn_revision=turn_revision, - speech_stopped_at_s=speech_stopped_at_s, - ) diff --git a/runtime/src/speech_to_speech/STT/whisper_stt_handler.py b/runtime/src/speech_to_speech/STT/whisper_stt_handler.py deleted file mode 100644 index ae1905e54f34d8301ccd41a86fef240e25438c05..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/STT/whisper_stt_handler.py +++ /dev/null @@ -1,147 +0,0 @@ -from __future__ import annotations - -import logging -from copy import copy -from typing import Any, Iterator, Optional - -import numpy as np -import torch -from rich.console import Console -from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor - -from speech_to_speech.pipeline.handler_types import STTIn, STTOut -from speech_to_speech.pipeline.messages import Transcription -from speech_to_speech.STT.base_stt_handler import BaseSTTHandler - -logger = logging.getLogger(__name__) -console = Console() - -SUPPORTED_LANGUAGES = [ - "en", - "fr", - "es", - "zh", - "ja", - "ko", - "hi", - "de", - "pt", - "pl", - "it", - "nl", -] - - -class WhisperSTTHandler(BaseSTTHandler): - """ - Handles the Speech To Text generation using a Whisper model. - """ - - def setup( - self, - model_name: str = "distil-whisper/distil-large-v3", - device: str = "cuda", - torch_dtype: str = "float16", - compile_mode: Optional[str] = None, - language: Optional[str] = None, - gen_kwargs: dict[str, Any] = {}, - ) -> None: - self.device = device - self.torch_dtype = getattr(torch, torch_dtype) - self.compile_mode = compile_mode - self.gen_kwargs = gen_kwargs - self.start_language = language - self.last_language = language if language != "auto" else None - if self.last_language is not None: - self.gen_kwargs["language"] = self.last_language - - self.processor = AutoProcessor.from_pretrained(model_name) - self.model = AutoModelForSpeechSeq2Seq.from_pretrained( - model_name, - torch_dtype=self.torch_dtype, - ).to(device) - - # compile - if self.compile_mode: - self.model.generation_config.cache_implementation = "static" - self.model.forward = torch.compile(self.model.forward, mode=self.compile_mode, fullgraph=True) - self.warmup() - - def prepare_model_inputs(self, audio: np.ndarray) -> torch.Tensor: - input_features = self.processor(audio, sampling_rate=16000, return_tensors="pt").input_features - input_features = input_features.to(self.device, dtype=self.torch_dtype) - - return input_features - - def warmup(self) -> None: - logger.info(f"Warming up {self.__class__.__name__}") - - # 2 warmup steps for no compile or compile mode with CUDA graphs capture - n_steps = 1 if self.compile_mode == "default" else 2 - dummy_input = torch.randn( - (1, self.model.config.num_mel_bins, 3000), - dtype=self.torch_dtype, - device=self.device, - ) - if self.compile_mode not in (None, "default"): - # generating more tokens than previously will trigger CUDA graphs capture - # one should warmup with a number of generated tokens above max tokens targeted for subsequent generation - # hence, having min_new_tokens < max_new_tokens in the future doesn't make sense - warmup_gen_kwargs = { - "min_new_tokens": self.gen_kwargs["max_new_tokens"], # Yes, assign max_new_tokens to min_new_tokens - "max_new_tokens": self.gen_kwargs["max_new_tokens"], - **self.gen_kwargs, - } - else: - warmup_gen_kwargs = self.gen_kwargs - - if self.device == "cuda": - start_event = torch.cuda.Event(enable_timing=True) - end_event = torch.cuda.Event(enable_timing=True) - torch.cuda.synchronize() - start_event.record() - - for _ in range(n_steps): - _ = self.model.generate(dummy_input, **warmup_gen_kwargs) - - if self.device == "cuda": - end_event.record() - torch.cuda.synchronize() - - logger.info( - f"{self.__class__.__name__}: warmed up! time: {start_event.elapsed_time(end_event) * 1e-3:.3f} s" - ) - - def process(self, vad_audio: STTIn) -> Iterator[STTOut]: - logger.debug("infering whisper...") - - input_features = self.prepare_model_inputs(vad_audio.audio) - pred_ids = self.model.generate(input_features, **self.gen_kwargs) - language_code = self.processor.tokenizer.decode(pred_ids[0, 1])[2:-2] # remove "<|" and "|>" - - if language_code not in SUPPORTED_LANGUAGES: # reprocess with the last language - logger.warning("Whisper detected unsupported language: %s", language_code) - gen_kwargs = copy(self.gen_kwargs) - gen_kwargs["language"] = self.last_language - language_code = self.last_language - pred_ids = self.model.generate(input_features, **gen_kwargs) - else: - self.last_language = language_code - - pred_text = self.processor.batch_decode(pred_ids, skip_special_tokens=True, decode_with_timestamps=False)[0] - language_code = self.processor.tokenizer.decode(pred_ids[0, 1])[2:-2] # remove "<|" and "|>" - - logger.debug("finished whisper inference") - console.print(f"[yellow]USER: {pred_text}") - logger.debug(f"Language Code Whisper: {language_code}") - - if self.start_language == "auto": - language_code += "-auto" - - yield Transcription( - text=pred_text, - language_code=language_code, - turn_id=vad_audio.turn_id, - turn_revision=vad_audio.turn_revision, - speech_stopped_at_s=vad_audio.created_at_s, - ) diff --git a/runtime/src/speech_to_speech/TTS/README.md b/runtime/src/speech_to_speech/TTS/README.md deleted file mode 100644 index 33686d40d926a07e1e74522193203c8d6ae63b2e..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/TTS/README.md +++ /dev/null @@ -1,155 +0,0 @@ -# TTS Summary - -## Available TTS modes (`--tts`) - -Runtime-supported values in `s2s_pipeline.py`: - -- `chatTTS` → `chatTTS_handler.py` -- `facebookMMS` → `facebookmms_handler.py` -- `pocket` → `pocket_tts_handler.py` -- `kokoro` → `kokoro_handler.py` -- `qwen3` → `qwen3_tts_handler.py` - -Deprecated TTS implementations, including MeloTTS, live in [`../../../archive/TTS`](../../../archive/TTS) and are no longer wired into `s2s_pipeline.py`. - -## Usage - -### 1) ChatTTS (`--tts chatTTS`) - -Primary args prefix: `--chat_tts_*` - -```bash -python s2s_pipeline.py \ - --tts chatTTS \ - --chat_tts_device cuda \ - --chat_tts_stream true \ - --chat_tts_chunk_size 512 -``` - -### 2) Facebook MMS (`--tts facebookMMS`) - -Primary args prefix: `--facebook_mms_*` plus `--tts_language` - -```bash -python s2s_pipeline.py \ - --tts facebookMMS \ - --facebook_mms_device cuda \ - --tts_language en -``` - -This handler maps STT language codes (e.g. `en`, `fr`, `es`) to MMS model suffixes (e.g. `eng`, `fra`, `spa`) and reloads the model on language changes. - -### 3) Pocket TTS (`--tts pocket`) - -Primary args prefix: `--pocket_tts_*` - -```bash -python s2s_pipeline.py \ - --tts pocket \ - --pocket_tts_voice jean \ - --pocket_tts_device cpu \ - --pocket_tts_sample_rate 16000 -``` - -Available preset voices include: -`alba`, `marius`, `javert`, `jean`, `fantine`, `cosette`, `eponine`, `azelma`. - -### 4) Kokoro (`--tts kokoro`) - -Primary args prefix: `--kokoro_*` - -```bash -python s2s_pipeline.py \ - --tts kokoro \ - --kokoro_device auto \ - --kokoro_voice bm_fable \ - --kokoro_lang_code b -``` - -Behavior: -- Uses MLX backend on Apple Silicon (`mlx-community/Kokoro-82M-bf16`) -- Uses native kokoro pipeline otherwise (`hexgrad/Kokoro-82M`) -- Can auto-switch voice/language based on STT language code mapping - -### 5) Qwen3-TTS (`--tts qwen3`) - -Primary args prefix: `--qwen3_tts_*` - -```bash -python s2s_pipeline.py \ - --tts qwen3 \ - --qwen3_tts_model_name Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \ - --qwen3_tts_device cuda \ - --qwen3_tts_backend ggml \ - --qwen3_tts_speaker Aiden \ - --qwen3_tts_language auto \ - --qwen3_tts_non_streaming_mode True -``` - -Behavior: -- Uses `faster-qwen3-tts` on non-macOS platforms, defaulting to the GGML backend. Pass `--qwen3_tts_backend torch` to use the CUDA-graphs backend instead. -- Uses `mlx-audio` on Apple Silicon and auto-maps `Qwen/...` model IDs to `mlx-community/...`, defaulting to the `6bit` MLX variant unless the model name already pins a suffix. -- Supports MLX quantization overrides on Apple Silicon via `--qwen3_tts_mlx_quantization bf16|4bit|6bit|8bit`. -- Keeps the existing voice-clone/custom-voice/voice-design handler flow intact. -- Defaults to the CustomVoice model with speaker `Aiden`, so no reference audio is required. Voice-clone/base models can still use `--qwen3_tts_ref_audio`. - -Install notes for Linux GGML: -- The default PyPI `qwentts-cpp-python` wheel targets CUDA 12.8. -- If that wheel does not match your CUDA runtime, install one of the Hugging Face wheelhouse builds before installing `speech-to-speech`. - -```bash -pip install "qwentts-cpp-python==0.3.1+cu130" \ - -f https://huggingface.co/datasets/andito/qwentts-cpp-python-wheels/tree/main/whl/cu130 -pip install speech-to-speech -``` - -Available wheelhouse directories include `cu124`, `cu128`, `cu130`, and `cpu`. - -Example for Apple Silicon using the default 6-bit MLX variant: - -```bash -python s2s_pipeline.py \ - --tts qwen3 \ - --qwen3_tts_model_name Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \ - --qwen3_tts_speaker Aiden -``` - -You can override the default and select `bf16`, `4bit`, or `8bit` explicitly: - -```bash -python s2s_pipeline.py \ - --tts qwen3 \ - --qwen3_tts_model_name Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \ - --qwen3_tts_mlx_quantization 4bit \ - --qwen3_tts_speaker Aiden -``` - -To benchmark the Apple Silicon MLX variants side by side: - -```bash -.venv/bin/python benchmark_tts.py \ - --handlers qwen3 \ - --iterations 3 \ - --qwen3_mlx_quantizations bf16 4bit 6bit 8bit -``` - -This will run separate benchmark entries for `qwen3[bf16]`, `qwen3[4bit]`, `qwen3[6bit]`, and `qwen3[8bit]`. - -## Setup - -### Low-latency GPU setup - -```bash -python s2s_pipeline.py \ - --stt whisper \ - --tts pocket -``` - -### Apple Silicon setup - -```bash -python s2s_pipeline.py \ - --local_mac_optimal_settings -``` - -`--tts pocket` and `--tts kokoro` are also valid options on macOS. diff --git a/runtime/src/speech_to_speech/TTS/__init__.py b/runtime/src/speech_to_speech/TTS/__init__.py deleted file mode 100644 index 2f028e293e2247a231dde19d0293cb1942376ff5..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/TTS/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# TTS module diff --git a/runtime/src/speech_to_speech/TTS/chatTTS_handler.py b/runtime/src/speech_to_speech/TTS/chatTTS_handler.py deleted file mode 100644 index 500752db350e348e457d82e9f448cd25bdb9178e..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/TTS/chatTTS_handler.py +++ /dev/null @@ -1,119 +0,0 @@ -from __future__ import annotations - -import logging -from threading import Event -from typing import Any, Iterator - -import ChatTTS -import librosa -import numpy as np -import torch -from rich.console import Console - -from speech_to_speech.baseHandler import BaseHandler -from speech_to_speech.pipeline.cancel_scope import CancelScope -from speech_to_speech.pipeline.handler_types import TTSIn, TTSOut -from speech_to_speech.pipeline.messages import AUDIO_RESPONSE_DONE, EndOfResponse -from speech_to_speech.pipeline.speculative_turns import SpeculativeTurnTracker - -logging.basicConfig( - format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", -) - -logger = logging.getLogger(__name__) - -console = Console() - - -class ChatTTSHandler(BaseHandler[TTSIn, TTSOut]): - def setup( - self, - should_listen: Event, - device: str = "cuda", - gen_kwargs: dict[str, Any] = {}, # Unused - stream: bool = True, - chunk_size: int = 512, - cancel_scope: CancelScope | None = None, - speculative_turns: SpeculativeTurnTracker | None = None, - ) -> None: - self.should_listen = should_listen - self.cancel_scope = cancel_scope - self.speculative_turns = speculative_turns - self.device = device - self.model = ChatTTS.Chat() - self.model.load(compile=False) # Doesn't work for me with True - self.chunk_size = chunk_size - self.stream = stream - rnd_spk_emb = self.model.sample_random_speaker() - self.params_infer_code = ChatTTS.Chat.InferCodeParams( - spk_emb=rnd_spk_emb, - ) - self.warmup() - - def warmup(self) -> None: - logger.info(f"Warming up {self.__class__.__name__}") - _ = self.model.infer("text") - - def process(self, tts_input: TTSIn) -> Iterator[TTSOut]: - speculative_turns = getattr(self, "speculative_turns", None) - if isinstance(tts_input, EndOfResponse): - if speculative_turns and not speculative_turns.is_latest_after_reopen_grace( - tts_input.turn_id, - tts_input.turn_revision, - ): - return - yield AUDIO_RESPONSE_DONE - return - - if speculative_turns and not speculative_turns.is_latest_after_reopen_grace( - tts_input.turn_id, - tts_input.turn_revision, - ): - logger.debug("Dropping stale TTS input for turn=%s rev=%s", tts_input.turn_id, tts_input.turn_revision) - return - if speculative_turns: - speculative_turns.commit(tts_input.turn_id, tts_input.turn_revision) - - text = tts_input.text - - _cancel_gen = self.cancel_scope.generation if self.cancel_scope else None - console.print(f"[green]ASSISTANT: {text}") - if self.device == "mps": - import time - - start = time.time() - torch.mps.synchronize() # Waits for all kernels in all streams on the MPS device to complete. - torch.mps.empty_cache() # Frees all memory allocated by the MPS device. - _ = time.time() - start # Removing this line makes it fail more often. I'm looking into it. - - wavs_gen = self.model.infer(text, params_infer_code=self.params_infer_code, stream=self.stream) - - if self.stream: - wavs = [np.array([])] - for gen in wavs_gen: - if ( - _cancel_gen is not None - and self.cancel_scope is not None - and self.cancel_scope.is_stale(_cancel_gen) - ): - logger.info("TTS generation cancelled (interruption)") - return - if gen[0] is None or len(gen[0]) == 0: - return - audio_chunk = librosa.resample(gen[0], orig_sr=24000, target_sr=16000) - audio_chunk = (audio_chunk * 32768).astype(np.int16)[0] - while len(audio_chunk) > self.chunk_size: - yield audio_chunk[: self.chunk_size] # Return the first chunk_size samples of the audio data - audio_chunk = audio_chunk[self.chunk_size :] # Remove the samples that have already been returned - yield np.pad(audio_chunk, (0, self.chunk_size - len(audio_chunk))) - else: - wavs = wavs_gen - if len(wavs[0]) == 0: - return - audio_chunk = librosa.resample(wavs[0], orig_sr=24000, target_sr=16000) - audio_chunk = (audio_chunk * 32768).astype(np.int16) - for i in range(0, len(audio_chunk), self.chunk_size): - yield np.pad( - audio_chunk[i : i + self.chunk_size], - (0, self.chunk_size - len(audio_chunk[i : i + self.chunk_size])), - ) diff --git a/runtime/src/speech_to_speech/TTS/facebookmms_handler.py b/runtime/src/speech_to_speech/TTS/facebookmms_handler.py deleted file mode 100644 index d5430afc54881673e1638ac96ddb735df26f8574..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/TTS/facebookmms_handler.py +++ /dev/null @@ -1,212 +0,0 @@ -from __future__ import annotations - -import logging -from threading import Event -from typing import Any, Iterator - -import librosa -import numpy as np -import torch -from rich.console import Console -from transformers import AutoTokenizer, VitsModel - -from speech_to_speech.baseHandler import BaseHandler -from speech_to_speech.pipeline.cancel_scope import CancelScope -from speech_to_speech.pipeline.handler_types import TTSIn, TTSOut -from speech_to_speech.pipeline.messages import AUDIO_RESPONSE_DONE, EndOfResponse -from speech_to_speech.pipeline.speculative_turns import SpeculativeTurnTracker - -logging.basicConfig(format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", level=logging.DEBUG) - -logger = logging.getLogger(__name__) - -console = Console() - -WHISPER_LANGUAGE_TO_FACEBOOK_LANGUAGE = { - "en": "eng", # English - "fr": "fra", # French - "es": "spa", # Spanish - "ko": "kor", # Korean - "hi": "hin", # Hindi - "ar": "ara", # Arabic - "hy": "hyw", # Armenian - "az": "azb", # Azerbaijani - "bu": "bul", # Bulgarian - "ca": "cat", # Catalan - "nl": "nld", # Dutch - "fi": "fin", # Finnish - "de": "deu", # German - "el": "ell", # Greek - "he": "heb", # Hebrew - "hu": "hun", # Hungarian - "is": "isl", # Icelandic - "id": "ind", # Indonesian - "ka": "kan", # Kannada - "kk": "kaz", # Kazakh - "lv": "lav", # Latvian - "zl": "zlm", # Malay - "ma": "mar", # Marathi - "fa": "fas", # Persian - "po": "pol", # Polish - "pt": "por", # Portuguese - "ro": "ron", # Romanian - "ru": "rus", # Russian - "sw": "swh", # Swahili - "sv": "swe", # Swedish - "tg": "tgl", # Tagalog - "ta": "tam", # Tamil - "th": "tha", # Thai - "tu": "tur", # Turkish - "uk": "ukr", # Ukrainian - "ur": "urd", # Urdu - "vi": "vie", # Vietnamese - "cy": "cym", # Welsh -} - - -class FacebookMMSTTSHandler(BaseHandler[TTSIn, TTSOut]): - def setup( - self, - should_listen: Event, - device: str = "cuda", - torch_dtype: str = "float32", - language: str = "en", - stream: bool = True, - chunk_size: int = 512, - cancel_scope: CancelScope | None = None, - speculative_turns: SpeculativeTurnTracker | None = None, - **kwargs: Any, - ) -> None: - self.should_listen = should_listen - self.cancel_scope = cancel_scope - self.speculative_turns = speculative_turns - self.device = device - self.torch_dtype = getattr(torch, torch_dtype) - self.stream = stream - self.chunk_size = chunk_size - self.language = language - - self._initial_language = self.language - self.load_model(self.language) - self.warmup() - - def load_model(self, language_code: str) -> None: - try: - model_name = f"facebook/mms-tts-{WHISPER_LANGUAGE_TO_FACEBOOK_LANGUAGE[language_code]}" - logger.info(f"Loading model: {model_name}") - self.model = VitsModel.from_pretrained(model_name).to(self.device) # type: ignore[arg-type] - self.tokenizer = AutoTokenizer.from_pretrained(model_name) - self.language = language_code - except KeyError: - logger.warning(f"Unsupported language: {language_code}. Falling back to English.") - self.load_model("en") - - def warmup(self) -> None: - logger.info(f"Warming up {self.__class__.__name__}") - self.generate_audio("Hello, this is a test") - - def generate_audio(self, text: str) -> torch.Tensor | None: - if not text: - logger.warning("Received empty text input") - return None - - try: - logger.debug(f"Tokenizing text: {text}") - logger.debug(f"Current language: {self.language}") - logger.debug(f"Tokenizer: {self.tokenizer}") - - inputs = self.tokenizer(text, return_tensors="pt", padding=True, truncation=True) - input_ids = inputs.input_ids.to(self.device).long() - attention_mask = inputs.attention_mask.to(self.device) - - logger.debug(f"Input IDs shape: {input_ids.shape}, dtype: {input_ids.dtype}") - logger.debug(f"Input IDs: {input_ids}") - - if input_ids.numel() == 0: - logger.error("Input IDs tensor is empty") - return None - - with torch.no_grad(): - output = self.model(input_ids=input_ids, attention_mask=attention_mask) - - logger.debug(f"Output waveform shape: {output.waveform.shape}") - return output.waveform - except Exception as e: - logger.error(f"Error in generate_audio: {str(e)}") - logger.exception("Full traceback:") - return None - - def process(self, tts_input: TTSIn) -> Iterator[TTSOut]: - speculative_turns = getattr(self, "speculative_turns", None) - if isinstance(tts_input, EndOfResponse): - if speculative_turns and not speculative_turns.is_latest_after_reopen_grace( - tts_input.turn_id, - tts_input.turn_revision, - ): - return - yield AUDIO_RESPONSE_DONE - return - - if speculative_turns and not speculative_turns.is_latest_after_reopen_grace( - tts_input.turn_id, - tts_input.turn_revision, - ): - logger.debug("Dropping stale TTS input for turn=%s rev=%s", tts_input.turn_id, tts_input.turn_revision) - return - if speculative_turns: - speculative_turns.commit(tts_input.turn_id, tts_input.turn_revision) - - gen = self.cancel_scope.generation if self.cancel_scope else None - language_code = tts_input.language_code - text = tts_input.text - - console.print(f"[green]ASSISTANT: {text}") - logger.debug(f"Processing text: {text}") - logger.debug(f"Language code: {language_code}") - - if language_code is not None and self.language != language_code: - try: - logger.info(f"Switching language from {self.language} to {language_code}") - self.load_model(language_code) - except KeyError: - console.print( - f"[red]Language {language_code} not supported by Facebook MMS. Using {self.language} instead." - ) - logger.warning(f"Unsupported language: {language_code}") - - audio_output = self.generate_audio(text) - - if audio_output is None or audio_output.numel() == 0: - logger.warning("No audio output generated") - return - - audio_numpy = audio_output.cpu().numpy().squeeze() - logger.debug(f"Raw audio shape: {audio_numpy.shape}, dtype: {audio_numpy.dtype}") - - audio_resampled = librosa.resample(audio_numpy, orig_sr=self.model.config.sampling_rate, target_sr=16000) - logger.debug(f"Resampled audio shape: {audio_resampled.shape}, dtype: {audio_resampled.dtype}") - - audio_int16 = (audio_resampled * 32768).astype(np.int16) - logger.debug(f"Final audio shape: {audio_int16.shape}, dtype: {audio_int16.dtype}") - - if self.stream: - for i in range(0, len(audio_int16), self.chunk_size): - if gen is not None and self.cancel_scope is not None and self.cancel_scope.is_stale(gen): - logger.info("TTS generation cancelled (interruption)") - return - chunk = audio_int16[i : i + self.chunk_size] - yield np.pad(chunk, (0, self.chunk_size - len(chunk))) - else: - for i in range(0, len(audio_int16), self.chunk_size): - if gen is not None and self.cancel_scope is not None and self.cancel_scope.is_stale(gen): - logger.info("TTS generation cancelled (interruption)") - return - yield np.pad( - audio_int16[i : i + self.chunk_size], - (0, self.chunk_size - len(audio_int16[i : i + self.chunk_size])), - ) - - def on_session_end(self) -> None: - if self.language != self._initial_language: - self.load_model(self._initial_language) - logger.debug("Facebook MMS TTS session state reset") diff --git a/runtime/src/speech_to_speech/TTS/kokoro_handler.py b/runtime/src/speech_to_speech/TTS/kokoro_handler.py deleted file mode 100644 index a93b7a4522df53bade4d34805d0c3063ce20bd57..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/TTS/kokoro_handler.py +++ /dev/null @@ -1,418 +0,0 @@ -""" -Kokoro TTS Handler - -Supports NVIDIA Kokoro TTS model for high-quality multilingual speech synthesis. -- On Apple Silicon (MPS): Uses mlx-audio with mlx-community/Kokoro-82M-bf16 -- On CUDA/CPU: Uses native kokoro library with hexgrad/Kokoro-82M - -Model supports 8 languages with multiple voices per language. -""" - -from __future__ import annotations - -import logging -from sys import platform -from threading import Event -from typing import Any, Iterator, Optional - -import numpy as np -from rich.console import Console - -from speech_to_speech.baseHandler import BaseHandler -from speech_to_speech.pipeline.cancel_scope import CancelScope -from speech_to_speech.pipeline.handler_types import TTSIn, TTSOut -from speech_to_speech.pipeline.messages import AUDIO_RESPONSE_DONE, EndOfResponse -from speech_to_speech.pipeline.speculative_turns import SpeculativeTurnTracker -from speech_to_speech.utils.mlx_lock import MLXLockContext - -logger = logging.getLogger(__name__) -console = Console() - -# Language code mapping from Whisper/langdetect language codes to Kokoro lang codes -WHISPER_LANGUAGE_TO_KOKORO_LANG = { - "en": "b", # British English - "ja": "j", # Japanese - "zh": "z", # Chinese - "fr": "f", # French - "es": "e", # Spanish - "it": "i", # Italian - "pt": "p", # Portuguese - "hi": "h", # Hindi - # Additional European languages from Parakeet - map to closest Kokoro language - "de": "b", # German -> British English (no German voice available) - "nl": "b", # Dutch -> British English - "pl": "b", # Polish -> British English - "ru": "b", # Russian -> British English - "uk": "b", # Ukrainian -> British English -} - -# Default voices for each Kokoro language code -# These are native voices that sound natural for each language -# Voice naming: first letter = language, second letter = gender (f=female, m=male) -# Available voices per language: -# a (American): af_alloy, af_aoede, af_bella, af_heart, af_jessica, af_kore, af_nicole, af_nova, af_river, af_sarah, af_sky -# am_adam, am_echo, am_eric, am_fenrir, am_liam, am_michael, am_onyx, am_puck, am_santa -# b (British): bf_alice, bf_emma, bf_isabella, bf_lily, bm_daniel, bm_fable, bm_george, bm_lewis -# e (Spanish): ef_dora, em_alex, em_santa -# f (French): ff_siwis -# h (Hindi): hf_alpha, hf_beta, hm_omega, hm_psi -# i (Italian): if_sara, im_nicola -# j (Japanese): jf_alpha, jf_gongitsune, jf_nezumi, jf_tebukuro, jm_kumo -# p (Portuguese): pf_dora, pm_alex, pm_santa -# z (Chinese): zf_xiaobei, zf_xiaoni, zf_xiaoxiao, zf_xiaoyi, zm_yunjian, zm_yunxi, zm_yunxia, zm_yunyang -KOKORO_LANG_DEFAULT_VOICES = { - "a": "af_heart", # American English female - "b": "bm_fable", # British English male - "e": "ef_dora", # Spanish female - "f": "ff_siwis", # French female - "h": "hf_alpha", # Hindi female - "i": "if_sara", # Italian female - "j": "jf_alpha", # Japanese female - "p": "pf_dora", # Portuguese female - "z": "zf_xiaobei", # Chinese female -} - - -class KokoroTTSHandler(BaseHandler[TTSIn, TTSOut]): - """ - Handles Text-to-Speech using Kokoro TTS model. - - On Apple Silicon (MPS): Uses mlx-audio with the MLX-converted model. - On CUDA/CPU: Uses native kokoro library for optimal performance. - - Kokoro 82M is a 82M parameter multilingual TTS model - supporting 8 languages with multiple voices per language. - """ - - def setup( - self, - should_listen: Event, - model_name: Optional[str] = None, - device: str = "auto", - voice: str = "bm_fable", - lang_code: str = "b", - speed: float = 1.0, - blocksize: int = 512, - gen_kwargs: dict[str, Any] | None = None, - cancel_scope: CancelScope | None = None, - speculative_turns: SpeculativeTurnTracker | None = None, - ) -> None: - """ - Initialize the Kokoro TTS model. - - Args: - model_name: Model identifier. Defaults are: - - MPS: "mlx-community/Kokoro-82M-bf16" - - CUDA/CPU: "hexgrad/Kokoro-82M" - device: Device to use ("auto", "cuda", "mps", "cpu") - voice: Voice identifier (e.g. "bm_fable") - lang_code: Language code (e.g. "b" for British English) - speed: Speech speed multiplier - blocksize: Audio chunk size for streaming - gen_kwargs: Unused, for pipeline compatibility - """ - self.should_listen = should_listen - self.voice = voice - self.lang_code = lang_code - self.speed = speed - self.blocksize = blocksize - self.cancel_scope = cancel_scope - self.speculative_turns = speculative_turns - - # Determine device - if device == "auto": - if platform == "darwin": - self.device = "mps" - else: - import torch - - self.device = "cuda" if torch.cuda.is_available() else "cpu" - else: - self.device = device - - # Set default model based on device - if model_name is None: - if self.device == "mps": - model_name = "mlx-community/Kokoro-82M-bf16" - else: - model_name = "hexgrad/Kokoro-82M" - - self.model_name = model_name - - logger.info(f"Loading Kokoro model: {model_name} on {self.device}") - - if self.device == "mps": - self._setup_mlx(model_name) - else: - self._setup_kokoro(model_name) - - self._initial_voice = self.voice - self._initial_lang_code = self.lang_code - - self.warmup() - - def _setup_mlx(self, model_name: str) -> None: - """Setup for Apple Silicon using mlx-audio.""" - try: - from mlx_audio.tts.utils import load_model - - self.backend = "mlx" - self.model = load_model(model_name) - logger.info("MLX Audio Kokoro model loaded successfully") - - # Get or create the pipeline for our language and preload the voice - # This avoids the voice being reloaded on every generate() call - self._pipeline = self.model._get_pipeline(self.lang_code) - self._voice_tensor = self._pipeline.load_voice(self.voice) - logger.info(f"Preloaded voice: {self.voice}") - - # Preload voices for common languages to avoid download delays during inference - self._preload_multilingual_voices() - except ImportError as e: - message = str(e) - if "misaki" in message or "espeakng_loader" in message: - raise ImportError( - "Kokoro TTS on Apple Silicon requires additional mlx-audio TTS dependencies. " - f"Missing dependency: {message}. " - "Install with: pip install misaki espeakng-loader" - ) from e - raise ImportError( - "mlx-audio is required for Kokoro TTS on Apple Silicon. Install with: pip install mlx-audio" - ) from e - - def _setup_kokoro(self, model_name: str) -> None: - """Setup for CUDA/CPU using native kokoro library.""" - try: - from kokoro import KPipeline - - self.backend = "kokoro" - self.pipeline = KPipeline(lang_code=self.lang_code) - logger.info("Native Kokoro pipeline loaded successfully") - except ImportError as e: - raise ImportError( - "kokoro is required for Kokoro TTS on CUDA/CPU. " - "Install with: pip install kokoro>=0.9.2 soundfile\n" - "Also ensure espeak-ng is installed: apt-get install espeak-ng (Linux) or brew install espeak-ng (macOS)" - ) from e - - def _preload_multilingual_voices(self) -> None: - """Preload voices for common languages to avoid download delays during inference.""" - # Only preload a few commonly used language voices to avoid excessive startup time - # Users speaking other languages will experience a one-time download delay - # preload_langs = ["a", "b", "e", "f", "h", "i", "j", "p", "z"] # English, Spanish, French, Hindi, Italian, Japanese, Portuguese - preload_langs = ["a", "e", "f"] # English, Spanish, French - - for lang_code in preload_langs: - voice = KOKORO_LANG_DEFAULT_VOICES.get(lang_code) - if voice: - try: - logger.info(f"Preloading voice for language '{lang_code}': {voice}") - pipeline = self.model._get_pipeline(lang_code) - pipeline.load_voice(voice) - except Exception as e: - logger.warning(f"Failed to preload voice {voice}: {e}") - - def warmup(self) -> None: - """Warm up the model with a dummy inference.""" - logger.info(f"Warming up {self.__class__.__name__}") - - # Run a short dummy inference to warm up the model - if self.backend == "mlx": - for _ in self.model.generate( - text="Hello", - voice=self.voice, - speed=self.speed, - lang_code=self.lang_code, - ): - pass - else: - for _ in self.pipeline("Hello", voice=self.voice, speed=self.speed): - pass - - logger.info(f"{self.__class__.__name__} warmed up") - - def process(self, tts_input: TTSIn) -> Iterator[TTSOut]: - """ - Process text input and generate audio output. - - Yields: - Audio chunks as numpy int16 arrays - """ - speculative_turns = getattr(self, "speculative_turns", None) - if isinstance(tts_input, EndOfResponse): - if speculative_turns and not speculative_turns.is_latest_after_reopen_grace( - tts_input.turn_id, - tts_input.turn_revision, - ): - return - yield AUDIO_RESPONSE_DONE - return - - if speculative_turns and not speculative_turns.is_latest_after_reopen_grace( - tts_input.turn_id, - tts_input.turn_revision, - ): - logger.debug("Dropping stale TTS input for turn=%s rev=%s", tts_input.turn_id, tts_input.turn_revision) - return - if speculative_turns: - speculative_turns.commit(tts_input.turn_id, tts_input.turn_revision) - - runtime_config = tts_input.runtime_config - response = tts_input.response - language_code = tts_input.language_code - text = tts_input.text - - voice: Optional[str] = None - if response and response.audio and response.audio.output and response.audio.output.voice: - voice = str(response.audio.output.voice) - if not voice and runtime_config: - audio_cfg = runtime_config.session.audio - audio_output = audio_cfg.output if audio_cfg is not None else None - voice = str(audio_output.voice) if audio_output is not None and audio_output.voice else None - if voice: - self.voice = voice - - if self.backend == "mlx": - yield from self._process_mlx(text, language_code) - else: - yield from self._process_kokoro(text, language_code) - - def _process_mlx(self, llm_sentence: str, language_code: Optional[str] = None) -> Iterator[np.ndarray]: - """Process using MLX backend with Apple Silicon optimizations.""" - from scipy.signal import resample_poly - - gen = self.cancel_scope.generation if self.cancel_scope else None - with MLXLockContext(handler_name="KokoroTTS", timeout=10.0): - if language_code is not None: - new_lang_code = WHISPER_LANGUAGE_TO_KOKORO_LANG.get(language_code, self.lang_code) - if new_lang_code != self.lang_code: - new_voice = KOKORO_LANG_DEFAULT_VOICES.get(new_lang_code, self.voice) - logger.info( - f"Language change detected: {self.lang_code} -> {new_lang_code}, voice: {self.voice} -> {new_voice}" - ) - try: - new_pipeline = self.model._get_pipeline(new_lang_code) - new_voice_tensor = new_pipeline.load_voice(new_voice) - self.lang_code = new_lang_code - self.voice = new_voice - self._pipeline = new_pipeline - self._voice_tensor = new_voice_tensor - except Exception as e: - logger.warning( - f"Failed to switch language/voice: {e}. Keeping current language: {self.lang_code}" - ) - - console.print(f"[green]ASSISTANT: {llm_sentence}") - - # Generate audio using the preloaded pipeline directly - # This avoids the voice reload that happens in model.generate() - for result in self._pipeline( - text=llm_sentence, - voice=self.voice, - speed=self.speed, - ): - if result.audio is None: - continue - # result.audio is an mx.array with shape (1, samples), convert to numpy and squeeze - audio = np.array(result.audio, dtype=np.float32).squeeze(0) - - # Trim silence from start and end of audio - # Kokoro generates ~250ms of silence at the start and variable silence at the end - threshold = 0.01 - abs_audio = np.abs(audio) - # Find first and last samples above threshold - above_threshold = abs_audio > threshold - if np.any(above_threshold): - start_idx = int(np.argmax(above_threshold)) - end_idx = len(audio) - int(np.argmax(above_threshold[::-1])) - # Add small padding (5ms = 120 samples at 24kHz) to avoid cutting speech - padding = 120 - start_idx = max(0, start_idx - padding) - end_idx = min(len(audio), end_idx + padding) - audio = audio[start_idx:end_idx] - - # Kokoro outputs at 24kHz, resample to 16kHz for the pipeline - # Using scipy's polyphase resampling (fast and high quality) - # 16000/24000 = 2/3, so up=2, down=3 - audio = resample_poly(audio, up=2, down=3) - - # Convert to int16 format - audio = (audio * 32768).astype(np.int16) - - # Yield audio in fixed-size chunks - for i in range(0, len(audio), self.blocksize): - if gen is not None and self.cancel_scope is not None and self.cancel_scope.is_stale(gen): - logger.info("TTS generation cancelled (interruption)") - return - chunk = audio[i : i + self.blocksize] - # Pad the last chunk if necessary - if len(chunk) < self.blocksize: - chunk = np.pad(chunk, (0, self.blocksize - len(chunk))) - logger.debug(f"TTS yielding audio chunk: {len(chunk)} samples") - yield chunk - - def _process_kokoro(self, llm_sentence: str, language_code: Optional[str] = None) -> Iterator[np.ndarray]: - """Process using native kokoro library.""" - from scipy.signal import resample_poly - - gen = self.cancel_scope.generation if self.cancel_scope else None - if language_code is not None: - new_lang_code = WHISPER_LANGUAGE_TO_KOKORO_LANG.get(language_code, self.lang_code) - if new_lang_code != self.lang_code: - new_voice = KOKORO_LANG_DEFAULT_VOICES.get(new_lang_code, self.voice) - logger.info( - f"Language change detected: {self.lang_code} -> {new_lang_code}, voice: {self.voice} -> {new_voice}" - ) - self.lang_code = new_lang_code - self.voice = new_voice - from kokoro import KPipeline - - self.pipeline = KPipeline(lang_code=self.lang_code) - - console.print(f"[green]ASSISTANT: {llm_sentence}") - - # Generate audio using Kokoro - # The pipeline yields tuples of (graphemes, phonemes, audio) - for gs, ps, audio in self.pipeline(llm_sentence, voice=self.voice, speed=self.speed): - if audio is None: - continue - - # Ensure audio is numpy array - if not isinstance(audio, np.ndarray): - audio = np.array(audio, dtype=np.float32) - else: - audio = audio.astype(np.float32) - - # Kokoro outputs at 24kHz, resample to 16kHz for the pipeline - # Using scipy's polyphase resampling (fast and high quality) - # 16000/24000 = 2/3, so up=2, down=3 - audio = resample_poly(audio, up=2, down=3) - - # Convert to int16 format - audio = (audio * 32768).astype(np.int16) - - # Yield audio in fixed-size chunks - for i in range(0, len(audio), self.blocksize): - if gen is not None and self.cancel_scope is not None and self.cancel_scope.is_stale(gen): - logger.info("TTS generation cancelled (interruption)") - return - chunk = audio[i : i + self.blocksize] - # Pad the last chunk if necessary - if len(chunk) < self.blocksize: - chunk = np.pad(chunk, (0, self.blocksize - len(chunk))) - yield chunk - - def on_session_end(self) -> None: - self.voice = self._initial_voice - self.lang_code = self._initial_lang_code - if self.backend == "mlx": - try: - self._pipeline = self.model._get_pipeline(self.lang_code) - self._voice_tensor = self._pipeline.load_voice(self.voice) - except Exception as e: - logger.warning(f"Failed to restore initial voice/language on session end: {e}") - else: - from kokoro import KPipeline - - self.pipeline = KPipeline(lang_code=self.lang_code) - logger.debug("Kokoro TTS session state reset") diff --git a/runtime/src/speech_to_speech/TTS/pocket_tts_handler.py b/runtime/src/speech_to_speech/TTS/pocket_tts_handler.py deleted file mode 100644 index a0121e9300aec3c47c15535fa0b034b87c937bb3..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/TTS/pocket_tts_handler.py +++ /dev/null @@ -1,227 +0,0 @@ -from __future__ import annotations - -import logging -from threading import Event -from time import perf_counter -from typing import Any, Iterator - -import numpy as np -from rich.console import Console - -from speech_to_speech.baseHandler import BaseHandler -from speech_to_speech.pipeline.cancel_scope import CancelScope -from speech_to_speech.pipeline.handler_types import TTSIn, TTSOut -from speech_to_speech.pipeline.messages import AUDIO_RESPONSE_DONE, EndOfResponse -from speech_to_speech.pipeline.speculative_turns import SpeculativeTurnTracker - -logger = logging.getLogger(__name__) -console = Console() - - -class PocketTTSHandler(BaseHandler[TTSIn, TTSOut]): - """ - Handler for Pocket TTS model from Kyutai Labs. - Supports streaming audio generation with voice cloning. - """ - - def setup( - self, - should_listen: Event, - device: str = "cpu", - voice: str = "alba", # Default voice from catalog - sample_rate: int = 16000, # Match the pipeline's audio output (LocalAudioStreamer uses 16kHz) - blocksize: int = 512, - max_tokens: int = 50, - gen_kwargs: dict[str, Any] | None = None, # For compatibility with pipeline, not used - cancel_scope: CancelScope | None = None, - speculative_turns: SpeculativeTurnTracker | None = None, - ) -> None: - """ - Initialize Pocket TTS handler. - - Args: - should_listen: Event to control when to start listening again - device: Device to run model on ('cpu', 'cuda', 'mps') - voice: Voice to use. Can be: - - A preset name: 'alba', 'marius', 'javert', 'jean', 'fantine', 'cosette', 'eponine', 'azelma' - - A local audio file path - - A Hugging Face path like "hf://kyutai/tts-voices/..." - sample_rate: Output sample rate (pocket-tts generates at 24kHz and will be resampled to this rate). Default 16kHz matches the pipeline's audio output. - blocksize: Size of audio blocks to yield - max_tokens: Maximum tokens to generate - """ - self.should_listen = should_listen - self.cancel_scope = cancel_scope - self.speculative_turns = speculative_turns - self.device = device - self.voice = voice - self.sample_rate = sample_rate - self.blocksize = blocksize - self.max_tokens = max_tokens - - # Suppress verbose logging from pocket_tts library - logging.getLogger("pocket_tts").setLevel(logging.WARNING) - logging.getLogger("pocket_tts.models.tts_model").setLevel(logging.WARNING) - logging.getLogger("pocket_tts.utils.utils").setLevel(logging.WARNING) - - # Import and load model - from pocket_tts import TTSModel - - logger.info("Loading Pocket TTS model") - self.model = TTSModel.load_model() - - # Move model to specified device - if device == "cuda": - self.model = self.model.cuda() - elif device == "mps": - self.model = self.model.to("mps") - elif device != "cpu": - self.model = self.model.to(device) - - logger.info(f"Pocket TTS model moved to {device}") - - # Load voice state - logger.info(f"Loading voice: {voice}") - self.voice_state = self.model.get_state_for_audio_prompt(voice) - - logger.info(f"Pocket TTS model sample rate: {self.model.sample_rate}") - - @property - def min_time_to_debug(self) -> float: - """ - Override to suppress logging for individual audio chunks. - Pocket TTS yields many small chunks (~10-20ms each), which would flood logs. - Only log if a chunk takes unusually long (>100ms). - """ - return 0.1 # 100ms threshold - - def process(self, tts_input: TTSIn) -> Iterator[TTSOut]: - speculative_turns = getattr(self, "speculative_turns", None) - if isinstance(tts_input, EndOfResponse): - if speculative_turns and not speculative_turns.is_latest_after_reopen_grace( - tts_input.turn_id, - tts_input.turn_revision, - ): - return - yield AUDIO_RESPONSE_DONE - return - - if speculative_turns and not speculative_turns.is_latest_after_reopen_grace( - tts_input.turn_id, - tts_input.turn_revision, - ): - logger.debug("Dropping stale TTS input for turn=%s rev=%s", tts_input.turn_id, tts_input.turn_revision) - return - if speculative_turns: - speculative_turns.commit(tts_input.turn_id, tts_input.turn_revision) - - gen = self.cancel_scope.generation if self.cancel_scope else None - language_code = tts_input.language_code - text = tts_input.text - logger.debug(f"Received language code: {language_code}") - - console.print(f"[green]ASSISTANT: {text}") - - # Generate audio stream - logger.debug(f"Generating audio for: {text[:50]}...") - - pipeline_start = perf_counter() - first_chunk = True - - # Calculate target chunk size in original sample rate - # We need enough samples so that after resampling we get blocksize samples - needs_resampling = self.model.sample_rate != self.sample_rate - if needs_resampling: - from scipy.signal import resample_poly - - # Target chunk size in original sample rate (before resampling) - resample_ratio = self.model.sample_rate / self.sample_rate - target_chunk_size = int(self.blocksize * resample_ratio) - # Calculate up/down factors for polyphase resampling - # For 24kHz -> 16kHz: 16000/24000 = 2/3 - from math import gcd - - g = gcd(self.sample_rate, self.model.sample_rate) - self._resample_up = self.sample_rate // g - self._resample_down = self.model.sample_rate // g - else: - target_chunk_size = self.blocksize - - # Buffer to accumulate audio until we have enough for one block - audio_buffer = [] - buffer_size = 0 - - for audio_chunk in self.model.generate_audio_stream( - self.voice_state, - text, - max_tokens=self.max_tokens, - copy_state=True, # Don't modify the original voice state - ): - if gen is not None and self.cancel_scope is not None and self.cancel_scope.is_stale(gen): - logger.info("TTS generation cancelled (interruption)") - return - if first_chunk: - logger.debug(f"Time to first audio: {perf_counter() - pipeline_start:.3f}s") - first_chunk = False - - # Convert from torch tensor to numpy and add to buffer - audio_np = audio_chunk.cpu().numpy() - audio_buffer.append(audio_np) - buffer_size += len(audio_np) - - # Process buffer when we have enough samples for at least one block - while buffer_size >= target_chunk_size: - # Concatenate buffer - concatenated = np.concatenate(audio_buffer) - - # Extract enough samples for one block - chunk_to_process = concatenated[:target_chunk_size] - remainder = concatenated[target_chunk_size:] - - # Resample this chunk if needed - if needs_resampling: - chunk_resampled = resample_poly( - chunk_to_process, - up=self._resample_up, - down=self._resample_down, - ) - else: - chunk_resampled = chunk_to_process - - # Convert to int16 format expected by audio output - audio_int16 = (chunk_resampled * 32768).astype(np.int16) - - # Ensure exact blocksize (pad or trim if resampling caused slight size differences) - if len(audio_int16) < self.blocksize: - audio_int16 = np.pad(audio_int16, (0, self.blocksize - len(audio_int16))) - elif len(audio_int16) > self.blocksize: - audio_int16 = audio_int16[: self.blocksize] - - yield audio_int16 - - # Update buffer with remainder - audio_buffer = [remainder] if len(remainder) > 0 else [] - buffer_size = len(remainder) - - # Process any remaining audio in buffer - if audio_buffer and buffer_size > 0: - concatenated = np.concatenate(audio_buffer) - - # Resample the remainder - if needs_resampling: - concatenated = resample_poly( - concatenated, - up=self._resample_up, - down=self._resample_down, - ) - - # Convert to int16 - audio_int16 = (concatenated * 32768).astype(np.int16) - - # Yield in blocks - for i in range(0, len(audio_int16), self.blocksize): - chunk = audio_int16[i : i + self.blocksize] - # Pad last chunk if needed - if len(chunk) < self.blocksize: - chunk = np.pad(chunk, (0, self.blocksize - len(chunk))) - yield chunk diff --git a/runtime/src/speech_to_speech/TTS/qwen3_tts_handler.py b/runtime/src/speech_to_speech/TTS/qwen3_tts_handler.py deleted file mode 100644 index e71d6b065d06737f56098968086f6fdb074aa4a0..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/TTS/qwen3_tts_handler.py +++ /dev/null @@ -1,903 +0,0 @@ -""" -Qwen3 TTS Handler - -- On Apple Silicon: Uses mlx-audio with MLX-converted Qwen3-TTS models. -- On CUDA/CPU: Uses faster-qwen3-tts for low-latency streaming. -""" - -from __future__ import annotations - -import logging -import math -import re -import tempfile -import unicodedata -from collections.abc import Callable -from pathlib import Path -from sys import platform -from threading import Event -from time import perf_counter -from typing import Any, Iterator, Optional - -import numpy as np -import torch -from openai.types.realtime.realtime_response_create_params import RealtimeResponseCreateParams -from rich.console import Console - -from speech_to_speech.api.openai_realtime.runtime_config import RuntimeConfig -from speech_to_speech.baseHandler import BaseHandler -from speech_to_speech.pipeline.cancel_scope import CancelScope -from speech_to_speech.pipeline.control import SESSION_END, is_control_message -from speech_to_speech.pipeline.handler_types import TTSIn, TTSOut -from speech_to_speech.pipeline.messages import AUDIO_RESPONSE_DONE, PIPELINE_END, EndOfResponse, TTSInput -from speech_to_speech.pipeline.speculative_turns import SpeculativeTurnTracker -from speech_to_speech.utils.mlx_lock import MLXLockContext - -logger = logging.getLogger(__name__) -console = Console() - -DEFAULT_MODEL = "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice" -DEFAULT_MLX_MODEL = "mlx-community/Qwen3-TTS-12Hz-1.7B-CustomVoice-6bit" -DEFAULT_REF_TEXT = "I'm confused why some people have super short timelines, yet at the same time are bullish on scaling up reinforcement learning atop LLMs. If we're actually close to a human-like learner, then this whole approach of training on verifiable outcomes." -DEFAULT_FASTER_STREAMING_CHUNK_SIZE = 8 -DEFAULT_MLX_STREAMING_CHUNK_SIZE = 4 -DEFAULT_QWEN3_TTS_MAX_NEW_TOKENS = 1536 -MIN_QWEN3_TTS_UTTERANCE_TOKENS = 360 -VALID_MLX_QUANTIZATION_SUFFIXES = ("bf16", "4bit", "6bit", "8bit") -VALID_FASTER_BACKENDS = ("ggml", "torch") -MLX_STREAMING_TOKENS_PER_SECOND = 12.5 -PIPELINE_SR = 16000 -ESTIMATED_QWEN3_WORDS_PER_SECOND = 2.6 -ESTIMATED_QWEN3_CHARS_PER_SECOND = 14.0 -QWEN3_TOKEN_SAFETY_MARGIN = 1.35 -QWEN3_BASE_PROMPT_SECONDS = 1.0 -QWEN3_PUNCTUATION_PAUSE_SECONDS = 0.5 -QWEN3_LANGUAGE_ALIASES = { - "zh": "chinese", - "zh-cn": "chinese", - "zh-hans": "chinese", - "zh-hant": "chinese", - "cmn": "chinese", - "en": "english", - "en-us": "english", - "en-gb": "english", - "ja": "japanese", - "jp": "japanese", - "ko": "korean", - "kr": "korean", - "de": "german", - "fr": "french", - "ru": "russian", - "pt": "portuguese", - "pt-br": "portuguese", - "pt-pt": "portuguese", - "es": "spanish", - "it": "italian", -} - - -class Qwen3TTSHandler(BaseHandler[TTSIn, TTSOut]): - """ - Handles Text-to-Speech using Qwen3-TTS. - - Backend selection: - - Apple Silicon (Darwin): mlx-audio - - Other platforms: faster-qwen3-tts - - Supports three generation modes depending on the loaded model: - - Voice cloning (ref_audio + ref_text) - - Custom voice (preset speakers) - - Voice design (instruct prompt) - """ - - def setup( - self, - should_listen: Event, - model_name: str = DEFAULT_MODEL, - device: str = "cuda", - dtype: str | torch.dtype = "auto", - attn_implementation: str = "eager", - backend: str = "ggml", - ref_audio: str | Path | None = None, - ref_text: str = DEFAULT_REF_TEXT, - language: str = "auto", - speaker: Optional[str] = "Aiden", - instruct: Optional[str] = None, - xvec_only: bool = False, - parity_mode: bool = False, - non_streaming_mode: bool | None = True, - mlx_quantization: Optional[str] = None, - streaming_chunk_size: int | None = None, - max_new_tokens: int = DEFAULT_QWEN3_TTS_MAX_NEW_TOKENS, - blocksize: int = 512, - gen_kwargs: dict[str, Any] | None = None, - cancel_scope: CancelScope | None = None, - speculative_turns: SpeculativeTurnTracker | None = None, - ) -> None: - self.cancel_scope = cancel_scope - self.speculative_turns = speculative_turns - self.should_listen = should_listen - self.requested_device = device - self.ref_audio = ref_audio - self.ref_text = ref_text - self.language = self._normalize_language(language) - self.speaker = speaker - self.instruct = instruct - self.xvec_only = xvec_only - self.parity_mode = parity_mode - self.non_streaming_mode = non_streaming_mode - self.faster_backend = self._normalize_faster_backend(backend) - self.mlx_quantization = self._normalize_mlx_quantization(mlx_quantization) - self.max_new_tokens = max_new_tokens - self.blocksize = blocksize - self.dtype: torch.dtype | None | str = None - self.gen_kwargs = gen_kwargs or {} - self._mlx_ref_audio_cache: dict[str, Any] = {} - self._mlx_temp_ref_audio_files: set[str] = set() - - self.backend = "mlx" if platform == "darwin" else "faster_qwen3_tts" - self.streaming_chunk_size = self._resolve_streaming_chunk_size(streaming_chunk_size) - - if self.backend == "mlx": - self.device = "mps" - self.model_name = self._resolve_mlx_model_name(model_name) - logger.info(f"Loading Qwen3-TTS model: {self.model_name} via mlx-audio on Apple Silicon") - if self.non_streaming_mode is not None: - logger.debug( - "qwen3_tts_non_streaming_mode=%s is ignored on Apple Silicon because " - "mlx-audio does not expose non_streaming_mode yet.", - self.non_streaming_mode, - ) - model_quantization = self._model_name_quantization_suffix(self.model_name) - if model_quantization and model_quantization != "bf16": - logger.info( - "Using MLX quantized Qwen3-TTS variant: %s", - model_quantization, - ) - self._setup_mlx(self.model_name) - else: - self.device = device - self.model_name = model_name - logger.info( - "Loading Qwen3-TTS model: %s via faster-qwen3-tts (%s backend)", - self.model_name, - self.faster_backend, - ) - self._setup_faster( - model_name=self.model_name, - dtype=dtype, - attn_implementation=attn_implementation, - backend=self.faster_backend, - ) - - logger.info( - "Using Qwen3-TTS streaming chunk size %d (~%.0fms audio per chunk) on %s", - self.streaming_chunk_size, - self.streaming_chunk_size / MLX_STREAMING_TOKENS_PER_SECOND * 1000, - self.backend, - ) - - self._initial_speaker = self.speaker - self._initial_ref_audio = self.ref_audio - - self.warmup() - - def _setup_faster(self, model_name: str, dtype: Any, attn_implementation: str, backend: str) -> None: - try: - import torch - except ImportError as e: - raise ImportError("torch is required. Install with: pip install torch") from e - - if dtype == "auto": - self.dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16 - elif isinstance(dtype, str): - self.dtype = getattr(torch, dtype) - else: - self.dtype = dtype - - try: - from faster_qwen3_tts import FasterQwen3TTS - except ImportError as e: - raise ImportError( - "faster-qwen3-tts is required for Qwen3 TTS on non-macOS platforms. " - "Install with: pip install 'faster-qwen3-tts[ggml]'" - ) from e - - self.model = FasterQwen3TTS.from_pretrained( - model_name, - device=self.device, - dtype=self.dtype, - attn_implementation=attn_implementation, - backend=backend, - ) - logger.info("Qwen3-TTS model loaded") - - def _setup_mlx(self, model_name: str) -> None: - try: - from mlx_audio.tts.utils import load_model - - self.model = load_model(model_name) - except ImportError as e: - message = str(e) - if any( - dep in message - for dep in ( - "misaki", - "spacy", - "phonemizer", - "espeakng_loader", - ) - ): - raise ImportError( - "Qwen3-TTS on Apple Silicon requires mlx-audio and its TTS dependencies. " - f"Missing dependency: {message}. " - "Install with: pip install mlx-audio misaki spacy phonemizer-fork espeakng-loader" - ) from e - raise ImportError( - "mlx-audio is required for Qwen3 TTS on Apple Silicon. Install with: pip install mlx-audio" - ) from e - - logger.info("MLX Audio Qwen3-TTS model loaded") - - def _normalize_faster_backend(self, backend: Any) -> str: - value = str(backend or "ggml").strip().lower() - if value not in VALID_FASTER_BACKENDS: - raise ValueError( - f"Unsupported qwen3_tts_backend value {backend!r}. Supported values: {', '.join(VALID_FASTER_BACKENDS)}" - ) - return value - - def _normalize_mlx_quantization(self, mlx_quantization: Any) -> Optional[str]: - if mlx_quantization is None: - return None - - value = str(mlx_quantization).strip().lower() - if value in ("", "none", "default"): - return None - if value not in VALID_MLX_QUANTIZATION_SUFFIXES: - raise ValueError( - "Unsupported qwen3_tts_mlx_quantization value " - f"{mlx_quantization!r}. Supported values: {', '.join(VALID_MLX_QUANTIZATION_SUFFIXES)}" - ) - return value - - def _apply_mlx_quantization_suffix(self, model_name: str) -> str: - if self.mlx_quantization is None: - return model_name - - desired_suffix = f"-{self.mlx_quantization}" - for suffix in VALID_MLX_QUANTIZATION_SUFFIXES: - current_suffix = f"-{suffix}" - if model_name.endswith(current_suffix): - return model_name[: -len(current_suffix)] + desired_suffix - - return f"{model_name}{desired_suffix}" - - def _model_name_quantization_suffix(self, model_name: str) -> Optional[str]: - if not model_name: - return None - - for suffix in VALID_MLX_QUANTIZATION_SUFFIXES: - if model_name.endswith(f"-{suffix}"): - return suffix - - return None - - def _resolve_mlx_model_name(self, model_name: str) -> str: - if not model_name: - return self._apply_mlx_quantization_suffix(DEFAULT_MLX_MODEL) - if model_name.startswith("mlx-community/"): - if self.mlx_quantization is None and self._model_name_quantization_suffix(model_name) is None: - return f"{model_name}-6bit" - return self._apply_mlx_quantization_suffix(model_name) - if model_name.startswith("Qwen/"): - mapped = model_name.replace("Qwen/", "mlx-community/", 1) - if self._model_name_quantization_suffix(mapped) is None: - if self.mlx_quantization is None: - return f"{mapped}-6bit" - mapped = f"{mapped}-bf16" - return self._apply_mlx_quantization_suffix(mapped) - return model_name - - def _resolve_streaming_chunk_size(self, streaming_chunk_size: int | None) -> int: - if streaming_chunk_size is not None: - return max(1, int(streaming_chunk_size)) - if self.backend == "mlx": - return DEFAULT_MLX_STREAMING_CHUNK_SIZE - return DEFAULT_FASTER_STREAMING_CHUNK_SIZE - - def _normalize_language(self, language: str | None) -> str: - if language is None: - return "auto" - normalized = str(language).strip().replace("_", "-").lower() - if not normalized: - return "auto" - return QWEN3_LANGUAGE_ALIASES.get(normalized, normalized) - - def _infer_model_type_from_name(self) -> str: - name = (self.model_name or "").lower() - if "voicedesign" in name: - return "voice_design" - if "customvoice" in name: - return "custom_voice" - return "base" - - def _resolve_audio_path(self, audio: Any) -> Path | None: - if not isinstance(audio, (str, Path)) or not audio: - return None - - candidate = Path(audio).expanduser() - repo_root = Path(__file__).resolve().parents[1] - search_paths = [] - - if candidate.is_absolute(): - search_paths.append(candidate) - else: - search_paths.append(Path.cwd() / candidate) - search_paths.append(repo_root / candidate) - - seen = set() - for path in search_paths: - normalized = str(path) - if normalized in seen: - continue - seen.add(normalized) - if path.exists(): - return path.resolve() - - return None - - def _prepare_mlx_ref_audio(self, ref_audio: Any) -> Any: - if self.backend != "mlx" or ref_audio is None: - return ref_audio - - if not isinstance(ref_audio, (str, Path)): - return ref_audio - - resolved_path = self._resolve_audio_path(ref_audio) - if resolved_path is None: - raise FileNotFoundError( - "Qwen3-TTS on Apple Silicon requires qwen3_tts_ref_audio to point to " - f"a readable audio file. Got: {ref_audio!r}" - ) - - cache_key = str(resolved_path) - cached_path = self._mlx_ref_audio_cache.get(cache_key) - if cached_path and Path(cached_path).exists(): - return cached_path - - try: - import soundfile as sf - from scipy.signal import resample_poly - - waveform, sample_rate = sf.read(str(resolved_path), always_2d=False, dtype="float32") - waveform = np.asarray(waveform, dtype=np.float32) - if waveform.ndim > 1: - waveform = waveform.mean(axis=1) - target_sample_rate = getattr(self.model, "sample_rate", 24000) - if sample_rate != target_sample_rate: - gcd = np.gcd(int(sample_rate), int(target_sample_rate)) - waveform = resample_poly( - waveform, - up=int(target_sample_rate) // gcd, - down=int(sample_rate) // gcd, - ) - sample_rate = target_sample_rate - - with tempfile.NamedTemporaryFile( - prefix="qwen3_ref_", - suffix=".wav", - delete=False, - ) as temp_file: - normalized_path = temp_file.name - - sf.write( - normalized_path, - waveform, - sample_rate, - format="WAV", - subtype="PCM_16", - ) - except Exception as e: - raise RuntimeError(f"Failed to normalize Qwen3-TTS reference audio {resolved_path}: {e}") from e - - self._mlx_ref_audio_cache[cache_key] = normalized_path - self._mlx_temp_ref_audio_files.add(normalized_path) - return normalized_path - - def _apply_session_voice_override( - self, - model_type: str, - runtime_config: RuntimeConfig | None = None, - response: RealtimeResponseCreateParams | None = None, - ) -> None: - session_voice: Optional[str] = None - if response and response.audio and response.audio.output: - resp_voice = response.audio.output.voice - session_voice = str(resp_voice) if resp_voice else None - if not session_voice and runtime_config is not None: - audio = runtime_config.session.audio - output = audio.output if audio is not None else None - sess_voice = output.voice if output is not None else None - session_voice = str(sess_voice) if sess_voice else None - if not session_voice: - return - - if model_type == "custom_voice": - supported_speakers = self._supported_speakers() - if supported_speakers is not None: - speakers_by_lower = {speaker.lower(): speaker for speaker in supported_speakers} - speaker = speakers_by_lower.get(session_voice.lower()) - if speaker is None: - supported_text = ", ".join(sorted(supported_speakers, key=str.lower)) or "unknown" - logger.warning( - "Ignoring Qwen3-TTS session voice override %r because it is not a supported " - "CustomVoice speaker. Supported speakers: %s", - session_voice, - supported_text, - ) - return - session_voice = speaker - - self.speaker = session_voice - self.ref_audio = None - return - - if self._resolve_audio_path(session_voice) is not None: - self.ref_audio = session_voice - return - - logger.warning( - "Ignoring Qwen3-TTS session voice override because it is not an audio file path: %r", - session_voice, - ) - - def warmup(self) -> None: - logger.info(f"Warming up {self.__class__.__name__}") - - if self.backend == "faster_qwen3_tts": - if self.parity_mode: - logger.info("Qwen3-TTS parity mode enabled: skipping backend warmup") - else: - try: - self.model.warmup(prefill_len=100) - except Exception as e: - logger.warning("Qwen3-TTS backend warmup failed: %s", e) - - try: - for _ in self._warmup_process("Hello, this is a warmup."): - pass - logger.info(f"{self.__class__.__name__} warmed up") - except Exception as e: - logger.warning(f"Warmup generation failed: {e}") - - def _model_type(self) -> str: - if self.backend == "mlx": - config = getattr(self.model, "config", None) - return getattr(config, "tts_model_type", None) or self._infer_model_type_from_name() - - inner = getattr(getattr(self.model, "model", None), "model", None) - return getattr(inner, "tts_model_type", None) or self._infer_model_type_from_name() - - def _supported_speakers(self) -> list[str] | None: - for candidate in ( - self.model, - getattr(self.model, "model", None), - getattr(getattr(self.model, "model", None), "model", None), - ): - get_speakers = getattr(candidate, "get_supported_speakers", None) - if callable(get_speakers): - speakers = get_speakers() - if speakers is None: - return None - return [str(speaker) for speaker in speakers if speaker] - return None - - def _resolve_speaker(self) -> Optional[str]: - if self.speaker: - return self.speaker - - for candidate in ( - self.model, - getattr(getattr(self.model, "model", None), "model", None), - ): - get_speakers = getattr(candidate, "get_supported_speakers", None) - if callable(get_speakers): - speakers = list(get_speakers() or []) - if speakers: - return speakers[0] - - return None - - def _to_int16(self, audio: np.ndarray) -> np.ndarray: - return np.clip(audio * 32768, -32768, 32767).astype(np.int16) - - def _estimate_max_new_tokens(self, text: Optional[str]) -> int: - text = (text or "").strip() - chunk_size = max(1, int(getattr(self, "streaming_chunk_size", 1))) - configured_cap = max(1, int(getattr(self, "max_new_tokens", DEFAULT_QWEN3_TTS_MAX_NEW_TOKENS))) - - if not text: - return min(configured_cap, MIN_QWEN3_TTS_UTTERANCE_TOKENS) - - word_count = len(re.findall(r"\w+", text, flags=re.UNICODE)) - char_count = len(re.sub(r"\s+", "", text)) - word_seconds = word_count / ESTIMATED_QWEN3_WORDS_PER_SECOND if word_count else 0.0 - char_seconds = char_count / ESTIMATED_QWEN3_CHARS_PER_SECOND if char_count else 0.0 - punctuation_count = sum(unicodedata.category(ch).startswith("P") for ch in text) - punctuation_seconds = punctuation_count * QWEN3_PUNCTUATION_PAUSE_SECONDS - estimated_seconds = max(word_seconds, char_seconds) + punctuation_seconds + QWEN3_BASE_PROMPT_SECONDS - estimated_tokens = math.ceil(estimated_seconds * MLX_STREAMING_TOKENS_PER_SECOND * QWEN3_TOKEN_SAFETY_MARGIN) - aligned_tokens = max( - chunk_size, - math.ceil(estimated_tokens / chunk_size) * chunk_size, - ) - requested_tokens = max(MIN_QWEN3_TTS_UTTERANCE_TOKENS, aligned_tokens) - resolved_tokens = min(configured_cap, requested_tokens) - - if resolved_tokens < requested_tokens: - logger.warning( - "Qwen3-TTS estimated %d codec tokens for a %d-character utterance, " - "but max_new_tokens is capped at %d; output may still truncate.", - requested_tokens, - len(text), - configured_cap, - ) - - logger.debug( - "Qwen3-TTS using max_new_tokens=%d for utterance with %d words and %d chars", - resolved_tokens, - word_count, - char_count, - ) - return resolved_tokens - - def _warmup_process(self, llm_sentence: str) -> Iterator[bytes | np.ndarray]: - model_type = self._model_type() - if self.ref_audio: - yield from self._process_voice_clone(llm_sentence) - elif model_type == "custom_voice": - yield from self._process_custom_voice(llm_sentence) - elif model_type == "voice_design": - yield from self._process_voice_design(llm_sentence) - else: - raise ValueError( - "Qwen3-TTS Base model requires ref_audio for voice cloning. " - "Provide qwen3_tts_ref_audio or use a CustomVoice/VoiceDesign model." - ) - - def _resample_to_pipeline_sr(self, audio: np.ndarray, sr: int) -> np.ndarray: - if sr == PIPELINE_SR: - return audio - from scipy.signal import resample_poly - - gcd = np.gcd(PIPELINE_SR, sr) - return resample_poly(audio, up=PIPELINE_SR // gcd, down=sr // gcd) - - def _prepare_audio_chunk(self, item: Any) -> tuple[np.ndarray | None, int | None]: - if isinstance(item, tuple): - audio_chunk, sr, _timing = item - return np.asarray(audio_chunk, dtype=np.float32), sr - - audio = getattr(item, "audio", None) - if audio is None: - return None, None - - audio_chunk = np.asarray(audio, dtype=np.float32).squeeze() - sr = getattr(item, "sample_rate", None) or PIPELINE_SR - return audio_chunk, sr - - def _stream(self, gen: Any, label: str) -> Iterator[bytes | np.ndarray]: - """Common streaming loop: log TTFA and RTF, yield int16 chunks.""" - cancel_gen = self.cancel_scope.generation if self.cancel_scope else None - start = perf_counter() - total_samples = 0 - first_chunk = True - found_speech = False - leftover = np.array([], dtype=np.int16) - - for item in gen: - if cancel_gen is not None and self.cancel_scope is not None and self.cancel_scope.is_stale(cancel_gen): - logger.info("TTS generation cancelled (interruption)") - return - - audio_chunk, sr = self._prepare_audio_chunk(item) - if audio_chunk is None or sr is None or audio_chunk.size == 0: - continue - - if first_chunk: - logger.info(f"Qwen3-TTS TTFA: {perf_counter() - start:.2f}s ({label})") - first_chunk = False - - audio_chunk = self._resample_to_pipeline_sr(audio_chunk, sr) - audio_chunk = self._to_int16(audio_chunk) - - # Trim the initial silent ramp-up, but keep enough preroll to avoid - # shaving soft initial phonemes at the start of the utterance. - if not found_speech: - threshold = int(32768 * 0.01) - above = np.abs(audio_chunk) > threshold - if not np.any(above): - continue - start_idx = max(0, int(np.argmax(above)) - int(PIPELINE_SR * 0.040)) - audio_chunk = audio_chunk[start_idx:] - found_speech = True - - audio_chunk = np.concatenate([leftover, audio_chunk]) - - n = (len(audio_chunk) // self.blocksize) * self.blocksize - for i in range(0, n, self.blocksize): - yield audio_chunk[i : i + self.blocksize] - total_samples += self.blocksize - leftover = audio_chunk[n:] - - if len(leftover) > 0: - chunk = np.pad(leftover, (0, self.blocksize - len(leftover))) - yield chunk - total_samples += len(leftover) - - generation_time = perf_counter() - start - audio_duration = total_samples / PIPELINE_SR - rtf = audio_duration / generation_time if generation_time > 0 else 0 - logger.info( - f"Qwen3-TTS generated {audio_duration:.2f}s audio in {generation_time:.2f}s (RTF: {rtf:.2f}, {label})" - ) - - def _coalesce_pending_tts_input(self, current_input: TTSInput) -> tuple[str, Optional[str], bool]: - """Combine already-queued text chunks before the next TTS synthesis call.""" - if not hasattr(self.queue_in, "mutex") or not hasattr(self.queue_in, "queue"): - return current_input.text, current_input.language_code, False - - text = current_input.text - language_code = current_input.language_code - - parts = [text.strip()] if text and text.strip() else [] - saw_end_of_response = False - - with self.queue_in.mutex: - while self.queue_in.queue: - next_item = self.queue_in.queue[0] - if is_control_message(next_item, SESSION_END.kind): - break - if isinstance(next_item, bytes) and next_item == PIPELINE_END: - break - if isinstance(next_item, EndOfResponse): - saw_end_of_response = True - break - if not isinstance(next_item, TTSInput): - break - if current_input.turn_id != next_item.turn_id or current_input.turn_revision != next_item.turn_revision: - break - if ( - language_code is not None - and next_item.language_code is not None - and next_item.language_code != language_code - ): - break - - self.queue_in.queue.popleft() - if next_item.text.strip(): - parts.append(next_item.text.strip()) - if language_code is None: - language_code = next_item.language_code - - combined_text = " ".join(parts).strip() - return combined_text, language_code, saw_end_of_response - - def process(self, tts_input: TTSIn) -> Iterator[TTSOut]: - speculative_turns = getattr(self, "speculative_turns", None) - if isinstance(tts_input, EndOfResponse): - if speculative_turns and not speculative_turns.is_latest_after_reopen_grace( - tts_input.turn_id, - tts_input.turn_revision, - ): - return - yield AUDIO_RESPONSE_DONE - return - - if speculative_turns and not speculative_turns.is_latest_after_reopen_grace( - tts_input.turn_id, - tts_input.turn_revision, - ): - logger.debug("Dropping stale TTS input for turn=%s rev=%s", tts_input.turn_id, tts_input.turn_revision) - return - if speculative_turns: - speculative_turns.commit(tts_input.turn_id, tts_input.turn_revision) - - runtime_config = tts_input.runtime_config - response = tts_input.response - - coalesced_text, _language_code, _saw_end_of_response = self._coalesce_pending_tts_input(tts_input) - - text = coalesced_text or "Hello." - - model_type = self._model_type() - self._apply_session_voice_override(model_type, runtime_config, response) - - console.print(f"[green]ASSISTANT: {text}") - - try: - if self.ref_audio: - audio_iter = self._process_voice_clone(text) - elif model_type == "custom_voice": - audio_iter = self._process_custom_voice(text) - elif model_type == "voice_design": - audio_iter = self._process_voice_design(text) - else: - raise ValueError( - "Qwen3-TTS Base model requires ref_audio for voice cloning. " - "Provide qwen3_tts_ref_audio or use a CustomVoice/VoiceDesign model." - ) - first_audio = True - for audio_chunk in audio_iter: - if first_audio: - self._log_first_audio_latency(tts_input) - first_audio = False - yield audio_chunk - except Exception as e: - logger.error(f"Error during Qwen3-TTS generation: {e}", exc_info=True) - - def _log_first_audio_latency(self, tts_input: TTSInput) -> None: - if tts_input.speech_stopped_at_s is None: - return - latency_s = perf_counter() - tts_input.speech_stopped_at_s - if latency_s < 0: - return - logger.info( - "Last speech detected to first speech out: %.3fs (turn=%s rev=%s)", - latency_s, - tts_input.turn_id, - tts_input.turn_revision, - ) - - def _mlx_streaming_interval(self) -> float: - return max(1, self.streaming_chunk_size) / MLX_STREAMING_TOKENS_PER_SECOND - - def _mlx_stream_kwargs(self, max_tokens: int) -> dict[str, Any]: - return { - "max_tokens": max_tokens, - "verbose": False, - "stream": True, - "streaming_interval": self._mlx_streaming_interval(), - **self.gen_kwargs, - } - - def _stream_mlx_generation( - self, - generation_fn: Callable, - label: str, - max_tokens: int, - **generation_kwargs: Any, - ) -> Iterator[bytes | np.ndarray]: - with MLXLockContext(handler_name="Qwen3TTS", timeout=10.0) as acquired: - if not acquired: - raise TimeoutError("Timed out waiting for MLX lock") - yield from self._stream( - generation_fn( - **self._mlx_stream_kwargs(max_tokens=max_tokens), - **generation_kwargs, - ), - label=label, - ) - - def _process_voice_clone(self, text: str) -> Iterator[bytes | np.ndarray]: - utterance_max_new_tokens = self._estimate_max_new_tokens(text) - if self.backend == "mlx": - if self.xvec_only: - logger.warning("mlx-audio Qwen3-TTS does not support xvec_only; ignoring it") - if self.parity_mode: - logger.info("Qwen3-TTS parity mode is CUDA-specific and is ignored on mlx-audio") - - yield from self._stream_mlx_generation( - self.model.generate, - label="voice_clone_mlx", - max_tokens=utterance_max_new_tokens, - text=text, - ref_audio=self._prepare_mlx_ref_audio(self.ref_audio), - ref_text=self.ref_text, - lang_code=self.language, - ) - return - - yield from self._stream( - self.model.generate_voice_clone_streaming( - text=text, - language=self.language, - ref_audio=self.ref_audio, - ref_text=self.ref_text, - xvec_only=self.xvec_only, - chunk_size=self.streaming_chunk_size, - max_new_tokens=utterance_max_new_tokens, - parity_mode=self.parity_mode, - non_streaming_mode=self.non_streaming_mode, - ), - label="voice_clone_parity" if self.parity_mode else "voice_clone", - ) - - def _process_custom_voice(self, text: str) -> Iterator[bytes | np.ndarray]: - utterance_max_new_tokens = self._estimate_max_new_tokens(text) - speaker = self._resolve_speaker() - if not speaker: - raise ValueError( - "CustomVoice generation requires a speaker. " - "Set qwen3_tts_speaker or use a voice-clone model with ref_audio." - ) - - if self.backend == "mlx": - yield from self._stream_mlx_generation( - self.model.generate_custom_voice, - label="custom_voice_mlx", - max_tokens=utterance_max_new_tokens, - text=text, - speaker=speaker, - language=self.language, - instruct=self.instruct, - ) - return - - yield from self._stream( - self.model.generate_custom_voice_streaming( - text=text, - speaker=speaker, - language=self.language, - instruct=self.instruct, - chunk_size=self.streaming_chunk_size, - max_new_tokens=utterance_max_new_tokens, - non_streaming_mode=self.non_streaming_mode, - ), - label="custom_voice", - ) - - def _process_voice_design(self, text: str) -> Iterator[bytes | np.ndarray]: - utterance_max_new_tokens = self._estimate_max_new_tokens(text) - if self.backend == "mlx": - yield from self._stream_mlx_generation( - self.model.generate_voice_design, - label="voice_design_mlx", - max_tokens=utterance_max_new_tokens, - text=text, - instruct=self.instruct, - language=self.language, - ) - return - - yield from self._stream( - self.model.generate_voice_design_streaming( - text=text, - instruct=self.instruct, - language=self.language, - chunk_size=self.streaming_chunk_size, - max_new_tokens=utterance_max_new_tokens, - non_streaming_mode=self.non_streaming_mode, - ), - label="voice_design", - ) - - def on_session_end(self) -> None: - self.speaker = self._initial_speaker - self.ref_audio = self._initial_ref_audio - logger.debug("Qwen3-TTS session state reset") - - def cleanup(self) -> None: - try: - del self.model - for path in list(getattr(self, "_mlx_temp_ref_audio_files", set())): - try: - Path(path).unlink(missing_ok=True) - except Exception: - pass - if self.backend == "mlx": - try: - import mlx.core as mx - - mx.clear_cache() - except Exception: - pass - else: - import torch - - if torch.cuda.is_available(): - torch.cuda.empty_cache() - logger.info("Qwen3-TTS handler cleaned up") - except Exception as e: - logger.warning(f"Cleanup error: {e}") diff --git a/runtime/src/speech_to_speech/TTS/ref_audio.wav b/runtime/src/speech_to_speech/TTS/ref_audio.wav deleted file mode 100644 index 9b3d864e358b30091890bbffd687c7d7b69d8837..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/TTS/ref_audio.wav +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:c8f3be0d4e32aa3ab4bbc859bb5c2ae176d2170853661edc65c94f232ecca232 -size 159092 diff --git a/runtime/src/speech_to_speech/VAD/__init__.py b/runtime/src/speech_to_speech/VAD/__init__.py deleted file mode 100644 index 8b137891791fe96927ad78e64b0aad7bded08bdc..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/VAD/__init__.py +++ /dev/null @@ -1 +0,0 @@ - diff --git a/runtime/src/speech_to_speech/VAD/vad_handler.py b/runtime/src/speech_to_speech/VAD/vad_handler.py deleted file mode 100644 index f70ce79f0da6565843bf3590ea59173f8b424758..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/VAD/vad_handler.py +++ /dev/null @@ -1,870 +0,0 @@ -from __future__ import annotations - -import logging -import time -from collections.abc import Iterator -from dataclasses import dataclass -from queue import Queue -from threading import Event -from typing import Any, TypeAlias - -import numpy as np -import torch - -from speech_to_speech.api.openai_realtime.runtime_config import RuntimeConfig -from speech_to_speech.baseHandler import BaseHandler -from speech_to_speech.pipeline.events import SpeechStartedEvent, SpeechStoppedEvent -from speech_to_speech.pipeline.handler_types import VADIn, VADOut -from speech_to_speech.pipeline.messages import VADAudio -from speech_to_speech.pipeline.queue_types import TextEventItem -from speech_to_speech.pipeline.speculative_turns import SpeculativeTurnTracker -from speech_to_speech.utils.utils import int2float -from speech_to_speech.VAD.vad_iterator import VADIterator - -logger = logging.getLogger(__name__) - -VADInput: TypeAlias = bytes | tuple[bytes, RuntimeConfig] - - -@dataclass -class _PendingShortSegment: - audio: np.ndarray - active_ms: float - start_ms: int - end_ms: int - - -# Fragments with less active speech than this are treated as noise and never -# held for stitching, so sub-threshold bursts cannot sum past min_speech_ms -# and fire a false barge-in. -_SHORT_SEGMENT_MIN_FRAGMENT_MS = 100 - - -# Optional import for audio enhancement -try: - from df.enhance import enhance, init_df - - HAS_DF = True -except (ImportError, ModuleNotFoundError) as e: - HAS_DF = False - logger.warning(f"DeepFilterNet not available for audio enhancement: {e}") - - -class VADHandler(BaseHandler[VADIn, VADOut]): - """ - Handles voice activity detection. When voice activity is detected, audio will be accumulated until the end of speech is detected and then passed - to the following part. - """ - - def setup( - self, - should_listen: Event, - thresh: float = 0.6, - sample_rate: int = 16000, - min_silence_ms: int = 64, - min_speech_ms: int = 384, - min_speech_continuation_ms: int = 192, - max_speech_ms: float = float("inf"), - speech_pad_ms: int = 30, - audio_enhancement: bool = False, - enable_realtime_transcription: bool = False, - realtime_processing_pause: float = 0.5, - text_output_queue: Queue[TextEventItem] | None = None, - speculative_turns: SpeculativeTurnTracker | None = None, - speculative_reopen_ms: int = 1000, - unanswered_reopen_ms: int = 7000, - short_segment_merge_ms: int = 0, - ) -> None: - self.should_listen = should_listen - self.sample_rate = sample_rate - self.min_silence_ms = min_silence_ms - self.min_speech_ms = min_speech_ms - self.min_speech_continuation_ms = self._resolve_min_speech_continuation_ms( - self.min_speech_ms, - min_speech_continuation_ms, - ) - self.max_speech_ms = max_speech_ms - self.enable_realtime_transcription = enable_realtime_transcription - self.realtime_processing_pause = realtime_processing_pause - self.text_output_queue = text_output_queue - self.speculative_turns = speculative_turns - self.speculative_reopen_ms = speculative_reopen_ms - self.unanswered_reopen_ms = max(self.speculative_reopen_ms, unanswered_reopen_ms) - self.short_segment_merge_ms = max(0, short_segment_merge_ms) - self._last_turn_detection: dict | None = None - self.model, _ = torch.hub.load( - "snakers4/silero-vad", - "silero_vad", - trust_repo=True, - skip_validation=True, - ) - self.iterator = VADIterator( - self.model, - threshold=thresh, - sampling_rate=sample_rate, - min_silence_duration_ms=min_silence_ms, - speech_pad_ms=speech_pad_ms, - ) - self.audio_enhancement = audio_enhancement - if audio_enhancement: - if not HAS_DF: - logger.error( - "Audio enhancement requested but DeepFilterNet is not available. Disabling audio enhancement." - ) - self.audio_enhancement = False - else: - self.enhanced_model, self.df_state, _ = init_df() - - # State for progressive audio release - self.last_process_time: float = 0.0 - - # Cumulative sample counter for audio_start_ms / audio_end_ms - self._total_samples: int = 0 - - # Throttled logging state (summary once per second) - self._last_log_time = 0.0 - self._log_chunks = 0 - self._log_speech_starts = 0 - self._log_speech_ends = 0 - self._log_progressive_yields = 0 - self._speech_started_emitted = False - self._turn_counter = 0 - self._current_turn_id: str | None = None - self._current_turn_revision: int | None = None - self._speculative_audio_prefix: np.ndarray | None = None - self._last_final_wall_time: float | None = None - self._last_final_audio_ms: int | None = None - self._pending_reopen_candidate: tuple[str, int, int] | None = None - self._pending_short_segment: _PendingShortSegment | None = None - - @property - def _audio_ms(self) -> int: - """Cumulative audio received so far, in milliseconds.""" - return int(self._total_samples / self.sample_rate * 1000) - - def _apply_runtime_turn_detection(self, runtime_config: RuntimeConfig | None = None) -> None: - """Check RuntimeConfig for turn_detection changes and apply them.""" - audio = runtime_config.session.audio if runtime_config else None - audio_input = audio.input if audio is not None else None - if not runtime_config or not audio_input or not audio_input.turn_detection: - return - td_raw = audio_input.turn_detection - - # Convert Pydantic models (e.g. OpenAI SDK ServerVad) to dict; - # plain dicts pass through unchanged. - if hasattr(td_raw, "model_dump"): - td = td_raw.model_dump(exclude_none=True) - elif isinstance(td_raw, dict): - td = td_raw - else: - logger.warning(f"Unexpected turn_detection type: {type(td_raw)}") - return - - # Compare normalized snapshot (identity on td_raw vs stored dict was wrong after first apply). - if td == self._last_turn_detection: - return - - self._last_turn_detection = dict(td) - - if "threshold" in td: - self.iterator.threshold = td["threshold"] - logger.info(f"VAD threshold updated to {td['threshold']}") - if "silence_duration_ms" in td: - self.iterator.min_silence_samples = self.sample_rate * td["silence_duration_ms"] / 1000 - logger.info(f"VAD silence duration updated to {td['silence_duration_ms']}ms") - - def _start_new_turn(self) -> tuple[str, int]: - self._cancel_pending_reopen() - self._turn_counter += 1 - self._current_turn_id = f"turn_{self._turn_counter}" - self._current_turn_revision = 0 - self._speculative_audio_prefix = None - self._last_final_wall_time = None - self._last_final_audio_ms = None - if self.speculative_turns: - self.speculative_turns.observe(self._current_turn_id, self._current_turn_revision) - return self._current_turn_id, self._current_turn_revision - - def _speech_buffer_duration_ms(self) -> float: - if not hasattr(self.iterator, "speech_buffer"): - return 0.0 - buffer_samples = sum(len(t) for t in self.iterator.speech_buffer()) - return buffer_samples / self.sample_rate * 1000 - - def _current_active_speech_duration_ms(self) -> float: - active_speech_samples = getattr(self.iterator, "active_speech_samples", 0) - return active_speech_samples / self.sample_rate * 1000 - - def _last_utterance_active_speech_duration_ms(self) -> float: - active_speech_samples = getattr(self.iterator, "last_utterance_active_speech_samples", 0) - return active_speech_samples / self.sample_rate * 1000 - - @staticmethod - def _resolve_min_speech_continuation_ms(min_speech_ms: int, min_speech_continuation_ms: int) -> int: - if min_speech_continuation_ms <= 0: - return min_speech_ms - return min( - min_speech_ms, - max(_SHORT_SEGMENT_MIN_FRAGMENT_MS, min_speech_continuation_ms), - ) - - def _uses_realtime_turn_handling(self) -> bool: - return self.enable_realtime_transcription or self.speculative_turns is not None - - def _active_speech_min_ms(self, start_ms: int) -> float: - """Duration hysteresis for speech that continues a reopenable turn.""" - if self._pending_reopen_candidate is not None or self._should_reopen_current_turn(start_ms): - return self.min_speech_continuation_ms - return self.min_speech_ms - - def _should_reopen_current_turn(self, audio_start_ms: int) -> bool: - if not self._uses_realtime_turn_handling(): - return False - if self._current_turn_id is None or self._current_turn_revision is None or self._last_final_audio_ms is None: - return False - - is_committed = self.speculative_turns is not None and self.speculative_turns.is_committed( - self._current_turn_id, - self._current_turn_revision, - ) - if is_committed: - return False - - # Elapsed is measured on the audio clock, so the window only advances - # while the client streams audio (continuous capture behaves like wall - # time; push-to-talk style gaps freeze it). - elapsed_ms = max(0, audio_start_ms - self._last_final_audio_ms) - - # Within the short grace window, any uncommitted turn may reopen. - # Beyond it, an unanswered turn (no assistant output committed yet) - # remains reopenable up to the unanswered_reopen_ms sanity cap, so a - # user pause longer than speculative_reopen_ms does not orphan a turn - # the assistant has not replied to. The cap also bounds the - # empty-transcript case, where no request is queued and the turn would - # otherwise never commit. - reopen_limit_ms = self.speculative_reopen_ms - if self.speculative_turns is not None: - reopen_limit_ms = self.unanswered_reopen_ms - return elapsed_ms <= reopen_limit_ms - - def _begin_pending_reopen_if_needed(self, audio_start_ms: int) -> None: - if self._pending_reopen_candidate is not None or not self._should_reopen_current_turn(audio_start_ms): - return - if self.speculative_turns is None: - return - candidate_revision = self.speculative_turns.begin_reopen_candidate( - self._current_turn_id, - self._current_turn_revision, - ) - if candidate_revision is None or self._current_turn_id is None or self._current_turn_revision is None: - return - self._pending_reopen_candidate = ( - self._current_turn_id, - self._current_turn_revision, - candidate_revision, - ) - logger.info( - "VAD: pending reopen candidate for speculative turn %s revision %d", - self._current_turn_id, - candidate_revision, - ) - - def _cancel_pending_reopen(self) -> None: - if self._pending_reopen_candidate is None: - return - turn_id, _base_revision, candidate_revision = self._pending_reopen_candidate - if self.speculative_turns: - self.speculative_turns.cancel_reopen_candidate(turn_id, candidate_revision) - self._pending_reopen_candidate = None - - def _confirm_pending_reopen(self) -> tuple[str, int, bool] | None: - if self._pending_reopen_candidate is None: - return None - turn_id, base_revision, candidate_revision = self._pending_reopen_candidate - self._pending_reopen_candidate = None - if self.speculative_turns and not self.speculative_turns.confirm_reopen_candidate( - turn_id, - base_revision, - candidate_revision, - ): - return None - self._current_turn_id = turn_id - self._current_turn_revision = candidate_revision - logger.info("VAD: reopened speculative turn %s revision %d", turn_id, candidate_revision) - return turn_id, candidate_revision, True - - def _reopen_current_turn(self) -> tuple[str, int, bool] | None: - if self._current_turn_id is None or self._current_turn_revision is None: - return None - - turn_id = self._current_turn_id - base_revision = self._current_turn_revision - if self.speculative_turns is not None: - candidate_revision = self.speculative_turns.begin_reopen_candidate(turn_id, base_revision) - if candidate_revision is None or not self.speculative_turns.confirm_reopen_candidate( - turn_id, - base_revision, - candidate_revision, - ): - return None - else: - candidate_revision = base_revision + 1 - - self._current_turn_id = turn_id - self._current_turn_revision = candidate_revision - logger.info("VAD: reopened speculative turn %s revision %d", turn_id, candidate_revision) - return turn_id, candidate_revision, True - - def _ensure_turn_for_speech_start(self, audio_start_ms: int) -> tuple[str, int, bool]: - if ( - self._speech_started_emitted - and self._current_turn_id is not None - and self._current_turn_revision is not None - ): - return self._current_turn_id, self._current_turn_revision, False - - confirmed_reopen = self._confirm_pending_reopen() - if confirmed_reopen is not None: - return confirmed_reopen - - reopened = False - if self._should_reopen_current_turn(audio_start_ms): - reopened_turn = self._reopen_current_turn() - if reopened_turn is not None: - return reopened_turn - - self._start_new_turn() - - if self._current_turn_id is None or self._current_turn_revision is None: - raise RuntimeError("VAD failed to allocate turn metadata") - if self.speculative_turns: - self.speculative_turns.observe(self._current_turn_id, self._current_turn_revision) - return self._current_turn_id, self._current_turn_revision, reopened - - def _current_turn_metadata(self) -> tuple[str | None, int | None]: - return self._current_turn_id, self._current_turn_revision - - def _combined_turn_audio(self, current_segment: np.ndarray) -> np.ndarray: - if self._speculative_audio_prefix is None: - return current_segment - return np.concatenate((self._speculative_audio_prefix, current_segment)) - - def _short_segment_merge_window_ms(self) -> int: - return int(getattr(self, "short_segment_merge_ms", 0)) - - def _segment_duration_ms(self, segment: np.ndarray) -> float: - return len(segment) / self.sample_rate * 1000 - - def _segment_start_ms(self, segment: np.ndarray, end_ms: int) -> int: - return max(0, end_ms - int(self._segment_duration_ms(segment))) - - def _short_segment_gap_ms(self, start_ms: int) -> float: - if self._pending_short_segment is None: - return float("inf") - return max(0, start_ms - self._pending_short_segment.end_ms) - - def _can_merge_pending_short_segment(self, start_ms: int) -> bool: - return ( - self._pending_short_segment is not None - and self._short_segment_merge_window_ms() > 0 - and self._short_segment_gap_ms(start_ms) <= self._short_segment_merge_window_ms() - ) - - def _effective_active_speech_for_start(self, start_ms: int, active_ms: float) -> tuple[int, float]: - # A live fragment below the noise floor never counts the held segment - # toward the speech-start threshold, mirroring the finalization path. - if active_ms < _SHORT_SEGMENT_MIN_FRAGMENT_MS: - return start_ms, active_ms - if not self._can_merge_pending_short_segment(start_ms): - return start_ms, active_ms - assert self._pending_short_segment is not None - return self._pending_short_segment.start_ms, self._pending_short_segment.active_ms + active_ms - - def _merge_pending_short_segment( - self, - segment: np.ndarray, - active_ms: float, - end_ms: int, - ) -> tuple[np.ndarray, float, int, bool]: - start_ms = self._segment_start_ms(segment, end_ms) - if not self._can_merge_pending_short_segment(start_ms): - self._discard_expired_pending_short_segment(start_ms) - return segment, active_ms, start_ms, False - - pending = self._pending_short_segment - assert pending is not None - # Reinsert the silence between the two segments so the stitched audio - # keeps its acoustic gap and its length matches the audio-clock span. - gap_samples = int(self._short_segment_gap_ms(start_ms) * self.sample_rate / 1000) - self._pending_short_segment = None - parts = [pending.audio] - if gap_samples > 0: - parts.append(np.zeros(gap_samples, dtype=segment.dtype)) - parts.append(segment) - merged = np.concatenate(parts) - return merged, pending.active_ms + active_ms, pending.start_ms, True - - def _hold_short_segment(self, segment: np.ndarray, active_ms: float, start_ms: int, end_ms: int) -> None: - self._pending_short_segment = _PendingShortSegment( - audio=segment, - active_ms=active_ms, - start_ms=start_ms, - end_ms=end_ms, - ) - logger.info( - "VAD: holding short segment=%.0fms active=%.0fms (active_min=%sms, merge_max=%sms)", - self._segment_duration_ms(segment), - active_ms, - self.min_speech_ms, - self._short_segment_merge_window_ms(), - ) - - def _discard_pending_short_segment(self, reason: str = "expired") -> None: - pending = self._pending_short_segment - if pending is None: - return - self._pending_short_segment = None - logger.info( - "VAD: discarding held short segment=%.0fms active=%.0fms (%s, active_min=%sms)", - self._segment_duration_ms(pending.audio), - pending.active_ms, - reason, - self.min_speech_ms, - ) - - def _discard_expired_pending_short_segment(self, next_start_ms: int | None = None) -> None: - pending = self._pending_short_segment - if pending is None or self._short_segment_merge_window_ms() <= 0: - return - reference_ms = self._audio_ms if next_start_ms is None else next_start_ms - gap_ms = max(0, reference_ms - pending.end_ms) - if gap_ms > self._short_segment_merge_window_ms(): - self._discard_pending_short_segment("merge window elapsed") - - def before_emit_output(self, output: VADOut) -> None: - if isinstance(output, VADAudio): - self._drop_superseded_vad_audio(output) - - def _drop_superseded_vad_audio(self, latest: VADAudio) -> int: - if not hasattr(self.queue_out, "mutex") or not hasattr(self.queue_out, "queue"): - return 0 - - dropped = 0 - with self.queue_out.mutex: - kept: list[Any] = [] - while self.queue_out.queue: - queued_item = self.queue_out.queue.popleft() - if isinstance(queued_item, VADAudio) and self._vad_audio_is_superseded( - queued_item, - latest, - ): - dropped += 1 - else: - kept.append(queued_item) - self.queue_out.queue.extend(kept) - if dropped: - self.queue_out.not_full.notify_all() - - if dropped: - logger.debug( - "VAD: dropped %d superseded audio chunk(s) before enqueueing turn=%s rev=%s mode=%s", - dropped, - latest.turn_id, - latest.turn_revision, - latest.mode, - ) - return dropped - - def _vad_audio_is_superseded(self, queued_item: VADAudio, latest: VADAudio) -> bool: - if queued_item.turn_id is None or queued_item.turn_revision is None: - return False - if self.speculative_turns is not None and not self.speculative_turns.is_latest( - queued_item.turn_id, - queued_item.turn_revision, - ): - return True - return ( - queued_item.mode == "progressive" - and queued_item.turn_id == latest.turn_id - and queued_item.turn_revision == latest.turn_revision - ) - - def process(self, audio_chunk: VADIn) -> Iterator[VADOut]: - runtime_config = None - if isinstance(audio_chunk, tuple): - audio_chunk, runtime_config = audio_chunk - self._apply_runtime_turn_detection(runtime_config) - - if not self.should_listen.is_set(): - return - - # Normal listening mode - self._log_chunks += 1 - audio_int16 = np.frombuffer(audio_chunk, dtype=np.int16) - self._total_samples += len(audio_int16) - audio_float32 = int2float(audio_int16) - - vad_output = self.iterator(torch.from_numpy(audio_float32)) - - # Deferred speech_started: only emit once active VAD speech reaches the valid speech threshold. - is_triggered_now = self.iterator.triggered - if is_triggered_now and not self._speech_started_emitted: - segment_samples = sum(len(t) for t in self.iterator.buffer) - segment_duration_ms = segment_samples / self.sample_rate * 1000 - active_speech_duration_ms = self._current_active_speech_duration_ms() - speech_buffer_duration_ms = self._speech_buffer_duration_ms() - start_ms = max(0, self._audio_ms - int(speech_buffer_duration_ms)) - effective_start_ms, effective_active_speech_duration_ms = self._effective_active_speech_for_start( - start_ms, - active_speech_duration_ms, - ) - self._begin_pending_reopen_if_needed(effective_start_ms) - active_speech_min_ms = self._active_speech_min_ms(effective_start_ms) - if effective_active_speech_duration_ms >= active_speech_min_ms: - turn_id, turn_revision, reopened = self._ensure_turn_for_speech_start(effective_start_ms) - self._speech_started_emitted = True - self._log_speech_starts += 1 - logger.info( - "Speech started (confirmed, active=%.0fms, min=%.0fms, segment=%.0fms, turn=%s rev=%s)", - effective_active_speech_duration_ms, - active_speech_min_ms, - segment_duration_ms, - turn_id, - turn_revision, - ) - if self.text_output_queue: - self.text_output_queue.put( - SpeechStartedEvent( - audio_start_ms=effective_start_ms, - turn_id=turn_id, - turn_revision=turn_revision, - reopened=reopened, - ) - ) - elif not is_triggered_now and vad_output is None: - self._discard_expired_pending_short_segment() - - # Log a summary once per second instead of every chunk - now = time.time() - if now - self._last_log_time >= 1.0: - state = "SPEAKING" if is_triggered_now else "silent" - logger.debug( - f"VAD: {self._log_chunks} chunks/s | {state} | " - f"starts={self._log_speech_starts} ends={self._log_speech_ends} progressive={self._log_progressive_yields}" - ) - self._log_chunks = 0 - self._log_speech_starts = 0 - self._log_speech_ends = 0 - self._log_progressive_yields = 0 - self._last_log_time = now - - if self._uses_realtime_turn_handling(): - # Realtime mode keeps turns reopenable; live transcription additionally - # emits progressive audio chunks while speaking. - yield from self._process_realtime(vad_output, runtime_config) - else: - # Original mode: yield only when speech ends - yield from self._process_normal(vad_output, runtime_config) - - def _process_realtime( - self, - vad_output: list[torch.Tensor] | None, - runtime_config: RuntimeConfig | None = None, - ) -> Iterator[VADOut]: - """Process with real-time progressive audio release.""" - # Check if we're currently in a speech segment. - if self.enable_realtime_transcription and hasattr(self.iterator, "buffer") and len(self.iterator.buffer) > 0: - current_time = time.time() - duration_ms = self._speech_buffer_duration_ms() - progressive_pause = self._progressive_processing_pause(duration_ms) - - # Yield accumulated audio periodically while speaking - if (current_time - self.last_process_time) >= progressive_pause: - array = torch.cat(self.iterator.speech_buffer()).cpu().numpy() - duration_ms = len(array) / self.sample_rate * 1000 - active_speech_duration_ms = self._current_active_speech_duration_ms() - start_ms = max(0, self._audio_ms - int(duration_ms)) - - if active_speech_duration_ms >= self._active_speech_min_ms(start_ms): - self._log_progressive_yields += 1 - logger.debug( - "VAD: yielding progressive audio (segment=%.0fms, active=%.0fms, interval=%.2fs)", - duration_ms, - active_speech_duration_ms, - progressive_pause, - ) - turn_id, turn_revision = self._current_turn_metadata() - yield VADAudio( - audio=self._combined_turn_audio(array), - runtime_config=runtime_config, - mode="progressive", - turn_id=turn_id, - turn_revision=turn_revision, - ) - self.last_process_time = current_time - - # Handle end of speech - if vad_output is not None: - if len(vad_output) == 0: - logger.info("VAD: phantom trigger (empty buffer), closing speech pair") - if self._speech_started_emitted and self.text_output_queue: - turn_id, turn_revision = self._current_turn_metadata() - self.text_output_queue.put( - SpeechStoppedEvent( - audio_end_ms=self._audio_ms, - turn_id=turn_id, - turn_revision=turn_revision, - ) - ) - if not self._speech_started_emitted: - self._cancel_pending_reopen() - self._speech_started_emitted = False - self._discard_expired_pending_short_segment() - return - - array = torch.cat(vad_output).cpu().numpy() - end_ms = self._audio_ms - raw_active_ms = self._last_utterance_active_speech_duration_ms() - active_speech_duration_ms = raw_active_ms - stitched_short_segment = False - # Fragments below the noise floor never merge with or replace a - # held segment; the pending segment's own expiry handles it. - if raw_active_ms >= _SHORT_SEGMENT_MIN_FRAGMENT_MS: - array, active_speech_duration_ms, start_ms, stitched_short_segment = self._merge_pending_short_segment( - array, - active_speech_duration_ms, - end_ms, - ) - else: - start_ms = self._segment_start_ms(array, end_ms) - duration_ms = self._segment_duration_ms(array) - min_active_ms = 0.0 if self._speech_started_emitted else self._active_speech_min_ms(start_ms) - - if active_speech_duration_ms < min_active_ms or duration_ms > self.max_speech_ms: - if ( - self._short_segment_merge_window_ms() > 0 - and raw_active_ms >= _SHORT_SEGMENT_MIN_FRAGMENT_MS - and active_speech_duration_ms < min_active_ms - and duration_ms <= self.max_speech_ms - ): - self._hold_short_segment(array, active_speech_duration_ms, start_ms, end_ms) - else: - logger.info( - "VAD: discarding segment=%.0fms active=%.0fms (active_min=%sms, segment_max=%sms)", - duration_ms, - active_speech_duration_ms, - min_active_ms, - self.max_speech_ms, - ) - if self._speech_started_emitted and self.text_output_queue: - turn_id, turn_revision = self._current_turn_metadata() - self.text_output_queue.put( - SpeechStoppedEvent( - audio_end_ms=self._audio_ms, - turn_id=turn_id, - turn_revision=turn_revision, - ) - ) - if not self._speech_started_emitted: - self._cancel_pending_reopen() - self._speech_started_emitted = False - else: - if stitched_short_segment: - logger.info( - "VAD: stitched short segment(s) into segment=%.0fms active=%.0fms", - duration_ms, - active_speech_duration_ms, - ) - if not self._speech_started_emitted: - turn_id, turn_revision, reopened = self._ensure_turn_for_speech_start(start_ms) - if self.text_output_queue: - self.text_output_queue.put( - SpeechStartedEvent( - audio_start_ms=start_ms, - turn_id=turn_id, - turn_revision=turn_revision, - reopened=reopened, - interrupt_response=False, - ) - ) - else: - turn_id, turn_revision = self._current_turn_metadata() - self._log_speech_ends += 1 - logger.info( - "Speech soft-ended (segment=%.0fms, active=%.0fms, turn=%s rev=%s)", - duration_ms, - active_speech_duration_ms, - turn_id, - turn_revision, - ) - if self.audio_enhancement: - array = self._apply_audio_enhancement(array) - output_array = self._combined_turn_audio(array) - combined_duration_s = len(output_array) / self.sample_rate - if self.text_output_queue: - self.text_output_queue.put( - SpeechStoppedEvent( - duration_s=combined_duration_s, - audio_end_ms=end_ms, - turn_id=turn_id, - turn_revision=turn_revision, - ) - ) - self._speculative_audio_prefix = output_array - self._last_final_wall_time = time.time() - self._last_final_audio_ms = end_ms - if self.speculative_turns: - # The grace window only delays response commits; reopen - # eligibility for unanswered turns is extended separately - # via unanswered_reopen_ms in _should_reopen_current_turn. - self.speculative_turns.start_reopen_grace( - turn_id, - turn_revision, - self.speculative_reopen_ms / 1000.0, - ) - else: - self.should_listen.clear() - yield VADAudio( - audio=output_array, - runtime_config=runtime_config, - mode="final", - turn_id=turn_id, - turn_revision=turn_revision, - ) - self.last_process_time = 0.0 - self._speech_started_emitted = False - - def _progressive_processing_pause(self, duration_ms: float) -> float: - base_pause = max(0.0, self.realtime_processing_pause) - duration_s = duration_ms / 1000.0 - if duration_s < 8.0: - multiplier = 1.0 - elif duration_s < 15.0: - multiplier = 2.0 - elif duration_s < 30.0: - multiplier = 4.0 - else: - multiplier = 6.0 - return min(base_pause * multiplier, 2.0) - - def _process_normal( - self, - vad_output: list[torch.Tensor] | None, - runtime_config: RuntimeConfig | None = None, - ) -> Iterator[VADOut]: - """Original processing: yield only when speech ends.""" - if vad_output is not None: - if len(vad_output) == 0: - logger.info("VAD: phantom trigger (empty buffer), closing speech pair") - if self._speech_started_emitted and self.text_output_queue: - self.text_output_queue.put(SpeechStoppedEvent(audio_end_ms=self._audio_ms)) - self._speech_started_emitted = False - self._discard_expired_pending_short_segment() - return - - array = torch.cat(vad_output).cpu().numpy() - end_ms = self._audio_ms - raw_active_ms = self._last_utterance_active_speech_duration_ms() - active_speech_duration_ms = raw_active_ms - stitched_short_segment = False - # Fragments below the noise floor never merge with or replace a - # held segment; the pending segment's own expiry handles it. - if raw_active_ms >= _SHORT_SEGMENT_MIN_FRAGMENT_MS: - array, active_speech_duration_ms, start_ms, stitched_short_segment = self._merge_pending_short_segment( - array, - active_speech_duration_ms, - end_ms, - ) - else: - start_ms = self._segment_start_ms(array, end_ms) - duration_ms = self._segment_duration_ms(array) - min_active_ms = 0.0 if self._speech_started_emitted else self._active_speech_min_ms(start_ms) - if active_speech_duration_ms < min_active_ms or duration_ms > self.max_speech_ms: - if ( - self._short_segment_merge_window_ms() > 0 - and raw_active_ms >= _SHORT_SEGMENT_MIN_FRAGMENT_MS - and active_speech_duration_ms < min_active_ms - and duration_ms <= self.max_speech_ms - ): - self._hold_short_segment(array, active_speech_duration_ms, start_ms, end_ms) - else: - logger.info( - "VAD: discarding segment=%.0fms active=%.0fms (active_min=%sms, segment_max=%sms)", - duration_ms, - active_speech_duration_ms, - min_active_ms, - self.max_speech_ms, - ) - if self._speech_started_emitted and self.text_output_queue: - self.text_output_queue.put(SpeechStoppedEvent(audio_end_ms=self._audio_ms)) - self._speech_started_emitted = False - else: - if stitched_short_segment: - logger.info( - "VAD: stitched short segment(s) into segment=%.0fms active=%.0fms", - duration_ms, - active_speech_duration_ms, - ) - if not self._speech_started_emitted and self.text_output_queue: - self.text_output_queue.put( - SpeechStartedEvent( - audio_start_ms=start_ms, - interrupt_response=False, - ) - ) - self._log_speech_ends += 1 - self.should_listen.clear() - logger.info(f"Speech ended ({duration_ms:.0f}ms), stop listening") - if self.text_output_queue: - self.text_output_queue.put(SpeechStoppedEvent(duration_s=duration_ms / 1000.0, audio_end_ms=end_ms)) - if self.audio_enhancement: - array = self._apply_audio_enhancement(array) - yield VADAudio(audio=array, runtime_config=runtime_config) - self._speech_started_emitted = False - - def _apply_audio_enhancement(self, array: np.ndarray) -> np.ndarray: - """Apply audio enhancement if enabled.""" - import torchaudio - - if self.sample_rate != self.df_state.sr(): - audio_float32 = torchaudio.functional.resample( - torch.from_numpy(array), - orig_freq=self.sample_rate, - new_freq=self.df_state.sr(), - ) - enhanced = enhance( - self.enhanced_model, - self.df_state, - audio_float32.unsqueeze(0), - ) - enhanced = torchaudio.functional.resample( - enhanced, - orig_freq=self.df_state.sr(), - new_freq=self.sample_rate, - ) - else: - enhanced = enhance(self.enhanced_model, self.df_state, torch.from_numpy(array)) - return enhanced.numpy().squeeze() - - def on_session_end(self): - self.iterator.reset_states() - self._pending_short_segment = None - self.iterator.buffer = [] - self.last_process_time = 0.0 - self._total_samples = 0 - self._speech_started_emitted = False - self._turn_counter = 0 - self._current_turn_id = None - self._current_turn_revision = None - self._speculative_audio_prefix = None - self._last_final_wall_time = None - self._last_final_audio_ms = None - self._pending_reopen_candidate = None - if self.speculative_turns: - self.speculative_turns.reset() - self.should_listen.set() - logger.debug("VAD session state reset") - - @property - def min_time_to_debug(self) -> float: - return 0.00001 diff --git a/runtime/src/speech_to_speech/VAD/vad_iterator.py b/runtime/src/speech_to_speech/VAD/vad_iterator.py deleted file mode 100644 index 6ac58fa42ec978193b4731c06c724ea3aa0b7efb..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/VAD/vad_iterator.py +++ /dev/null @@ -1,170 +0,0 @@ -from collections import deque - -import torch - - -class VADIterator: - def __init__( - self, - model, - threshold: float = 0.5, - sampling_rate: int = 16000, - min_silence_duration_ms: int = 300, - speech_pad_ms: int = 30, - ) -> None: - """ - Mainly taken from https://github.com/snakers4/silero-vad - Class for stream imitation - - Parameters - ---------- - model: preloaded .jit/.onnx silero VAD model - - threshold: float (default - 0.5) - Speech threshold. Silero VAD outputs speech probabilities for each audio chunk, probabilities ABOVE this value are considered as SPEECH. - It is better to tune this parameter for each dataset separately, but "lazy" 0.5 is pretty good for most datasets. - - sampling_rate: int (default - 16000) - Currently silero VAD models support 8000 and 16000 sample rates - - min_silence_duration_ms: int (default - 300 milliseconds) - In the end of each speech chunk wait for min_silence_duration_ms before separating it - - speech_pad_ms: int (default - 30 milliseconds) - Retain up to speech_pad_ms of audio before VAD triggers and prepend it - to the detected speech chunk - """ - - self.model = model - self.threshold = threshold - self.sampling_rate = sampling_rate - self.is_speaking = False - self.buffer: list[torch.Tensor] = [] - self.prefix_buffer: list[torch.Tensor] = [] - self.active_speech_samples = 0 - self.last_utterance_active_speech_samples = 0 - self._pre_speech_buffer: deque[torch.Tensor] = deque() - self._pre_speech_samples = 0 - - if sampling_rate not in [8000, 16000]: - raise ValueError("VADIterator does not support sampling rates other than [8000, 16000]") - - self.min_silence_samples = int(sampling_rate * min_silence_duration_ms / 1000) - self.speech_pad_samples = int(sampling_rate * speech_pad_ms / 1000) - self.reset_states() - - def reset_states(self) -> None: - self.model.reset_states() - self.triggered = False - self.temp_end = 0 - self.current_sample = 0 - self.buffer = [] - self.prefix_buffer = [] - self.active_speech_samples = 0 - self.last_utterance_active_speech_samples = 0 - self._pre_speech_buffer.clear() - self._pre_speech_samples = 0 - - def _num_samples(self, chunk: torch.Tensor) -> int: - return len(chunk[0]) if chunk.dim() == 2 else len(chunk) - - def _trim_pre_speech_buffer(self) -> None: - while ( - self.speech_pad_samples > 0 - and self._pre_speech_buffer - and self._pre_speech_samples > self.speech_pad_samples - ): - first = self._pre_speech_buffer[0] - first_samples = self._num_samples(first) - excess = self._pre_speech_samples - self.speech_pad_samples - - if excess >= first_samples: - self._pre_speech_buffer.popleft() - self._pre_speech_samples -= first_samples - continue - - if first.dim() == 2: - self._pre_speech_buffer[0] = first[:, excess:] - else: - self._pre_speech_buffer[0] = first[excess:] - self._pre_speech_samples -= excess - - def _remember_pre_speech(self, chunk: torch.Tensor) -> None: - if self.speech_pad_samples <= 0: - self._pre_speech_buffer.clear() - self._pre_speech_samples = 0 - return - - self._pre_speech_buffer.append(chunk) - self._pre_speech_samples += self._num_samples(chunk) - self._trim_pre_speech_buffer() - - def _speech_buffer(self) -> list[torch.Tensor]: - if not self.prefix_buffer: - return list(self.buffer) - return [*self.prefix_buffer, *self.buffer] - - def speech_buffer(self) -> list[torch.Tensor]: - return self._speech_buffer() - - @torch.no_grad() - def __call__(self, x: torch.Tensor) -> list[torch.Tensor] | None: - """ - x: torch.Tensor - audio chunk (see examples in repo) - - return_seconds: bool (default - False) - whether return timestamps in seconds (default - samples) - """ - - if not torch.is_tensor(x): - try: - x = torch.Tensor(x) - except Exception: - raise TypeError("Audio cannot be casted to tensor. Cast it manually") - - window_size_samples = len(x[0]) if x.dim() == 2 else len(x) - self.current_sample += window_size_samples - - speech_prob = self.model(x, self.sampling_rate).item() - - if (speech_prob >= self.threshold) and not self.triggered: - self.triggered = True - self.prefix_buffer = list(self._pre_speech_buffer) - self._pre_speech_buffer.clear() - self._pre_speech_samples = 0 - self.buffer.append(x) - self.active_speech_samples = window_size_samples - self.last_utterance_active_speech_samples = 0 - return None - - if not self.triggered: - self._remember_pre_speech(x) - return None - - if self.triggered: - self.buffer.append(x) - if speech_prob >= self.threshold - 0.15: - self.active_speech_samples += window_size_samples - if self.temp_end: - self.temp_end = 0 - return None - - if speech_prob < self.threshold - 0.15: - if not self.temp_end: - self.temp_end = self.current_sample - if self.current_sample - self.temp_end < self.min_silence_samples: - return None - - # End of speech: keep the final low-confidence chunks that were - # observed before VAD decided the utterance was done. - self.temp_end = 0 - self.triggered = False - spoken_utterance = self.speech_buffer() - self.last_utterance_active_speech_samples = self.active_speech_samples - self.active_speech_samples = 0 - self.buffer = [] - self.prefix_buffer = [] - return spoken_utterance - - return None diff --git a/runtime/src/speech_to_speech/__init__.py b/runtime/src/speech_to_speech/__init__.py deleted file mode 100644 index e6a420113e2c1044ea63a89be661c2b357644bc9..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -"""speech-to-speech: low-latency speech-to-speech pipeline.""" - -__version__ = "0.2.11" diff --git a/runtime/src/speech_to_speech/api/__init__.py b/runtime/src/speech_to_speech/api/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/runtime/src/speech_to_speech/api/openai_realtime/README.md b/runtime/src/speech_to_speech/api/openai_realtime/README.md deleted file mode 100644 index 35ba141a25584550a00da5ddaefc520a778fca86..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/api/openai_realtime/README.md +++ /dev/null @@ -1,244 +0,0 @@ -## Realtime Engine -- High-Level Architecture - -The server runs as a handler thread inside the existing pipeline. A FastAPI/uvicorn instance serves the WebSocket endpoint at `/v1/realtime`, while the same queue-driven handler chain (VAD, STT, LLM, TTS) processes audio in parallel threads. - -```mermaid -flowchart LR - subgraph client [Client] - WS["WebSocket /v1/realtime"] - end - - subgraph server [Realtime Server] - Router["WebSocket Router"] - Service["RealtimeService"] - Config["RuntimeConfig"] - end - - subgraph pipeline [Pipeline Threads] - VAD["VAD"] - STT["STT"] - TN["TranscriptionNotifier"] - LLM["LLM"] - Proc["LMOutputProcessor"] - TTS["TTS"] - end - - WS -- "client events (JSON)" --> Router - Router -- "parse / dispatch" --> Service - Service -- "PCM chunks" --> VAD - Service -- "session config" --> Config - Config -. "read by" .-> VAD & LLM & TTS - VAD -- "speech segments" --> STT - STT -- "transcript" --> TN - TN -- "transcript + notify" --> LLM - LLM -- "text + tools" --> Proc - Proc -- "clean text" --> TTS - TTS -- "PCM audio" --> Router - Proc -- "assistant_text / tools" --> Router - VAD -- "speech_started/stopped" --> Router - TN -- "transcription events" --> Router - Router -- "server events (JSON)" --> WS -``` - -**Key flow:** - -1. **Inbound audio**: Client sends `input_audio_buffer.append` with base64 PCM. `RealtimeService` decodes, resamples to 16 kHz, splits into 512-sample chunks, and puts them on the `recv_audio_chunks_queue` for VAD. -2. **Speech detection**: VAD detects speech boundaries and emits `speech_started` / `speech_stopped` events on the `text_output_queue`. Full utterance audio goes to STT. -3. **Transcription**: STT output passes through `TranscriptionNotifier`, which taps the transcript for `transcription.delta` / `transcription.completed` events before forwarding to the LLM. -4. **Generation**: The LLM generates text (and optional tool calls). `LMOutputProcessor` splits the output: clean text goes to TTS, and `assistant_text` + tool call dicts go to the `text_output_queue`. -5. **Outbound audio**: TTS writes PCM chunks to `send_audio_chunks_queue`. The router's async `_send_loop` drains both queues, encoding PCM as `response.output_audio.delta` events and translating internal messages into protocol events. -6. **Session config**: `session.update` events deep-merge into `RuntimeConfig`, which is a shared Pydantic model read by VAD (turn detection thresholds), LLM (instructions, tools), and TTS (voice) at processing time. - ---- - -## Supported OpenAI Realtime Events - -### Client -> Server - -| Event | Description | -|---|---| -| `input_audio_buffer.append` | Stream base64 PCM audio. Decoded, resampled to 16 kHz, and chunked for the VAD. | -| `session.update` | Deep-merge session config (instructions, tools, voice, turn detection, audio format). | -| `conversation.item.create` | Inject `input_text` or `function_call_output` into the LLM context without triggering generation. | -| `response.create` | Trigger LLM generation. Supports per-response `instructions` and `tool_choice` overrides. | -| `response.cancel` | Cancel the in-progress response and re-enable listening. | - -### Server -> Client - -| Event | Description | -|---|---| -| `session.created` | Sent on connection with current session config. | -| `error` | Protocol errors (`session_limit_reached`, `unknown_or_invalid_event`, `invalid_session_type`, `conversation_already_has_active_response`, etc.) | -| `input_audio_buffer.speech_started` | VAD detected user speech. | -| `input_audio_buffer.speech_stopped` | End of user speech segment. | -| `conversation.item.created` | Acknowledges injected `input_text` from `conversation.item.create`. | -| `conversation.item.input_audio_transcription.delta` | Streaming partial transcript (when live transcription is enabled). | -| `conversation.item.input_audio_transcription.completed` | Final transcript for the user turn (with duration usage). | -| `response.created` | Emitted on the first outbound audio chunk (response is `in_progress`). | -| `response.output_audio.delta` | Base64 PCM audio chunk from TTS. | -| `response.output_audio.done` | Audio stream complete for the current output item. | -| `response.output_audio_transcript.done` | Full assistant text transcript for the turn. | -| `response.function_call_arguments.done` | Tool call with `call_id`, `name`, and JSON `arguments`. | -| `response.done` | Response finished (`completed`, `cancelled` with reason `turn_detected` or `client_cancelled`). | - ---- - -## WebRTC Transport - -Alongside the WebSocket endpoint, the server supports the OpenAI GA WebRTC handshake (requires the `webrtc` extra: `pip install 'speech-to-speech[webrtc]'`): - -``` -POST /v1/realtime/calls Content-Type: application/sdp -``` - -The client POSTs an SDP offer and receives an SDP answer (`201`, with a `Location: /v1/realtime/calls/{call_id}` header). Audio then flows over RTP media tracks (Opus at 48 kHz, resampled to and from the 16 kHz pipeline rate with a stateful resampler), while all JSON events use the same protocol as WebSocket mode, carried on the `oai-events` data channel. - -A WebRTC session claims a pipeline unit from the same pool as WebSocket clients; the per-unit send loop stays the sole consumer of the pipeline output queues and hands PCM to the transport, which paces it out as 20 ms RTP frames (silence when idle). - -Differences from WebSocket mode: - -- `input_audio_buffer.append` is rejected with `invalid_event_for_transport` — audio arrives on the media track. -- `output_audio_buffer.clear` is supported (WebRTC only): unplayed audio buffers server-side, so barge-in and cancellation flush it there. The server also flushes it automatically on `response.cancel` and VAD interruption. -- `session.created` is sent when the data channel opens rather than on connection. - -ICE servers (STUN/TURN) can be configured via the `SPEECH_TO_SPEECH_ICE_SERVERS` env var, a JSON list of server entries: - -```bash -export SPEECH_TO_SPEECH_ICE_SERVERS='[{"urls": "stun:stun.example.com:3478"}, {"urls": "turn:turn.example.com", "username": "u", "credential": "c"}]' -``` - -Without it, aiortc defaults apply (host candidates + Google STUN). Deployments where clients cannot reach the server directly (symmetric NAT, containers without exposed UDP) need a TURN server. - ---- - -## Tool Calling Design - -Tool calling works through two distinct paths depending on the LLM backend, but both converge to the same wire protocol for the client. - -### Local LLM path (`LanguageModelHandler` -- transformers / mlx-lm) - -Tools defined in `session.update` are converted to `FunctionTool` objects. Each tool's JSON Schema `parameters` are turned into a Python `inspect.Signature` (via `signature_from_schema`), and `to_code_prompt()` renders a human-readable `def name(...): """docstring"""` block. - -These tool signatures are injected into the system prompt using a Jinja2 template (`tool_prompt.py`) that instructs the model to wrap tool calls in `...` delimiters: - -``` - -function_name(arg_name_1=value1, arg_name_2='string_value') - -``` - -After generation, `_extract_tools` uses a regex to find `` blocks, then `extract_function_calls_from_text` parses each `name(kwargs)` call and validates it against registered tools. Valid calls become `ResponseFunctionToolCall` dicts with generated `call_id`s. - -### OpenAI API path (`ResponsesApiModelHandler`) - -Tools are passed natively as the `tools=` parameter to `client.responses.create`. The API returns structured `function_call` items directly -- no prompt engineering or regex parsing needed. Per-response `tool_choice` overrides from `response.create` are supported. - -### Common output path - -Both handlers yield `(text, language_code, tools)` tuples. `LMOutputProcessor` forwards clean text to TTS and puts `{"type": "assistant_text", "text": ..., "tools": [...]}` on the `text_output_queue`. The router's `_send_loop` translates these into: -- `response.output_audio_transcript.done` for the text -- `response.function_call_arguments.done` for each tool call - -### Tool result flow - -1. Client executes the tool and sends `conversation.item.create` with `type: "function_call_output"` and `output: ""` -2. `RealtimeService` appends the tool output to the chat context and emits `conversation.item.created`; this does not trigger generation. -3. If the tool result needs to be spoken to the user, such as camera/search/data results, the client sends `response.create` to trigger follow-up generation. -4. For fire-and-forget robot actions such as dance, emotion, head movement, stop, or idle tools, the client can stop after `conversation.item.created`; the assistant should already have spoken the natural lead-in before the tool call. - ---- - -## Interruption Handling - -Barge-in (user speaks while the assistant is playing audio) is handled cooperatively between the VAD, the `_send_loop`, and the LLM/TTS handlers via a shared `CancelScope` object (`cancel_scope.py`). - -### CancelScope design - -`CancelScope` replaces the old two-signal pattern (`cancel_response` Event + `discard_stale_output` boolean) with a single object that manages: - -- **Generation counter** (`cancel_scope.generation`): pipeline threads (LLM, TTS) capture the current generation at the start of each response and check `cancel_scope.is_stale(gen)` on every streaming token. When `cancel()` is called, the generation increments and all prior generations become stale -- no timing games required. -- **Discard flag** (`cancel_scope.discarding`): set by `cancel()`, checked by the async `_send_loop` to drop output from superseded generations that arrives between `cancel()` and `response_done()`. Cleared by `response_done(generation)` (only when the sentinel's generation matches the discarded or current one -- sentinels from unrelated older generations are ignored), by `new_response()` on an explicit `response.create`, or by `reset()` on session claim/release. - -Pipeline output is **generation-tagged**: `AudioOutput` chunks and `AssistantTextEvent`s carry a `cancel_generation` field stamped by the handler that produced them. The send loop's `_generation_is_discardable` drops an item if its generation is stale, or if `discarding` is set and the item is not from the current generation. Output from the *current* generation always passes through, so a fresh response is never swallowed by a lingering discard window (e.g. a superseded speculative turn whose TTS never emitted a `__RESPONSE_DONE__` sentinel). - -```mermaid -sequenceDiagram - participant User - participant VAD - participant SendLoop as _send_loop - participant LLM - participant TTS - participant Client - - Note over TTS,Client: Response active or pending (in_response / response_pending) - User->>VAD: speaks - VAD->>SendLoop: speech_started on text_output_queue - SendLoop->>Client: input_audio_buffer.speech_started - SendLoop->>Client: response.done (status=cancelled, reason=turn_detected) - SendLoop->>SendLoop: cancel_scope.cancel() (gen++ & discarding=True) - SendLoop->>SendLoop: flush output_queue + text_output_queue - SendLoop->>SendLoop: response_playing.clear() - LLM->>LLM: is_stale(gen) → True, aborts generation - TTS->>TTS: is_stale(gen) → True, aborts generation - TTS->>SendLoop: __RESPONSE_DONE__ (tagged with gen) - SendLoop->>SendLoop: cancel_scope.response_done(gen) (discarding=False) - Note over VAD,Client: Pipeline is now processing the new user utterance -``` - -**Step by step:** - -1. **VAD detects speech**: puts a `SpeechStartedEvent` on `text_output_queue`. -2. **`_send_loop` processes text events first** (priority over audio): translates `speech_started` into protocol events. If an active response was in progress, `RealtimeService.dispatch_pipeline_event` emits `response.output_audio.done` + `response.done` with `status="cancelled"` and `reason="turn_detected"`. -3. **Cancel + queue flush**: if a response is active (`in_response`) *or* pending (`response_pending` -- a model request is queued, but no output has started), and interrupts are enabled (see step 4), the send loop calls `cancel_scope.cancel()` (increments generation, enables discard), clears `response_pending`, drains `output_queue` (preserving `__RESPONSE_DONE__` sentinels) and `text_output_queue` (preserving user-side events: `speech_stopped`, direct-audio completions, partial/completed transcriptions, token usage), then clears `response_playing`. -4. **Interrupt gating**: the cancel only fires if the `SpeechStartedEvent.interrupt_response` flag is set *and* the session config allows it (`turn_detection.interrupt_response`, read via `RuntimeConfig.interrupt_response_enabled`, default true). When disabled, user speech during a response is transcribed but the response keeps playing. -5. **LLM/TTS cancellation**: handlers capture `gen = cancel_scope.generation` at the start of each response and check `cancel_scope.is_stale(gen)` on every streaming token, aborting early when stale. -6. **Discard guard**: while `cancel_scope.discarding` is True, the send loop drops audio chunks and assistant text whose `cancel_generation` is not current (see `_generation_is_discardable` above). The guard clears when a `__RESPONSE_DONE__` with a matching generation arrives (via `cancel_scope.response_done(gen)`), or when an explicit `response.create` starts a new response (`cancel_scope.new_response()`). -7. **Client-initiated cancel**: `response.cancel` calls `cancel_scope.cancel()` (only if a response was active), flushes both queues with the same preservation rules, triggers `finish_response(status="cancelled", reason="client_cancelled")`, re-enables `should_listen`, and clears `response_playing`. -8. **Spurious cancel safety**: if no response is active, `cancel_scope.cancel()` is not called, preventing the discard guard from being set without a `__RESPONSE_DONE__` to clear it. - ---- - -## Testing - -### Local LLM with Transformers - -```bash -uv run speech-to-speech \ - --mode realtime \ - --stt parakeet-tdt \ - --llm_backend transformers \ - --tts kokoro \ - --model_name "Qwen/Qwen3-4B-Instruct-2507" \ - --llm_device mps \ - --llm_torch_dtype float16 \ - --enable_live_transcription -``` - -### Local LLM with MLX-LM - -```bash -uv run speech-to-speech \ - --mode realtime \ - --stt parakeet-tdt \ - --llm_backend mlx-lm \ - --tts kokoro \ - --model_name "mlx-community/Qwen3-4B-Instruct-2507-bf16" \ - --llm_device mps \ - --llm_torch_dtype float16 \ - --enable_live_transcription -``` - -### Remote LLM with OpenAI-compatible API - -```bash -uv run speech-to-speech \ - --mode realtime \ - --stt parakeet-tdt \ - --llm_backend responses-api \ - --tts kokoro \ - --model_name "openai/gpt-oss-20b:groq" \ - --responses_api_base_url "https://router.huggingface.co/v1" \ - --responses_api_api_key "$HF_TOKEN" \ - --responses_api_stream \ - --enable_live_transcription -``` diff --git a/runtime/src/speech_to_speech/api/openai_realtime/__init__.py b/runtime/src/speech_to_speech/api/openai_realtime/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/runtime/src/speech_to_speech/api/openai_realtime/handlers/__init__.py b/runtime/src/speech_to_speech/api/openai_realtime/handlers/__init__.py deleted file mode 100644 index a8d90db17cf18f4606deb87f24e6db5ce1fef82b..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/api/openai_realtime/handlers/__init__.py +++ /dev/null @@ -1,11 +0,0 @@ -from speech_to_speech.api.openai_realtime.handlers.audio import AudioHandler -from speech_to_speech.api.openai_realtime.handlers.conversation import ConversationHandler -from speech_to_speech.api.openai_realtime.handlers.response import ResponseHandler -from speech_to_speech.api.openai_realtime.handlers.session import SessionHandler - -__all__ = [ - "AudioHandler", - "ConversationHandler", - "ResponseHandler", - "SessionHandler", -] diff --git a/runtime/src/speech_to_speech/api/openai_realtime/handlers/audio.py b/runtime/src/speech_to_speech/api/openai_realtime/handlers/audio.py deleted file mode 100644 index 208f6f21578382d5c533838d16ceb7193133c92f..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/api/openai_realtime/handlers/audio.py +++ /dev/null @@ -1,224 +0,0 @@ -from __future__ import annotations - -import base64 -import logging -from typing import TYPE_CHECKING - -from openai.types.realtime import ( - InputAudioBufferAppendEvent, - InputAudioBufferSpeechStartedEvent, - InputAudioBufferSpeechStoppedEvent, - RealtimeErrorEvent, - ResponseAudioDeltaEvent, - ResponseCreatedEvent, -) - -from speech_to_speech.api.openai_realtime.handlers.base import RealtimeBaseHandler -from speech_to_speech.api.openai_realtime.utils import resample -from speech_to_speech.pipeline.events import SpeechStartedEvent, SpeechStoppedEvent - -if TYPE_CHECKING: - from speech_to_speech.api.openai_realtime.service import ServerEvent - -logger = logging.getLogger(__name__) - -PIPELINE_SAMPLE_RATE = 16000 -CHUNK_SAMPLES = 512 -BYTES_PER_SAMPLE = 2 -CHUNK_SIZE_BYTES = CHUNK_SAMPLES * BYTES_PER_SAMPLE - - -class AudioHandler(RealtimeBaseHandler): - """Owns inbound audio decoding/chunking and outbound audio encoding.""" - - def _start_input_item(self, conn_id: str, *, preserve_active_response: bool = False) -> str: - response = self._service.response - st = self._state(conn_id) - if not preserve_active_response: - item_id = response._start_item(conn_id) - else: - response_item_id = st.current_item_id - response_content_index = st.content_index - item_id = response._start_item(conn_id) - st.current_item_id = response_item_id - st.content_index = response_content_index - st.input_content_index = 0 - return item_id - - def handle_audio_append(self, conn_id: str, event: InputAudioBufferAppendEvent) -> list[bytes]: - """Decode base64 audio, resample to pipeline rate, and split into 512-sample PCM16 chunks for the VAD.""" - try: - pcm_bytes = base64.b64decode(event.audio) - except Exception as e: - logger.error(f"Base64 decode error: {e}") - return [] - - st = self._state(conn_id) - - audio_cfg = st.runtime_config.session.audio - if audio_cfg is not None and audio_cfg.input is not None: - client_in_rate = getattr(audio_cfg.input.format, "rate", None) or PIPELINE_SAMPLE_RATE - else: - client_in_rate = PIPELINE_SAMPLE_RATE - return self.append_pcm(conn_id, pcm_bytes, client_in_rate) - - def append_pcm(self, conn_id: str, pcm_bytes: bytes, src_rate: int) -> list[bytes]: - """Resample raw PCM16 to the pipeline rate and split into 512-sample chunks for the VAD. - - Shared by both transports: the WebSocket route feeds it decoded - ``input_audio_buffer.append`` payloads, the WebRTC transport feeds it - PCM decoded from inbound media-track frames. Keeps the sub-chunk - remainder and the commit bookkeeping (``audio_buffer_has_data``) in - one place regardless of how audio arrives. - """ - st = self._state(conn_id) - pcm_bytes = resample(pcm_bytes, src_rate, PIPELINE_SAMPLE_RATE) - - pcm_bytes = st.audio_remainder + pcm_bytes - - chunks = [] - for i in range(0, len(pcm_bytes), CHUNK_SIZE_BYTES): - chunk = pcm_bytes[i : i + CHUNK_SIZE_BYTES] - if len(chunk) == CHUNK_SIZE_BYTES: - chunks.append(chunk) - else: - st.audio_remainder = chunk - break - else: - st.audio_remainder = b"" - - if chunks: - st.audio_buffer_has_data = True - return chunks - - def handle_audio_commit(self, conn_id: str) -> RealtimeErrorEvent | None: - """Commit the audio buffer. Returns an error if no audio was appended.""" - st = self._state(conn_id) - if not st.audio_buffer_has_data: - return self.make_error( - message="Input audio buffer is empty, nothing to commit.", - _type="input_audio_buffer_commit_empty", - ) - st.audio_buffer_has_data = False - logger.debug("Audio buffer committed") - return None - - # ── Pipeline event handlers ──────────────────── - - def on_speech_started(self, conn_id: str, event: SpeechStartedEvent) -> list[ServerEvent]: - """Handle VAD speech_started: cancel active response if interrupts enabled, start new input item.""" - response = self._service.response - events: list[ServerEvent] = [] - st = self._state(conn_id) - if st.in_response and event.interrupt_response and st.runtime_config.interrupt_response_enabled: - events.extend(response.finish_response(conn_id, status="cancelled", reason="turn_detected")) - is_reopen = bool(event.reopened and event.turn_id is not None and event.turn_id == st.speculative_turn_id) - preserve_active_response = st.in_response - if is_reopen: - input_item_id = st.speculative_input_item_id - if input_item_id is None: - input_item_id = self._start_input_item( - conn_id, - preserve_active_response=preserve_active_response, - ) - st.speculative_input_item_id = input_item_id - elif not preserve_active_response: - st.current_item_id = input_item_id - st.content_index = 0 - st.input_audio_duration_s = 0.0 - st.input_content_index = 0 - else: - input_item_id = self._start_input_item( - conn_id, - preserve_active_response=preserve_active_response, - ) - st.speculative_input_item_id = input_item_id - st.response_usage.turns += 1 - st.speculative_turn_id = event.turn_id - st.speculative_turn_revision = event.turn_revision - st.last_item_id = input_item_id - events.append( - InputAudioBufferSpeechStartedEvent( - type="input_audio_buffer.speech_started", - event_id=self._next_event_id(), - audio_start_ms=event.audio_start_ms, - item_id=input_item_id, - ) - ) - return events - - def on_speech_stopped(self, conn_id: str, event: SpeechStoppedEvent) -> list[ServerEvent]: - """Handle VAD speech_stopped: record duration and emit stopped event.""" - if event.duration_s: - self._state(conn_id).input_audio_duration_s = event.duration_s - return [ - InputAudioBufferSpeechStoppedEvent( - type="input_audio_buffer.speech_stopped", - event_id=self._next_event_id(), - audio_end_ms=event.audio_end_ms, - item_id=self._input_item_id(conn_id), - ) - ] - - # ── Outbound audio encoding ────────────────── - - def begin_audio_response(self, conn_id: str) -> tuple[str, str, list[ServerEvent]]: - """Ensure a response exists for outbound audio, emitting ResponseCreated once. - - When ``handle_response_create`` already allocated the response, - ``current_response_id`` is set and no duplicate event is emitted. - For the implicit-response path (VAD -> STT -> LLM -> TTS, no - ``response.create``), ``current_response_id`` is still ``None`` - and the event is emitted here on the first audio chunk. - - Returns ``(response_id, item_id, events)``. Shared by both - transports: the WebSocket path appends the base64 audio delta to the - returned events, the WebRTC path sends only the bookkeeping events - over the data channel while audio travels on the media track. - """ - response = self._service.response - st = self._state(conn_id) - - events: list[ServerEvent] = [] - need_created = st.current_response_id is None - resp_id, item_id = response._ensure_response(conn_id) - if need_created: - events.append( - ResponseCreatedEvent( - type="response.created", - event_id=self._next_event_id(), - response=response._build_response(conn_id, "in_progress"), - ) - ) - return resp_id, item_id, events - - def encode_audio_chunk(self, conn_id: str, audio: bytes) -> list[ServerEvent]: - """Encode a raw PCM audio chunk as a base64 delta event for the WebSocket transport.""" - response = self._service.response - st = self._state(conn_id) - - resp_id, item_id, events = self.begin_audio_response(conn_id) - rp = st.current_response_params - client_out_rate = None - if rp and rp.audio and rp.audio.output and rp.audio.output.format: - client_out_rate = getattr(rp.audio.output.format, "rate", None) - if client_out_rate is None: - audio_cfg = st.runtime_config.session.audio - if audio_cfg is not None and audio_cfg.output is not None: - client_out_rate = getattr(audio_cfg.output.format, "rate", None) or PIPELINE_SAMPLE_RATE - else: - client_out_rate = PIPELINE_SAMPLE_RATE - audio = resample(audio, PIPELINE_SAMPLE_RATE, client_out_rate) - b64 = base64.b64encode(audio).decode("ascii") - events.append( - ResponseAudioDeltaEvent( - type="response.output_audio.delta", - event_id=self._next_event_id(), - content_index=response._next_content_index(conn_id), - delta=b64, - item_id=item_id, - output_index=0, - response_id=resp_id, - ) - ) - return events diff --git a/runtime/src/speech_to_speech/api/openai_realtime/handlers/base.py b/runtime/src/speech_to_speech/api/openai_realtime/handlers/base.py deleted file mode 100644 index abc9083c9aee687db630b3d3c74ac0ac946216d9..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/api/openai_realtime/handlers/base.py +++ /dev/null @@ -1,57 +0,0 @@ -from __future__ import annotations - -import logging -from queue import Queue -from threading import Event as ThreadingEvent -from typing import TYPE_CHECKING - -from openai.types.realtime import RealtimeErrorEvent - -from speech_to_speech.utils.utils import _generate_id - -if TYPE_CHECKING: - from speech_to_speech.api.openai_realtime.service import ConnState, RealtimeService - -logger = logging.getLogger(__name__) - - -class RealtimeBaseHandler: - """Shared base for domain handlers. - - Provides conn_id-keyed accessors for per-connection state, config, - and queues. Each connection owns its own ``RuntimeConfig`` via - ``ConnState``. - """ - - def __init__(self, service: RealtimeService) -> None: - self._service = service - - # ── conn_id-keyed accessors ── - - def _queue(self, conn_id: str) -> Queue | None: - return self._service.text_prompt_queue - - def _should_listen(self, conn_id: str) -> ThreadingEvent | None: - return self._service.should_listen - - # ── shared helpers ──────────────────────────────── - - def _state(self, conn_id: str) -> ConnState: - return self._service._state(conn_id) - - def _input_item_id(self, conn_id: str) -> str: - st = self._state(conn_id) - return st.speculative_input_item_id or self._service.response._current_item_id(conn_id) - - def _next_input_content_index(self, conn_id: str) -> int: - st = self._state(conn_id) - idx = st.input_content_index - st.input_content_index += 1 - return idx - - @staticmethod - def _next_event_id() -> str: - return _generate_id("event") - - def make_error(self, message: str, _type: str) -> RealtimeErrorEvent: - return self._service.make_error(message, _type) diff --git a/runtime/src/speech_to_speech/api/openai_realtime/handlers/conversation.py b/runtime/src/speech_to_speech/api/openai_realtime/handlers/conversation.py deleted file mode 100644 index 38279e165d08a4588a94df3b64f322bb0d233867..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/api/openai_realtime/handlers/conversation.py +++ /dev/null @@ -1,128 +0,0 @@ -from __future__ import annotations - -import logging -from typing import TYPE_CHECKING - -from openai.types.realtime import ( - ConversationItem, - ConversationItemCreatedEvent, - ConversationItemCreateEvent, - ConversationItemInputAudioTranscriptionCompletedEvent, - ConversationItemInputAudioTranscriptionDeltaEvent, -) -from openai.types.realtime.conversation_item_input_audio_transcription_completed_event import ( - UsageTranscriptTextUsageDuration, -) - -from speech_to_speech.api.openai_realtime.handlers.base import RealtimeBaseHandler -from speech_to_speech.LLM.chat import ChatItemError, add_supported_item -from speech_to_speech.pipeline.events import PartialTranscriptionEvent, TranscriptionCompletedEvent - -if TYPE_CHECKING: - from speech_to_speech.api.openai_realtime.service import ServerEvent - -logger = logging.getLogger(__name__) - - -class ConversationHandler(RealtimeBaseHandler): - """Owns conversation item injection and pipeline-to-protocol translation.""" - - def handle_conversation_item_create( - self, - conn_id: str, - event: ConversationItemCreateEvent, - ) -> list[ServerEvent]: - """Inject a text message or function-call output into the LLM context. - - Items are added to the LLM chat context but do NOT trigger response - generation on their own. A subsequent ``response.create`` event is - required to trigger the model. - - While a response is generating, the item is *deferred*: applying it now - would race the LLM handler's end-of-turn chat write-back, which runs on - the pipeline thread (e.g. a ``function_call_output`` arriving before its - ``function_call`` is recorded, or an image stripped before the next - turn reads it). Deferred items are flushed, in order, once the response - completes — see :meth:`flush_deferred_items`. - """ - st = self._state(conn_id) - if st.in_response: - st.deferred_items.append(event.item) - logger.debug("Deferred conversation item until the active response completes") - return [] - return self._apply_item(conn_id, event.item) - - def _apply_item(self, conn_id: str, item: ConversationItem) -> list[ServerEvent]: - """Add one item to the chat and build its ``conversation.item.created``.""" - try: - self._append_item(conn_id, item) - except ChatItemError as exc: - return [self.make_error(str(exc), "invalid_conversation_item")] - - if not item: - return [] - st = self._state(conn_id) - event = ConversationItemCreatedEvent( - type="conversation.item.created", - event_id=self._next_event_id(), - previous_item_id=st.last_item_id, - item=item, - ) - st.last_item_id = item.id - return [event] - - def flush_deferred_items(self, conn_id: str) -> list[ServerEvent]: - """Apply items buffered during a response, in arrival order. - - Called at response completion (after the generation's own write-back), - so a ``function_call_output`` pairs with its now-recorded ``function_call`` - and an image survives the just-finished response's ``strip_images``. - """ - st = self._state(conn_id) - if not st.deferred_items: - return [] - items = st.deferred_items - st.deferred_items = [] - events: list[ServerEvent] = [] - for item in items: - events.extend(self._apply_item(conn_id, item)) - return events - - def _append_item(self, conn_id: str, item: ConversationItem) -> None: - """Narrow ``ConversationItem`` to ``SupportedItem`` and delegate to ``Chat.add_item``. - - Raises :class:`ChatItemError` on validation failure or unsupported type. - """ - add_supported_item(self._state(conn_id).runtime_config.chat, item) - - # ── Pipeline event handlers ──────────────────── - - def on_partial_transcription(self, conn_id: str, event: PartialTranscriptionEvent) -> list[ServerEvent]: - """Handle partial_transcription: emit transcription delta event.""" - return [ - ConversationItemInputAudioTranscriptionDeltaEvent( - type="conversation.item.input_audio_transcription.delta", - event_id=self._next_event_id(), - content_index=self._next_input_content_index(conn_id), - item_id=self._input_item_id(conn_id), - delta=event.delta, - ) - ] - - def on_transcription_completed(self, conn_id: str, event: TranscriptionCompletedEvent) -> list[ServerEvent]: - """Handle transcription_completed: accumulate duration and emit completed event.""" - st = self._state(conn_id) - st.response_usage.audio_duration_s += st.input_audio_duration_s - return [ - ConversationItemInputAudioTranscriptionCompletedEvent( - type="conversation.item.input_audio_transcription.completed", - event_id=self._next_event_id(), - content_index=0, - item_id=self._input_item_id(conn_id), - transcript=event.transcript, - usage=UsageTranscriptTextUsageDuration( - seconds=st.input_audio_duration_s, - type="duration", - ), - ) - ] diff --git a/runtime/src/speech_to_speech/api/openai_realtime/handlers/response.py b/runtime/src/speech_to_speech/api/openai_realtime/handlers/response.py deleted file mode 100644 index 2ee9fdd08251842b4a7293a7ed7baa601280ba6e..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/api/openai_realtime/handlers/response.py +++ /dev/null @@ -1,339 +0,0 @@ -from __future__ import annotations - -import logging -from typing import TYPE_CHECKING, Optional - -from openai.types.realtime import ( - RealtimeResponse, - ResponseAudioDoneEvent, - ResponseAudioTranscriptDoneEvent, - ResponseCreatedEvent, - ResponseCreateEvent, - ResponseDoneEvent, - ResponseFunctionCallArgumentsDoneEvent, - ResponseTextDeltaEvent, - ResponseTextDoneEvent, -) -from openai.types.realtime.realtime_response import Audio, AudioOutput -from openai.types.realtime.realtime_response_status import RealtimeResponseStatus -from openai.types.realtime.realtime_response_usage import RealtimeResponseUsage - -from speech_to_speech.api.openai_realtime.handlers.base import RealtimeBaseHandler -from speech_to_speech.LLM.chat import ChatItemError -from speech_to_speech.pipeline.events import AssistantTextEvent -from speech_to_speech.pipeline.messages import GenerateResponseRequest -from speech_to_speech.utils.utils import _generate_id, is_out_of_band, response_wants_audio - -if TYPE_CHECKING: - from speech_to_speech.api.openai_realtime.service import ServerEvent, _ResponseStatus, _StatusReason - -logger = logging.getLogger(__name__) - - -class ResponseHandler(RealtimeBaseHandler): - """Owns the response lifecycle: create, cancel, finish, and ID management.""" - - # ── ID / state helpers ──────────────────────── - - def _ensure_response(self, conn_id: str) -> tuple[str, str]: - """Ensure a response and output item exist, creating them if needed.""" - st = self._state(conn_id) - if st.current_response_id is None: - st.current_response_id = _generate_id("resp") - self._start_item(conn_id) - st.in_response = True - st.response_pending = False - return st.current_response_id, self._current_item_id(conn_id) - - def _end_response(self, conn_id: str, status: _ResponseStatus = "completed") -> None: - st = self._state(conn_id) - if status == "cancelled": - st.response_usage.responses_cancelled += 1 - else: - st.response_usage.responses_completed += 1 - self._service.total_usage += st.response_usage - logger.info( - "Response done (status=%s) — this response: input_tokens=%d, output_tokens=%d, audio=%.2fs" - " | cumulative: input_tokens=%d, output_tokens=%d, audio=%.2fs", - status, - st.response_usage.input_tokens, - st.response_usage.output_tokens, - st.response_usage.audio_duration_s, - self._service.total_usage.input_tokens, - self._service.total_usage.output_tokens, - self._service.total_usage.audio_duration_s, - ) - st.response_usage.reset() - st.current_response_id = None - st.current_item_id = None - st.content_index = 0 - st.in_response = False - st.response_pending = False - st.current_response_params = None - st.pending_output_text_parts = [] - - def _start_item(self, conn_id: str) -> str: - """Generate a new item ID, reset content index, and store it.""" - st = self._state(conn_id) - item_id = _generate_id("item") - st.current_item_id = item_id - st.content_index = 0 - st.input_audio_duration_s = 0.0 - return item_id - - def _current_item_id(self, conn_id: str) -> str: - return self._state(conn_id).current_item_id or self._start_item(conn_id) - - def _next_content_index(self, conn_id: str) -> int: - """Return the current content index and advance it.""" - st = self._state(conn_id) - idx = st.content_index - st.content_index += 1 - return idx - - def _build_response( - self, - conn_id: str, - status: _ResponseStatus, - reason: _StatusReason | None = None, - ) -> RealtimeResponse: - """Build a fully-populated RealtimeResponse from the current connection state.""" - st = self._state(conn_id) - status_details = None - if reason or status in ("completed", "cancelled", "incomplete", "failed"): - status_details = RealtimeResponseStatus(type=status, reason=reason) # type: ignore[arg-type] - - rp = st.current_response_params - metadata = rp.metadata if rp and rp.metadata else None - - voice: Optional[str] = None - if rp and rp.audio and rp.audio.output and rp.audio.output.voice: - voice = str(rp.audio.output.voice) - if not voice: - audio_cfg = st.runtime_config.session.audio - audio_output = audio_cfg.output if audio_cfg is not None else None - voice = str(audio_output.voice) if audio_output is not None and audio_output.voice else None - - # Out-of-band responses are not threaded into any conversation: report a null id. - conversation_id = None if is_out_of_band(rp) else st.conversation_id - - return RealtimeResponse( - id=st.current_response_id, - object="realtime.response", - status=status, - status_details=status_details, - audio=Audio(output=AudioOutput(voice=str(voice) if voice else None)), # type: ignore[arg-type] - conversation_id=conversation_id, - metadata=metadata, - usage=RealtimeResponseUsage( - input_tokens=st.response_usage.input_tokens, - output_tokens=st.response_usage.output_tokens, - total_tokens=st.response_usage.input_tokens + st.response_usage.output_tokens, - ), - ) - - # ── Public handlers ─────────────────────────── - - def handle_response_create(self, conn_id: str, event: ResponseCreateEvent) -> ServerEvent | None: - """Trigger a response. - - Returns a ``ResponseCreatedEvent`` on success, a ``RealtimeErrorEvent`` - on failure, or ``None`` if there is no text_prompt_queue. - """ - st = self._state(conn_id) - if event.response: - if event.response.tool_choice and not isinstance(event.response.tool_choice, str): - return self.make_error( - message="Only string tool_choice values are supported for now (auto, required, none).", - _type="tool_choice_not_supported", - ) - if st.in_response: - return self.make_error( - message="Cannot create response while another response is in progress.", - _type="conversation_already_has_active_response", - ) - - out_of_band = is_out_of_band(event.response) - - # In-band: response.input items are added to the default conversation here so - # they appear in history. Out-of-band: leave the default conversation untouched — - # the input rides along on the request and seeds a throwaway chat in the LM. - if not out_of_band and event.response and event.response.input: - for input_item in event.response.input: - try: - self._service.conversation._append_item(conn_id, input_item) - except ChatItemError as exc: - return self.make_error(message=str(exc), _type="invalid_input_item") - - st.in_response = True - st.response_pending = False - - st.current_response_params = event.response - st.current_response_id = _generate_id("resp") - self._start_item(conn_id) - - cfg = st.runtime_config - queue = self._queue(conn_id) - if queue: - # Out-of-band responses carry no turn identity: a null turn_id makes every - # speculative-turn staleness gate treat them as always-latest, so a new user - # turn mid-generation can never silently drop their output. - queue.put( - GenerateResponseRequest( - runtime_config=cfg, - response=event.response, - turn_id=None if out_of_band else st.speculative_user_turn_id, - turn_revision=None if out_of_band else st.speculative_user_turn_revision, - speech_stopped_at_s=None if out_of_band else st.speculative_user_speech_stopped_at_s, - ) - ) - logger.debug("response.create received, LLM generation triggered") - return ResponseCreatedEvent( - type="response.created", - event_id=self._next_event_id(), - response=self._build_response(conn_id, "in_progress"), - ) - - def handle_response_cancel(self, conn_id: str) -> list[ServerEvent]: - """Cancel the in-progress response and re-enable listening.""" - events = self.finish_response(conn_id, status="cancelled", reason="client_cancelled") - should_listen = self._should_listen(conn_id) - if should_listen: - should_listen.set() - logger.info("Response cancelled, listening re-enabled") - return events - - def finish_response( - self, - conn_id: str, - status: _ResponseStatus = "completed", - reason: _StatusReason | None = None, - ) -> list[ServerEvent]: - """Close the current response (audio/text done + response done). - - Audio responses emit ``response.output_audio.done`` for any terminal - status. Text-only responses emit a single ``response.output_text.done`` - carrying the full streamed text, but only on ``status="completed"`` — - a cancelled or failed text response sends no audio, so it just closes - with ``response.done``. - """ - st = self._state(conn_id) - events: list[ServerEvent] = [] - if st.in_response: - resp_id, item_id = self._ensure_response(conn_id) - if response_wants_audio(st.current_response_params): - events.append( - ResponseAudioDoneEvent( - type="response.output_audio.done", - event_id=self._next_event_id(), - content_index=0, - item_id=item_id, - output_index=0, - response_id=resp_id, - ) - ) - elif status == "completed" and st.pending_output_text_parts: - events.append( - ResponseTextDoneEvent( - type="response.output_text.done", - event_id=self._next_event_id(), - content_index=0, - item_id=item_id, - output_index=0, - response_id=resp_id, - text="".join(st.pending_output_text_parts), - ) - ) - events.append( - ResponseDoneEvent( - type="response.done", - event_id=self._next_event_id(), - response=self._build_response(conn_id, status, reason), - ) - ) - self._end_response(conn_id, status) - # Apply any client items that arrived mid-generation now that in_response - # is cleared and the generation's own write-back has landed. Done outside - # the in_response guard so a stray terminal call still drains the buffer. - events.extend(self._service.conversation.flush_deferred_items(conn_id)) - return events - - # ── Pipeline event handlers ─────────────────── - - def on_assistant_text( - self, - conn_id: str, - event: AssistantTextEvent, - *, - wait_for_pending_reopen: bool = True, - ) -> list[ServerEvent] | None: - """Handle assistant_text: emit transcript and/or tool-call events.""" - if self._service.speculative_turns: - commit_result: bool | None - if wait_for_pending_reopen: - commit_result = self._service.speculative_turns.commit_if_latest_after_reopen_grace( - event.turn_id, - event.turn_revision, - ) - else: - commit_result = self._service.speculative_turns.try_commit_if_latest_after_reopen_grace( - event.turn_id, - event.turn_revision, - ) - if commit_result is None: - return None - if not commit_result: - logger.debug("Dropping stale assistant text for turn=%s rev=%s", event.turn_id, event.turn_revision) - return [] - st = self._state(conn_id) - events: list[ServerEvent] = [] - resp_id, item_id = self._ensure_response(conn_id) - st.last_item_id = item_id - output_idx = 0 - if event.text: - if response_wants_audio(st.current_response_params): - events.append( - ResponseAudioTranscriptDoneEvent( - type="response.output_audio_transcript.done", - event_id=self._next_event_id(), - content_index=0, - item_id=item_id, - output_index=output_idx, - response_id=resp_id, - transcript=event.text, - ) - ) - else: - # Stream the delta now; the matching response.output_text.done is - # emitted once at close in finish_response, carrying the per-chunk - # parts collected here, space-joined ("" when there were none). - st.pending_output_text_parts.append(event.text) - events.append( - ResponseTextDeltaEvent( - type="response.output_text.delta", - event_id=self._next_event_id(), - content_index=0, - item_id=item_id, - output_index=output_idx, - response_id=resp_id, - delta=event.text, - ) - ) - output_idx += 1 - if event.tools: - st.response_usage.tool_calls += len(event.tools) - for tool in event.tools: - events.append( - ResponseFunctionCallArgumentsDoneEvent( - type="response.function_call_arguments.done", - event_id=self._next_event_id(), - call_id=tool.call_id, - name=tool.name, - arguments=tool.arguments, - item_id=item_id, - output_index=output_idx, - response_id=resp_id, - ) - ) - output_idx += 1 - return events diff --git a/runtime/src/speech_to_speech/api/openai_realtime/handlers/session.py b/runtime/src/speech_to_speech/api/openai_realtime/handlers/session.py deleted file mode 100644 index cf9f055c23f2a90145a2e67f8ae4c8c6fb1018cf..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/api/openai_realtime/handlers/session.py +++ /dev/null @@ -1,62 +0,0 @@ -from __future__ import annotations - -import logging -from typing import Optional - -from openai.types.realtime import ( - RealtimeErrorEvent, - SessionCreatedEvent, - SessionUpdateEvent, -) -from openai.types.realtime.realtime_transcription_session_create_request import ( - RealtimeTranscriptionSessionCreateRequest, -) - -from speech_to_speech.api.openai_realtime.handlers.base import RealtimeBaseHandler - -logger = logging.getLogger(__name__) - - -class SessionHandler(RealtimeBaseHandler): - """Owns session lifecycle: config updates and session.created events.""" - - def handle_session_update(self, conn_id: str, event: SessionUpdateEvent) -> Optional[RealtimeErrorEvent]: - """Apply session config changes. - - Only ``RealtimeSessionCreateRequest`` sessions are accepted; - ``RealtimeTranscriptionSessionCreateRequest`` sessions not yet supported. - Incoming fields are deep-merged into the existing session so that - partial updates preserve previously-set values. - """ - s = event.session - if s is None: - return None - - if isinstance(s, RealtimeTranscriptionSessionCreateRequest): - return self.make_error( - message="Only 'realtime' session type is supported; transcription sessions are not.", - _type="invalid_session_type", - ) - - model = getattr(s, "model", None) - if model is not None: - logger.info(f"Session model set to: {model}") - - cfg = self._state(conn_id).runtime_config - current = cfg.session - if current is None: - cfg.session = s - else: - cfg.apply_session_update(s) - logger.info("Session configuration updated") - return None - - def build_session_created(self, conn_id: str) -> SessionCreatedEvent: - """Build a SessionCreatedEvent populated with the current config.""" - cfg = self._state(conn_id).runtime_config - session = cfg.session - return SessionCreatedEvent( - type="session.created", - event_id=self._next_event_id(), - session=session, - ) diff --git a/runtime/src/speech_to_speech/api/openai_realtime/pipeline_unit.py b/runtime/src/speech_to_speech/api/openai_realtime/pipeline_unit.py deleted file mode 100644 index 84c844f619f08ca3a732c2575ad9ddbb8188f360..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/api/openai_realtime/pipeline_unit.py +++ /dev/null @@ -1,72 +0,0 @@ -import asyncio -from queue import Queue -from threading import Event -from typing import Any, Optional - -from pydantic import BaseModel, ConfigDict, Field - -from speech_to_speech.api.openai_realtime.service import RealtimeService -from speech_to_speech.api.openai_realtime.transports import SessionTransport -from speech_to_speech.pipeline.cancel_scope import CancelScope - - -class SessionState(BaseModel): - """Per-client ephemeral state. - - Created when a route handler claims a PipelineUnit (WebSocket accept or - WebRTC SDP offer); dropped when the client disconnects. Holding the - transport reference, the service session id, and any send-loop scratch - (pending_output_item) here ensures these fields share one lifecycle — a - stale value can't outlive its session. - - `transport` may briefly be None while a WebRTC claim finishes - constructing the session, so the send loop must tolerate a - transport-less snapshot. - - `drained` is set by the send loop when SESSION_END travels through the handler - chain back to the output queue; the release path awaits it before clearing - `PipelineUnit.session`, so a new client cannot claim the unit until in-flight - work from this session has fully reset. - """ - - model_config = ConfigDict(arbitrary_types_allowed=True) - - transport: Optional[SessionTransport] = None - session_id: str = "" - pending_output_item: Any = None - drained: asyncio.Event = Field(default_factory=asyncio.Event) - # Wall-clock time when the client disconnected (route handler released its - # claim). `None` while the client is still active. Used by /v1/pool to - # surface stuck units (handlers haven't finished propagating SESSION_END). - released_at: Optional[float] = None - # Wall-clock time when the drain wait gave up and quarantined the unit - # (SESSION_END_QUARANTINE_TIMEOUT_S elapsed). The unit stays unclaimable — - # its handlers may still emit this session's output — until SESSION_END - # actually drains. Reported as "stuck" by /v1/pool. - quarantined_at: Optional[float] = None - - -class PipelineUnit(BaseModel): - """One isolated realtime pipeline. - - Each unit owns its queues, events, RealtimeService, and the chain of handler - instances (VAD, STT, transcription notifier, LM, LM output processor, TTS). - Lives inside the pool managed by RealtimeServer; the websocket route handler - claims a free unit (`session is None`) on `accept` and releases it on disconnect - by setting `session` back to None. - """ - - model_config = ConfigDict(arbitrary_types_allowed=True) - - index: int - service: RealtimeService - cancel_scope: CancelScope - should_listen: Event - response_playing: Event - input_queue: Queue - output_queue: Queue - text_output_queue: Queue - text_prompt_queue: Queue - handlers: list[Any] - - session: Optional[SessionState] = None diff --git a/runtime/src/speech_to_speech/api/openai_realtime/runtime_config.py b/runtime/src/speech_to_speech/api/openai_realtime/runtime_config.py deleted file mode 100644 index f79b4adcffac7b23d1dbe7d83cb9114201e61f53..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/api/openai_realtime/runtime_config.py +++ /dev/null @@ -1,81 +0,0 @@ -from openai.types.realtime import RealtimeSessionCreateRequest -from openai.types.realtime.realtime_audio_config import RealtimeAudioConfig -from openai.types.realtime.realtime_audio_config_input import RealtimeAudioConfigInput -from openai.types.realtime.realtime_audio_config_output import RealtimeAudioConfigOutput -from pydantic import BaseModel, ConfigDict, Field, field_validator - -from speech_to_speech.LLM.chat import Chat - - -def _apply_update(current: BaseModel, update: BaseModel) -> None: - """Apply explicitly-set fields from *update* onto *current* in-place, - recursing into nested BaseModel children so partial nested updates - don't overwrite unset fields. - - Only fields present in update.model_fields_set (i.e. actually - sent by the client) are considered. - """ - for field_name in update.model_fields_set: - new_val = getattr(update, field_name) - old_val = getattr(current, field_name, None) - if isinstance(new_val, BaseModel) and isinstance(old_val, BaseModel): - _apply_update(old_val, new_val) - else: - setattr(current, field_name, new_val) - - -class RuntimeConfig(BaseModel): - """ - Shared mutable configuration written by the RealtimeService on - session.update and read by pipeline handlers (VAD, LLM, TTS) during - processing. Python's GIL makes simple attribute reads/writes atomic, - so no explicit locking is needed for primitive values. - - The canonical state lives in 'session' (a full - 'RealtimeSessionCreateRequest'). - """ - - model_config = ConfigDict(validate_assignment=True, arbitrary_types_allowed=True) - - chat: Chat = Field(default_factory=lambda: Chat(10)) - session: RealtimeSessionCreateRequest = Field( - default_factory=lambda: RealtimeSessionCreateRequest(type="realtime"), - validate_default=True, - ) - - @field_validator("session", mode="after") - @classmethod - def _ensure_audio_structure(cls, v: RealtimeSessionCreateRequest) -> RealtimeSessionCreateRequest: - """Guarantee 'audio.input' and 'audio.output' are never None.""" - if v.audio is None: - v.audio = RealtimeAudioConfig() - if v.audio.input is None: - v.audio.input = RealtimeAudioConfigInput() - if v.audio.output is None: - v.audio.output = RealtimeAudioConfigOutput() - return v - - @property - def interrupt_response_enabled(self) -> bool: - """Whether barge-in should cancel an active response. - - Reads 'turn_detection.interrupt_response' from the session config, - handling both Pydantic models ('ServerVad') and plain dicts. - Defaults to 'True' (OpenAI API default). - """ - assert self.session.audio is not None and self.session.audio.input is not None - td = self.session.audio.input.turn_detection - if td is None: - return True - if hasattr(td, "interrupt_response"): - val = td.interrupt_response - elif isinstance(td, dict): - val = td.get("interrupt_response", True) - else: - return True - return val if val is not None else True - - def apply_session_update(self, update: RealtimeSessionCreateRequest) -> None: - """Merge non-None, explicitly-set fields from 'update' into the - current 'session', preserving any fields not present in the update.""" - _apply_update(self.session, update) diff --git a/runtime/src/speech_to_speech/api/openai_realtime/server.py b/runtime/src/speech_to_speech/api/openai_realtime/server.py deleted file mode 100644 index 4b3b5deeb98ea6f4dcd12b9af6c1fd07d0b538c7..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/api/openai_realtime/server.py +++ /dev/null @@ -1,61 +0,0 @@ -import logging -import threading -from threading import Event - -import uvicorn - -from speech_to_speech.api.openai_realtime.pipeline_unit import PipelineUnit -from speech_to_speech.api.openai_realtime.websocket_router import create_app - -logger = logging.getLogger(__name__) - - -class RealtimeServer: - """ - Pipeline handler for the OpenAI Realtime API mode. - - Owns a pool of isolated PipelineUnits and a single uvicorn server. - The websocket route claims the next free unit on each accept and releases - it on disconnect; once all units are in use, further connections are rejected. - """ - - def __init__( - self, - stop_event: Event, - pool: list[PipelineUnit], - host: str = "0.0.0.0", - port: int = 8765, - ) -> None: - if not pool: - raise ValueError("RealtimeServer requires at least one PipelineUnit in the pool") - self.stop_event = stop_event - self.pool = pool - self.host = host - self.port = port - - def run(self) -> None: - """Start the FastAPI/uvicorn server (called from a ThreadManager thread).""" - app = create_app(pool=self.pool, stop_event=self.stop_event) - - logger.info( - f"OpenAI Realtime API starting on ws://{self.host}:{self.port}/v1/realtime (pool size {len(self.pool)})" - ) - - config = uvicorn.Config( - app, - host=self.host, - port=self.port, - log_level="info", - ) - server = uvicorn.Server(config) - - server.install_signal_handlers = lambda: None # type: ignore[attr-defined] - - def _watch_stop() -> None: - self.stop_event.wait() - server.should_exit = True - - watcher = threading.Thread(target=_watch_stop, daemon=True) - watcher.start() - - server.run() diff --git a/runtime/src/speech_to_speech/api/openai_realtime/service.py b/runtime/src/speech_to_speech/api/openai_realtime/service.py deleted file mode 100644 index d9c27b770386a6a2f5804511327ead16d4c9ab98..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/api/openai_realtime/service.py +++ /dev/null @@ -1,566 +0,0 @@ -import logging -from collections.abc import Mapping -from queue import Queue -from threading import Event as ThreadingEvent -from typing import Any, Callable, Literal, Optional, TypeVar, Union - -from openai.types.realtime import ( - ConversationItem, - ConversationItemCreatedEvent, - ConversationItemCreateEvent, - ConversationItemInputAudioTranscriptionCompletedEvent, - ConversationItemInputAudioTranscriptionDeltaEvent, - InputAudioBufferAppendEvent, - InputAudioBufferCommitEvent, - InputAudioBufferSpeechStartedEvent, - InputAudioBufferSpeechStoppedEvent, - OutputAudioBufferClearEvent, - RealtimeError, - RealtimeErrorEvent, - ResponseAudioDeltaEvent, - ResponseAudioDoneEvent, - ResponseAudioTranscriptDoneEvent, - ResponseCancelEvent, - ResponseCreatedEvent, - ResponseCreateEvent, - ResponseDoneEvent, - ResponseFunctionCallArgumentsDoneEvent, - ResponseTextDeltaEvent, - ResponseTextDoneEvent, - SessionCreatedEvent, - SessionUpdateEvent, -) -from openai.types.realtime.realtime_response_create_params import RealtimeResponseCreateParams -from pydantic import BaseModel, ConfigDict, Field, ValidationError - -from speech_to_speech.api.openai_realtime.handlers import ( - AudioHandler, - ConversationHandler, - ResponseHandler, - SessionHandler, -) -from speech_to_speech.api.openai_realtime.runtime_config import RuntimeConfig -from speech_to_speech.LLM.chat import Chat, make_user_message -from speech_to_speech.pipeline.events import ( - AssistantTextEvent, - AudioInputCompletedEvent, - PartialTranscriptionEvent, - PipelineEvent, - ResponseFailedEvent, - SpeechStartedEvent, - SpeechStoppedEvent, - TokenUsageEvent, - TranscriptionCompletedEvent, -) -from speech_to_speech.pipeline.messages import GenerateResponseRequest -from speech_to_speech.pipeline.queue_types import TextPromptItem -from speech_to_speech.pipeline.speculative_turns import SpeculativeTurnTracker -from speech_to_speech.utils.utils import _generate_id - -logger = logging.getLogger(__name__) - -PIPELINE_SAMPLE_RATE = 16000 -CHUNK_SAMPLES = 512 -BYTES_PER_SAMPLE = 2 -CHUNK_SIZE_BYTES = CHUNK_SAMPLES * BYTES_PER_SAMPLE - -_ResponseStatus = Literal["completed", "cancelled", "failed", "incomplete", "in_progress"] -_StatusReason = Literal["turn_detected", "client_cancelled", "max_output_tokens", "content_filter"] - -_EVENT_TYPE_TO_MODEL: dict[str, type[BaseModel]] = { - "input_audio_buffer.append": InputAudioBufferAppendEvent, - "input_audio_buffer.commit": InputAudioBufferCommitEvent, - "output_audio_buffer.clear": OutputAudioBufferClearEvent, - "session.update": SessionUpdateEvent, - "conversation.item.create": ConversationItemCreateEvent, - "response.create": ResponseCreateEvent, - "response.cancel": ResponseCancelEvent, -} - -ClientEvent = Union[ - InputAudioBufferAppendEvent, - InputAudioBufferCommitEvent, - OutputAudioBufferClearEvent, - SessionUpdateEvent, - ConversationItemCreateEvent, - ResponseCreateEvent, - ResponseCancelEvent, -] - -ServerEvent = Union[ - SessionCreatedEvent, - RealtimeErrorEvent, - InputAudioBufferSpeechStartedEvent, - InputAudioBufferSpeechStoppedEvent, - ConversationItemCreatedEvent, - ConversationItemInputAudioTranscriptionDeltaEvent, - ConversationItemInputAudioTranscriptionCompletedEvent, - ResponseCreatedEvent, - ResponseDoneEvent, - ResponseAudioDeltaEvent, - ResponseAudioDoneEvent, - ResponseAudioTranscriptDoneEvent, - ResponseFunctionCallArgumentsDoneEvent, - ResponseTextDeltaEvent, - ResponseTextDoneEvent, -] - -RealtimeEvent = Union[ClientEvent, ServerEvent] - - -_UsageMetricsT = TypeVar("_UsageMetricsT", bound="UsageMetrics") - - -class UsageMetrics(BaseModel): - """Per-response usage counters. - - Supports ``+=`` for rolling per-response metrics into a global total - and ``reset()`` for clearing per-response state after rollup. - """ - - input_tokens: int = 0 - output_tokens: int = 0 - audio_duration_s: float = 0.0 - responses_completed: int = 0 - responses_cancelled: int = 0 - tool_calls: int = 0 - turns: int = 0 - - def __iadd__(self: _UsageMetricsT, other: "UsageMetrics") -> _UsageMetricsT: - for field in UsageMetrics.model_fields: - setattr(self, field, getattr(self, field) + getattr(other, field)) - return self - - def reset(self) -> None: - for field, info in UsageMetrics.model_fields.items(): - setattr(self, field, info.default) - - -class GlobalUsageMetrics(UsageMetrics): - """Server-wide metrics that extend per-response counters with - connection and error tracking.""" - - connections: int = 0 - # connection duration in seconds. - # latency tts, llm, vad, stt (mean, max, p90) - errors_by_type: dict[str, int] = Field(default_factory=dict) - - def record_error(self, error_type: str) -> None: - self.errors_by_type[error_type] = self.errors_by_type.get(error_type, 0) + 1 - - @property - def total_errors(self) -> int: - return sum(self.errors_by_type.values()) - - -class ConnState(BaseModel): - """Per-connection mutable state, including all protocol-level IDs.""" - - model_config = ConfigDict(arbitrary_types_allowed=True) - - session_id: str = Field(default_factory=lambda: _generate_id("session")) - conversation_id: str = Field(default_factory=lambda: _generate_id("conv")) - runtime_config: RuntimeConfig = Field(default_factory=RuntimeConfig) - in_response: bool = False - response_pending: bool = False - audio_buffer_has_data: bool = False - audio_remainder: bytes = b"" - current_response_id: Optional[str] = None - current_item_id: Optional[str] = None - content_index: int = 0 - input_content_index: int = 0 - input_audio_duration_s: float = 0.0 - last_item_id: Optional[str] = None - current_response_params: RealtimeResponseCreateParams | None = None - pending_output_text_parts: list[str] = Field(default_factory=list) - response_usage: UsageMetrics = Field(default_factory=UsageMetrics) - speculative_turn_id: Optional[str] = None - speculative_turn_revision: Optional[int] = None - speculative_user_turn_id: Optional[str] = None - speculative_user_turn_revision: Optional[int] = None - speculative_user_speech_stopped_at_s: Optional[float] = None - speculative_user_item_id: Optional[str] = None - speculative_input_item_id: Optional[str] = None - speculative_audio_duration_s: float = 0.0 - # Client conversation.item.create items that arrived while a response was - # generating. Applying them mid-generation races the LLM handler's chat - # write-back (cross-thread), so they are buffered here and flushed in order - # once the response completes. See ConversationHandler.flush_deferred_items. - deferred_items: list[ConversationItem] = Field(default_factory=list) - - -class RealtimeService: - """Translates between OpenAI Realtime protocol events and internal pipeline messages. - - One instance is shared across all WebSocket connections. Per-connection - state (response lifecycle, audio buffer) is tracked internally by - connection id. - """ - - def __init__( - self, - text_prompt_queue: Queue[TextPromptItem] | None = None, - should_listen: ThreadingEvent | None = None, - chat_size: int = 10, - speculative_turns: SpeculativeTurnTracker | None = None, - ) -> None: - self.text_prompt_queue = text_prompt_queue - self.should_listen = should_listen - self._chat_size = chat_size - self.speculative_turns = speculative_turns - self._conns: dict[str, ConnState] = {} - self.total_usage = GlobalUsageMetrics() - - self.audio = AudioHandler(self) - self.session = SessionHandler(self) - self.response = ResponseHandler(self) - self.conversation = ConversationHandler(self) - - self._pipeline_dispatch: dict[type[PipelineEvent], Callable[..., list[ServerEvent]]] = { - SpeechStartedEvent: self.audio.on_speech_started, - SpeechStoppedEvent: self.audio.on_speech_stopped, - TokenUsageEvent: self._on_token_usage, - PartialTranscriptionEvent: self.conversation.on_partial_transcription, - TranscriptionCompletedEvent: self._on_transcription_completed, - AudioInputCompletedEvent: self._on_audio_input_completed, - ResponseFailedEvent: self._on_response_failed, - } - - # ── Connection lifecycle ───────────────────── - - def register(self) -> str: - """Register a new connection and return its session_id.""" - if self.speculative_turns: - self.speculative_turns.reset() - state = ConnState(runtime_config=RuntimeConfig(chat=Chat(self._chat_size))) - self._conns[state.session_id] = state - self.total_usage.connections += 1 - return state.session_id - - def unregister(self, conn_id: str) -> None: - st = self._conns.pop(conn_id, None) - if st is not None: - # Suppress any in-flight compaction splice so a daemon worker can't - # mutate a Chat tied to a closed session, and don't make further - # billable LLM calls on its behalf once the splice is suppressed. - st.runtime_config.chat.close() - self.total_usage += st.response_usage - logger.info( - "Session %s unregistered — cumulative: input_tokens=%d, output_tokens=%d, audio=%.2fs", - conn_id, - self.total_usage.input_tokens, - self.total_usage.output_tokens, - self.total_usage.audio_duration_s, - ) - - def _state(self, conn_id: str) -> ConnState: - return self._conns[conn_id] - - @property - def connection_ids(self) -> list[str]: - return list(self._conns) - - # ── Client event parsing ───────────────────── - - @staticmethod - def _next_event_id() -> str: - return _generate_id("event") - - def parse_client_event(self, raw: Mapping[str, object]) -> Optional[ClientEvent]: - raw_type = raw.get("type") - event_type: Optional[str] = raw_type if isinstance(raw_type, str) else None - if event_type is None: - logger.warning("Client event missing 'type' field") - return None - model_cls = _EVENT_TYPE_TO_MODEL.get(event_type) - if model_cls is None: - logger.warning(f"Unknown client event type: {event_type}") - return None - try: - return model_cls.model_validate(raw) # type: ignore[return-value] - except ValidationError as e: - logger.error(f"Invalid {event_type} payload: {e}") - return None - - # ── Client event handlers ──────────────────── - - def build_session_created(self, conn_id: str) -> SessionCreatedEvent: - return self.session.build_session_created(conn_id) - - def handle_session_update(self, conn_id: str, event: SessionUpdateEvent) -> Optional[RealtimeErrorEvent]: - return self.session.handle_session_update(conn_id, event) - - def handle_audio_append(self, conn_id: str, event: InputAudioBufferAppendEvent) -> list[bytes]: - return self.audio.handle_audio_append(conn_id, event) - - def append_pcm(self, conn_id: str, pcm_bytes: bytes, src_rate: int) -> list[bytes]: - return self.audio.append_pcm(conn_id, pcm_bytes, src_rate) - - def handle_audio_commit(self, conn_id: str) -> RealtimeErrorEvent | None: - return self.audio.handle_audio_commit(conn_id) - - def begin_audio_response(self, conn_id: str) -> tuple[str, str, list[ServerEvent]]: - return self.audio.begin_audio_response(conn_id) - - def encode_audio_chunk(self, conn_id: str, audio: bytes) -> list[ServerEvent]: - return self.audio.encode_audio_chunk(conn_id, audio) - - def handle_response_create(self, conn_id: str, event: ResponseCreateEvent) -> ServerEvent | None: - return self.response.handle_response_create(conn_id, event) - - def handle_response_cancel(self, conn_id: str) -> list[ServerEvent]: - return self.response.handle_response_cancel(conn_id) - - def finish_response( - self, - conn_id: str, - status: _ResponseStatus = "completed", - reason: _StatusReason | None = None, - ) -> list[ServerEvent]: - return self.response.finish_response(conn_id, status, reason) - - def handle_conversation_item_create(self, conn_id: str, event: ConversationItemCreateEvent) -> list[ServerEvent]: - return self.conversation.handle_conversation_item_create(conn_id, event) - - def dispatch_pipeline_event(self, conn_id: str, event: PipelineEvent) -> list[ServerEvent]: - """Route a pipeline text_output_queue event to the appropriate handler.""" - events = self._dispatch_pipeline_event(conn_id, event, wait_for_pending_reopen=True) - return [] if events is None else events - - def try_dispatch_pipeline_event(self, conn_id: str, event: PipelineEvent) -> list[ServerEvent] | None: - """Non-blocking dispatch. - - Returns ``None`` when dispatch must be retried after a speculative - reopen candidate resolves. - """ - return self._dispatch_pipeline_event(conn_id, event, wait_for_pending_reopen=False) - - def should_defer_pipeline_event(self, event: PipelineEvent) -> bool: - if self.speculative_turns is None or not isinstance(event, (AssistantTextEvent, TokenUsageEvent)): - return False - return self.speculative_turns.has_pending_reopen_or_grace( - getattr(event, "turn_id", None), - getattr(event, "turn_revision", None), - ) - - def _dispatch_pipeline_event( - self, - conn_id: str, - event: PipelineEvent, - *, - wait_for_pending_reopen: bool, - ) -> list[ServerEvent] | None: - is_stale = self._is_stale_turn_event(event, wait_for_pending_reopen=wait_for_pending_reopen) - if is_stale is None: - return None - if is_stale: - logger.info( - "Ignoring stale %s for turn=%s rev=%s", - event.type, - getattr(event, "turn_id", None), - getattr(event, "turn_revision", None), - ) - return [] - - self._observe_turn_event(event) - if isinstance(event, AssistantTextEvent): - return self.response.on_assistant_text( - conn_id, - event, - wait_for_pending_reopen=wait_for_pending_reopen, - ) - handler = self._pipeline_dispatch.get(type(event)) - if handler is None: - logger.debug("Unhandled pipeline event type: %s", type(event).__name__) - return [] - return handler(conn_id, event) - - def _is_stale_turn_event(self, event: PipelineEvent, *, wait_for_pending_reopen: bool = True) -> bool | None: - if self.speculative_turns is None: - return False - if not isinstance( - event, - ( - PartialTranscriptionEvent, - TranscriptionCompletedEvent, - AudioInputCompletedEvent, - AssistantTextEvent, - TokenUsageEvent, - ), - ): - return False - turn_id = getattr(event, "turn_id", None) - turn_revision = getattr(event, "turn_revision", None) - if isinstance(event, (AssistantTextEvent, TokenUsageEvent)): - is_latest: bool | None - if wait_for_pending_reopen: - is_latest = self.speculative_turns.is_latest_after_reopen_grace(turn_id, turn_revision) - else: - is_latest = self.speculative_turns.try_is_latest_after_reopen_grace(turn_id, turn_revision) - if is_latest is None: - return None - return not is_latest - return not self.speculative_turns.is_latest(turn_id, turn_revision) - - def _observe_turn_event(self, event: PipelineEvent) -> None: - if self.speculative_turns is None: - return - self.speculative_turns.observe( - getattr(event, "turn_id", None), - getattr(event, "turn_revision", None), - ) - - # ── STT → LM bridge ──────────────────────────── - - def _on_transcription_completed(self, conn_id: str, event: TranscriptionCompletedEvent) -> list[ServerEvent]: - """Handle a final STT transcription: emit protocol event, append to chat, trigger LM.""" - st = self._state(conn_id) - same_speculative_turn = event.turn_id is not None and event.turn_id == st.speculative_user_turn_id - if same_speculative_turn: - st.response_usage.audio_duration_s -= st.speculative_audio_duration_s - else: - st.speculative_audio_duration_s = 0.0 - - events = self.conversation.on_transcription_completed(conn_id, event) - if event.turn_id is not None: - st.speculative_audio_duration_s = st.input_audio_duration_s - - cfg = st.runtime_config - transcript = event.transcript - if transcript: - if same_speculative_turn and st.speculative_user_item_id: - replaced = cfg.chat.replace_user_message_text(st.speculative_user_item_id, transcript) - if not replaced: - item = cfg.chat.add_item(make_user_message(transcript)) - st.speculative_user_item_id = item.id - else: - item = cfg.chat.add_item(make_user_message(transcript)) - st.speculative_user_item_id = item.id - elif same_speculative_turn and st.speculative_user_item_id: - cfg.chat.remove_user_message(st.speculative_user_item_id) - st.speculative_user_item_id = None - elif event.turn_id is not None and event.turn_id != st.speculative_user_turn_id: - st.speculative_user_item_id = None - - if event.turn_id is not None: - st.speculative_user_turn_id = event.turn_id - st.speculative_user_turn_revision = event.turn_revision - st.speculative_user_speech_stopped_at_s = event.speech_stopped_at_s - - queue = self.text_prompt_queue - if queue and transcript: - st.response_pending = True - queue.put( - GenerateResponseRequest( - runtime_config=cfg, - language_code=event.language_code, - turn_id=event.turn_id, - turn_revision=event.turn_revision, - speech_stopped_at_s=event.speech_stopped_at_s, - ) - ) - - return events - - def _on_audio_input_completed(self, conn_id: str, event: AudioInputCompletedEvent) -> list[ServerEvent]: - """Record final input audio and queue its realtime LM request.""" - st = self._state(conn_id) - same_speculative_turn = event.turn_id is not None and event.turn_id == st.speculative_user_turn_id - if same_speculative_turn: - st.response_usage.audio_duration_s -= st.speculative_audio_duration_s - else: - st.speculative_audio_duration_s = 0.0 - - st.input_audio_duration_s = event.audio_duration_s - st.response_usage.audio_duration_s += event.audio_duration_s - if event.turn_id is not None: - st.speculative_audio_duration_s = event.audio_duration_s - st.speculative_user_turn_id = event.turn_id - st.speculative_user_turn_revision = event.turn_revision - st.speculative_user_speech_stopped_at_s = event.speech_stopped_at_s - - queue = self.text_prompt_queue - if queue: - st.response_pending = True - queue.put( - GenerateResponseRequest( - runtime_config=st.runtime_config, - audio=event.audio, - audio_sample_rate=event.audio_sample_rate, - turn_id=event.turn_id, - turn_revision=event.turn_revision, - speech_stopped_at_s=event.speech_stopped_at_s, - ) - ) - return [] - - # ── Metrics ──────────────────────────────────── - - def _on_token_usage(self, conn_id: str, event: TokenUsageEvent) -> list[ServerEvent]: - """Accumulate input/output token counts on the connection's usage metrics.""" - if self.speculative_turns and not self.speculative_turns.is_latest( - event.turn_id, - event.turn_revision, - ): - logger.debug("Dropping stale token usage for turn=%s rev=%s", event.turn_id, event.turn_revision) - return [] - st = self._state(conn_id) - st.response_usage.input_tokens += event.input_tokens - st.response_usage.output_tokens += event.output_tokens - logger.info( - "Token usage (response): input=%d, output=%d", - st.response_usage.input_tokens, - st.response_usage.output_tokens, - ) - return [] - - def _on_response_failed(self, conn_id: str, event: ResponseFailedEvent) -> list[ServerEvent]: - """Surface the failure to the client and close the response as ``failed``. - - Emitted when generation failed (e.g. invalid out-of-band input, or the - provider rejecting an empty context). A top-level ``error`` event carries - the human-readable reason — ``response.done.status_details.error`` only - has code/type, no message — then ``finish_response`` closes the slot. - - Idempotent: gated on an active or pending response. Pending implicit - responses are announced and promoted to active so ``finish_response`` - can emit terminal events for a response the client already knows about. - Once closed, a later EndOfResponse-driven close does nothing. - """ - logger.info("Response failed: %s", event.message) - st = self._state(conn_id) - if not (st.in_response or st.response_pending): - return [] - events: list[ServerEvent] = [] - if st.response_pending: - _, _, created_events = self.audio.begin_audio_response(conn_id) - events.extend(created_events) - events.append(self.make_error(event.message, "response_failed")) - events.extend(self.response.finish_response(conn_id, status="failed")) - return events - - def get_usage(self) -> dict[str, Any]: - """Return cumulative usage metrics across all completed responses.""" - data = self.total_usage.model_dump() - data["total_tokens"] = data["input_tokens"] + data["output_tokens"] - data["total_errors"] = self.total_usage.total_errors - return data - - # ── Error ─────────────────────────────────── - - def make_error(self, message: str, _type: str) -> RealtimeErrorEvent: - self.total_usage.record_error(_type) - return build_error_event(message, _type) - - -def build_error_event(message: str, error_type: str) -> RealtimeErrorEvent: - """Construct a RealtimeErrorEvent without touching any service-instance state. - - Used by the websocket route handler on pool rejection, where no unit's - service should be charged with the error in its usage metrics. - """ - return RealtimeErrorEvent( - type="error", - error=RealtimeError(message=message, type=error_type), - event_id=_generate_id("event"), - ) diff --git a/runtime/src/speech_to_speech/api/openai_realtime/transports.py b/runtime/src/speech_to_speech/api/openai_realtime/transports.py deleted file mode 100644 index 95b6d924907d406bee0bc8386d26fdc8aa695021..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/api/openai_realtime/transports.py +++ /dev/null @@ -1,100 +0,0 @@ -"""Session transports: how server events and outbound audio reach a client. - -The per-unit send loop in ``websocket_router`` is transport-agnostic: it owns -the pipeline output queues (sentinels, generation discards, SESSION_END drain) -and hands client-visible traffic to the transport attached to the current -``SessionState``. Two implementations exist: - -- ``WebSocketTransport`` (here): events and base64 audio deltas as JSON frames. -- ``WebRTCSession`` (in ``webrtc_session``, requires the ``webrtc`` extra): - events over the ``oai-events`` data channel, audio over the RTP media track. -""" - -from __future__ import annotations - -import logging -from abc import ABC, abstractmethod -from typing import TYPE_CHECKING - -from fastapi import WebSocket, WebSocketDisconnect -from starlette.websockets import WebSocketState - -if TYPE_CHECKING: - from speech_to_speech.api.openai_realtime.service import RealtimeService, ServerEvent - -logger = logging.getLogger(__name__) - - -class SessionTransport(ABC): - """What the send loop and client-event dispatch need from a transport.""" - - kind: str - - @abstractmethod - async def send_events(self, events: list[ServerEvent]) -> None: ... - - @abstractmethod - async def send_audio_chunk(self, service: RealtimeService, session_id: str, pcm: bytes) -> None: - """Deliver a pipeline-rate PCM16 chunk to the client.""" - - @abstractmethod - def discard_pending_audio(self) -> None: - """Drop transport-buffered audio that has not reached the client yet. - - WebSocket clients buffer audio on their side, so this is a no-op there; - the WebRTC transport paces playback server-side and must flush its - track buffer for barge-in to actually silence the assistant. - """ - - @abstractmethod - async def close(self) -> None: ... - - -async def send_ws_event(ws: WebSocket, event: ServerEvent) -> None: - # Skip cleanly when the ws is already closing/closed — happens during Ctrl-C - # shutdown, where the lifespan starts closing sockets while the route handler - # or send loop is still in flight pushing events. - if ws.application_state != WebSocketState.CONNECTED: - return - try: - await ws.send_json(event.model_dump()) - except WebSocketDisconnect: - logger.debug("Skipped event: ws disconnected mid-send") - except RuntimeError as e: - # Race: ws closed between the state check above and the send. Starlette - # raises a plain RuntimeError("Unexpected ASGI message 'websocket.send' - # after sending 'websocket.close' ...") — harmless during shutdown. - msg = str(e) - if "websocket.close" in msg or "websocket.disconnect" in msg or "response already completed" in msg: - logger.debug(f"Skipped event: ws already closed ({msg})") - else: - logger.error(f"Failed to send event to client: {e}") - except Exception as e: # noqa: BLE001 - logger.error(f"Failed to send event to client: {e}") - - -class WebSocketTransport(SessionTransport): - """JSON-over-WebSocket transport: audio is sent as base64 delta events.""" - - kind = "websocket" - - def __init__(self, websocket: WebSocket) -> None: - self.websocket = websocket - - async def send_events(self, events: list[ServerEvent]) -> None: - for event in events: - await send_ws_event(self.websocket, event) - - async def send_audio_chunk(self, service: RealtimeService, session_id: str, pcm: bytes) -> None: - await self.send_events(service.encode_audio_chunk(session_id, pcm)) - - def discard_pending_audio(self) -> None: - # Unplayed audio lives client-side over WebSocket; truncation is the - # client's responsibility. - pass - - async def close(self) -> None: - try: - await self.websocket.close() - except Exception: # noqa: BLE001 - pass diff --git a/runtime/src/speech_to_speech/api/openai_realtime/utils.py b/runtime/src/speech_to_speech/api/openai_realtime/utils.py deleted file mode 100644 index 369b4e1a4f70763a52e9577971d9a9943e19dede..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/api/openai_realtime/utils.py +++ /dev/null @@ -1,12 +0,0 @@ -import numpy as np -from scipy.signal import resample_poly - - -def resample(audio_int16: bytes, from_rate: int, to_rate: int) -> bytes: - """Resample int16 PCM audio between sample rates using polyphase filtering.""" - if from_rate == to_rate: - return audio_int16 - samples = np.frombuffer(audio_int16, dtype=np.int16).astype(np.float32) / 32768.0 - gcd = np.gcd(to_rate, from_rate) - resampled = resample_poly(samples, up=to_rate // gcd, down=from_rate // gcd) - return np.clip(resampled * 32768, -32768, 32767).astype(np.int16).tobytes() diff --git a/runtime/src/speech_to_speech/api/openai_realtime/webrtc_session.py b/runtime/src/speech_to_speech/api/openai_realtime/webrtc_session.py deleted file mode 100644 index 2d4d0c200d4f6308b56697dc70177b20d7c5d541..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/api/openai_realtime/webrtc_session.py +++ /dev/null @@ -1,351 +0,0 @@ -"""WebRTC transport for the OpenAI Realtime API emulation. - -Requires the ``webrtc`` extra (aiortc). Audio travels over RTP media tracks -(Opus at 48 kHz, resampled to/from the 16 kHz pipeline rate); all JSON events -use the same protocol as the WebSocket transport, carried on the -``oai-events`` data channel. - -``WebRTCSession`` subclasses ``SessionTransport`` from ``transports``, -so the per-unit send loop in ``websocket_router`` drives it -exactly like a WebSocket session: it stays the sole consumer of the pipeline -output queues, and this module only turns delivered PCM into paced RTP frames. -""" - -from __future__ import annotations - -import asyncio -import json -import logging -import os -import time -from collections.abc import Awaitable -from fractions import Fraction -from typing import TYPE_CHECKING, Callable, Optional - -import av -import numpy as np -from aiortc import RTCConfiguration, RTCIceServer, RTCPeerConnection, RTCSessionDescription -from aiortc.mediastreams import MediaStreamError, MediaStreamTrack - -from speech_to_speech.api.openai_realtime.service import PIPELINE_SAMPLE_RATE -from speech_to_speech.api.openai_realtime.transports import SessionTransport - -if TYPE_CHECKING: - from speech_to_speech.api.openai_realtime.service import RealtimeService, ServerEvent - -logger = logging.getLogger(__name__) - -WEBRTC_SAMPLE_RATE = 48_000 -AUDIO_PTIME = 0.02 # 20 ms frames -WEBRTC_FRAME_SAMPLES = int(WEBRTC_SAMPLE_RATE * AUDIO_PTIME) -DATA_CHANNEL_LABEL = "oai-events" -ICE_SERVERS_ENV = "SPEECH_TO_SPEECH_ICE_SERVERS" -ICE_GATHERING_TIMEOUT_S = 5.0 -# How long a negotiated session may sit without the peer connection reaching -# "connected" before we release its pipeline unit. Without this, a client that -# receives the SDP answer and never completes ICE would hold the unit forever. -CONNECT_TIMEOUT_S = 30.0 - - -def rtc_configuration_from_env() -> Optional[RTCConfiguration]: - """Build an RTCConfiguration from the SPEECH_TO_SPEECH_ICE_SERVERS env var. - - The variable holds a JSON list of RTCIceServer kwargs, e.g. - ``[{"urls": "stun:stun.example.com:3478"}, - {"urls": "turn:turn.example.com", "username": "u", "credential": "c"}]``. - Returns None (aiortc defaults) when unset or invalid. - """ - raw = os.environ.get(ICE_SERVERS_ENV) - if not raw: - return None - try: - entries = json.loads(raw) - servers = [RTCIceServer(**entry) for entry in entries] - except (json.JSONDecodeError, TypeError) as e: - logger.error(f"Ignoring invalid {ICE_SERVERS_ENV}: {e}") - return None - return RTCConfiguration(iceServers=servers) - - -class PcmResampler: - """Stateful mono/s16 resampler around av.AudioResampler. - - One instance per direction per session: the libswresample filter state - carries across calls, so 20 ms frames resample without the boundary - artifacts a stateless per-chunk resample would introduce. Also downmixes - multi-channel input (browser Opus is typically stereo) to mono. - """ - - def __init__(self, target_rate: int) -> None: - self._resampler = av.AudioResampler(format="s16", layout="mono", rate=target_rate) - self._pts = 0 - - def resample_frame(self, frame: av.AudioFrame) -> bytes: - out = bytearray() - for resampled in self._resampler.resample(frame): - out += resampled.to_ndarray().tobytes() - return bytes(out) - - def resample_pcm(self, pcm: bytes, src_rate: int) -> bytes: - samples = np.frombuffer(pcm, dtype=np.int16) - frame = av.AudioFrame.from_ndarray(samples[np.newaxis, :], format="s16", layout="mono") - frame.sample_rate = src_rate - frame.pts = self._pts - frame.time_base = Fraction(1, src_rate) - self._pts += samples.shape[0] - return self.resample_frame(frame) - - -class PipelineAudioTrack(MediaStreamTrack): - """Outbound audio track: paced 20 ms 48 kHz frames from a PCM buffer. - - The send loop pushes generated audio in via ``write()`` (faster than - real time); ``recv()`` paces delivery against the wall clock like - aiortc's built-in AudioStreamTrack, emitting silence when the buffer is - empty so the RTP stream stays continuous. ``clear()`` drops unplayed - audio — this is the server-side equivalent of the client's speaker - buffer, so barge-in must flush it for interruption to be audible. - """ - - kind = "audio" - - def __init__(self) -> None: - super().__init__() - self._buffer = bytearray() - self._start: Optional[float] = None - self._timestamp = 0 - - def write(self, pcm: bytes) -> None: - self._buffer.extend(pcm) - - def clear(self) -> None: - del self._buffer[:] - - @property - def buffered_bytes(self) -> int: - return len(self._buffer) - - async def recv(self) -> av.AudioFrame: - if self.readyState != "live": - raise MediaStreamError - - if self._start is None: - self._start = time.time() - self._timestamp = 0 - else: - self._timestamp += WEBRTC_FRAME_SAMPLES - wait = self._start + (self._timestamp / WEBRTC_SAMPLE_RATE) - time.time() - if wait > 0: - await asyncio.sleep(wait) - - needed = WEBRTC_FRAME_SAMPLES * 2 # bytes of s16 mono - payload = bytes(self._buffer[:needed]) - del self._buffer[: len(payload)] - if len(payload) < needed: - payload += b"\x00" * (needed - len(payload)) - - samples = np.frombuffer(payload, dtype=np.int16) - frame = av.AudioFrame.from_ndarray(samples[np.newaxis, :], format="s16", layout="mono") - frame.sample_rate = WEBRTC_SAMPLE_RATE - frame.pts = self._timestamp - frame.time_base = Fraction(1, WEBRTC_SAMPLE_RATE) - return frame - - -class WebRTCSession(SessionTransport): - """One WebRTC peer connection, used as the SessionState transport. - - All pipeline integration arrives through callbacks supplied by the route - handler (which owns the PipelineUnit): parsed client events, inbound PCM, - channel-open, and close. This module never touches queues or services - directly except through the SessionTransport methods the send loop calls. - """ - - kind = "webrtc" - - def __init__( - self, - pc: RTCPeerConnection, - *, - on_client_event: Callable[[dict], Awaitable[None]], - on_audio: Callable[[bytes], None], - on_open: Callable[[], Awaitable[None]], - on_closed: Callable[[], None], - ) -> None: - self._pc = pc - self._on_client_event = on_client_event - self._on_audio = on_audio - self._on_open = on_open - self._on_closed = on_closed - self._dc = None - self._closed = False - self._track = PipelineAudioTrack() - self._out_resampler = PcmResampler(WEBRTC_SAMPLE_RATE) - self._in_resampler = PcmResampler(PIPELINE_SAMPLE_RATE) - # Data-channel messages funnel through one queue + consumer task so - # client events apply in arrival order; dispatching each message in - # its own task could reorder e.g. session.update vs response.create. - self._dc_messages: asyncio.Queue[str] = asyncio.Queue() - self._tasks: list[asyncio.Task] = [] - - # ── Lifecycle ───────────────────────────────── - - def setup(self) -> None: - """Wire aiortc event callbacks. Call before negotiate().""" - self._pc.addTrack(self._track) - - @self._pc.on("datachannel") - def on_datachannel(dc) -> None: - if dc.label != DATA_CHANNEL_LABEL: - logger.warning(f"[WebRTC] Ignoring unexpected data channel: {dc.label}") - return - self._dc = dc - self._spawn(self._consume_dc_messages()) - logger.info(f"[WebRTC] Data channel '{DATA_CHANNEL_LABEL}' received") - - # aiortc may deliver the channel already open, in which case the - # "open" event never fires. - if dc.readyState == "open": - self._spawn(self._on_open()) - else: - - @dc.on("open") - def on_dc_open() -> None: - self._spawn(self._on_open()) - - @dc.on("message") - def on_message(msg) -> None: - if isinstance(msg, str): - self._dc_messages.put_nowait(msg) - else: - logger.warning("[WebRTC] Ignoring binary data-channel message") - - @dc.on("close") - def on_dc_close() -> None: - logger.info("[WebRTC] Data channel closed") - self._spawn(self.close()) - - @self._pc.on("track") - def on_track(track) -> None: - if track.kind == "audio": - logger.info("[WebRTC] Inbound audio track received") - self._spawn(self._consume_inbound_audio(track)) - - @self._pc.on("connectionstatechange") - async def on_connection_state_change() -> None: - state = self._pc.connectionState - logger.info(f"[WebRTC] Connection state: {state}") - if state in ("failed", "closed"): - await self.close() - - async def negotiate(self, offer_sdp: str) -> str: - """Apply the client's SDP offer and return the SDP answer. - - Waits for ICE gathering so the answer carries the server's candidates - — there is no trickle-ICE channel in the HTTP handshake. - """ - await self._pc.setRemoteDescription(RTCSessionDescription(sdp=offer_sdp, type="offer")) - answer = await self._pc.createAnswer() - await self._pc.setLocalDescription(answer) - - if self._pc.iceGatheringState != "complete": - done: asyncio.Event = asyncio.Event() - - @self._pc.on("icegatheringstatechange") - def on_ice_change() -> None: - if self._pc.iceGatheringState == "complete": - done.set() - - if self._pc.iceGatheringState == "complete": # raced to completion - done.set() - try: - await asyncio.wait_for(done.wait(), timeout=ICE_GATHERING_TIMEOUT_S) - except asyncio.TimeoutError: - logger.warning("[WebRTC] ICE gathering timed out, returning partial SDP") - - self._spawn(self._connect_watchdog()) - return self._pc.localDescription.sdp - - async def close(self) -> None: - if self._closed: - return - self._closed = True - # close() itself often runs as a _spawn()ed task (dc close handler), - # so it is in _tasks — cancelling the current task here would abort - # this method before the release callback runs. - current = asyncio.current_task() - for task in self._tasks: - if task is not current: - task.cancel() - self._track.stop() - try: - await self._pc.close() - except Exception as e: # noqa: BLE001 - logger.warning(f"[WebRTC] Error closing peer connection: {e}") - self._on_closed() - logger.info("[WebRTC] Session closed") - - # ── SessionTransport interface ──────────────── - - async def send_events(self, events: list[ServerEvent]) -> None: - dc = self._dc - if dc is None or dc.readyState != "open": - return - for event in events: - try: - dc.send(json.dumps(event.model_dump())) - except Exception as e: # noqa: BLE001 - logger.error(f"[WebRTC] Data channel send error: {e}") - - async def send_audio_chunk(self, service: RealtimeService, session_id: str, pcm: bytes) -> None: - # Bookkeeping events (response.created on the implicit VAD path) go - # over the data channel; the audio itself goes on the media track. - _resp_id, _item_id, events = service.begin_audio_response(session_id) - if events: - await self.send_events(events) - self._track.write(self._out_resampler.resample_pcm(pcm, PIPELINE_SAMPLE_RATE)) - - def discard_pending_audio(self) -> None: - self._track.clear() - - # ── Internals ───────────────────────────────── - - def _spawn(self, coro: Awaitable[None]) -> None: - task = asyncio.ensure_future(coro) - self._tasks.append(task) - - async def _connect_watchdog(self) -> None: - await asyncio.sleep(CONNECT_TIMEOUT_S) - if not self._closed and self._pc.connectionState != "connected": - logger.warning( - f"[WebRTC] Peer not connected after {CONNECT_TIMEOUT_S:.0f}s " - f"(state: {self._pc.connectionState}); releasing session" - ) - await self.close() - - async def _consume_dc_messages(self) -> None: - while not self._closed: - msg = await self._dc_messages.get() - try: - raw = json.loads(msg) - except json.JSONDecodeError: - logger.error(f"[WebRTC] Invalid JSON on data channel: {msg!r}") - continue - if not isinstance(raw, dict): - logger.error(f"[WebRTC] Non-object event on data channel: {msg!r}") - continue - try: - await self._on_client_event(raw) - except Exception: # noqa: BLE001 - logger.exception("[WebRTC] Error handling client event") - - async def _consume_inbound_audio(self, track) -> None: - while not self._closed: - try: - frame = await track.recv() - except MediaStreamError: - logger.info("[WebRTC] Inbound audio track ended") - break - pcm = self._in_resampler.resample_frame(frame) - if pcm: - self._on_audio(pcm) diff --git a/runtime/src/speech_to_speech/api/openai_realtime/websocket_router.py b/runtime/src/speech_to_speech/api/openai_realtime/websocket_router.py deleted file mode 100644 index 7b9e538100c4594ba26207e43c8f1a7100710f43..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/api/openai_realtime/websocket_router.py +++ /dev/null @@ -1,878 +0,0 @@ -import asyncio -import logging -import time -from collections.abc import AsyncIterator -from contextlib import asynccontextmanager -from queue import Empty, Queue -from threading import Event as ThreadingEvent -from typing import Any, Callable, TypeVar - -import numpy as np -from fastapi import FastAPI, Request, Response, WebSocket, WebSocketDisconnect -from openai.types.realtime import ( - ConversationItemCreateEvent, - InputAudioBufferAppendEvent, - InputAudioBufferCommitEvent, - OutputAudioBufferClearEvent, - ResponseCancelEvent, - ResponseCreateEvent, - SessionUpdateEvent, -) - -from speech_to_speech.api.openai_realtime.pipeline_unit import PipelineUnit, SessionState -from speech_to_speech.api.openai_realtime.service import ( - PIPELINE_SAMPLE_RATE, - build_error_event, -) -from speech_to_speech.api.openai_realtime.transports import ( - SessionTransport, - WebSocketTransport, - send_ws_event, -) -from speech_to_speech.pipeline.control import SESSION_END, PipelineControlMessage, is_control_message -from speech_to_speech.pipeline.events import ( - AssistantTextEvent, - AudioInputCompletedEvent, - PartialTranscriptionEvent, - PipelineEvent, - SpeechStartedEvent, - SpeechStoppedEvent, - TokenUsageEvent, - TranscriptionCompletedEvent, -) -from speech_to_speech.pipeline.log_context import pipeline_log_ctx -from speech_to_speech.pipeline.messages import AUDIO_RESPONSE_DONE, PIPELINE_END, AudioOutput - -# aiortc (the 'webrtc' extra) is optional. Import it here, at module load, -# rather than lazily in the calls endpoint: the av/cryptography C extensions -# take up to a second to load cold, which would block the shared event loop — -# and every live conversation's audio — on the first WebRTC handshake. -try: - from aiortc import RTCPeerConnection - - from speech_to_speech.api.openai_realtime.webrtc_session import ( - WebRTCSession, - rtc_configuration_from_env, - ) - - WEBRTC_AVAILABLE = True -except ImportError: - WEBRTC_AVAILABLE = False - -logger = logging.getLogger(__name__) -MAX_AUDIO_BATCH_BYTES = 6400 -# How long the release path waits for SESSION_END to propagate through the -# handler chain back to output_queue before warning that the unit is stuck. -# Tests monkeypatch this to a small value since their fixtures usually skip -# the real handler chain. -SESSION_END_DRAIN_TIMEOUT_S = 10.0 -# Past this, the unit is quarantined: its service session is unregistered -# (closing the chat so late handler output can't mutate or bill it), but the -# unit stays unclaimable. Releasing it instead would let a new client claim a -# unit whose handlers may still emit the previous session's output (e.g. a -# transcript, which carries no session identity and would be appended to the -# new session's conversation) — a cross-session leak. If SESSION_END does -# eventually drain, the chain has proven itself clean and the unit returns to -# the pool; a dead handler keeps it quarantined forever, visible in /v1/pool -# as "stuck". -SESSION_END_QUARANTINE_TIMEOUT_S = 180.0 -QItem = TypeVar("QItem") - - -def _keep_audio_sentinel(item: Any) -> bool: - # SESSION_END must survive barge-in flushes of output_queue: dropping it - # would leave the release path waiting forever for the drain signal. - return _is_audio_done(item) or is_control_message(item, SESSION_END.kind) - - -def _keep_user_text_event(item: Any) -> bool: - return isinstance( - item, - ( - SpeechStoppedEvent, - PartialTranscriptionEvent, - TranscriptionCompletedEvent, - AudioInputCompletedEvent, - TokenUsageEvent, - ), - ) - - -def _audio_payload(item: Any) -> Any: - return item.audio if isinstance(item, AudioOutput) else item - - -def _audio_generation(item: Any) -> int | None: - return item.cancel_generation if isinstance(item, AudioOutput) else None - - -def _flush_queue(q: Queue[QItem], *, preserve: Callable[[QItem], bool] | None = None) -> None: - """Drain a queue, optionally preserving items matching *preserve*. - - Preserved items are re-inserted at the **front** of the queue - (atomically under the queue's mutex) so they are processed before - anything a pipeline thread may have enqueued during the drain. - """ - preserved: list[QItem] = [] - while True: - try: - item = q.get_nowait() - if preserve and preserve(item): - preserved.append(item) - except Empty: - break - if preserved: - with q.mutex: - for item in reversed(preserved): - q.queue.appendleft(item) - q.not_empty.notify(len(preserved)) - - -async def _drain_pending_response_events( - transport: SessionTransport | None, - unit: PipelineUnit, - session_id: str | None, -) -> None: - if session_id is None: - return - - preserved: list[Any] = [] - drained_assistant = 0 - drained_usage = 0 - drain_assistant_events = True - try: - while True: - try: - item = unit.text_output_queue.get_nowait() - except Empty: - break - # Usage is accounting-only, so keep the old whole-queue drain behavior. - # Assistant events are client-visible response output and stop at the - # first non-response boundary to preserve normal text-event ordering. - if isinstance(item, TokenUsageEvent): - unit.service.dispatch_pipeline_event(session_id, item) - drained_usage += 1 - elif drain_assistant_events and isinstance(item, AssistantTextEvent): - drained_assistant += 1 - if _generation_is_discardable(unit, item.cancel_generation): - continue - events = unit.service.dispatch_pipeline_event(session_id, item) - if transport is not None and events: - await transport.send_events(events) - else: - preserved.append(item) - drain_assistant_events = False - finally: - if preserved: - with unit.text_output_queue.mutex: - for item in reversed(preserved): - unit.text_output_queue.queue.appendleft(item) - unit.text_output_queue.not_empty.notify(len(preserved)) - - if drained_assistant or drained_usage: - logger.debug( - "Pipeline %d: drained %d assistant event(s) and %d token usage event(s) before response completion", - unit.index, - drained_assistant, - drained_usage, - ) - - -def _clean_unit(unit: PipelineUnit, preserve: Callable[[Any], bool] | None = None) -> None: - """Cancel in-flight work and flush queues for a single pipeline unit. - - All four pipeline queues are drained — input audio, transcript-to-LM, - LM-to-TTS output, and the text-event side channel — so pending work from - a released session cannot be picked up by handlers and leak into the next - session that claims this unit. SESSION_END is enqueued by the route - handler *after* this returns to serve as the soft reset signal for - stateful handlers. - """ - unit.cancel_scope.cancel() - _flush_queue(unit.input_queue) - _flush_queue(unit.text_prompt_queue) - _flush_queue(unit.output_queue, preserve=preserve) - _flush_queue(unit.text_output_queue, preserve=preserve) - unit.response_playing.clear() - unit.cancel_scope.reset() - unit.should_listen.set() - - -def _to_audio_bytes(chunk: Any) -> bytes: - chunk = _audio_payload(chunk) - if isinstance(chunk, PipelineControlMessage): - raise TypeError(f"unexpected control message on audio output queue: {chunk!r}") - if isinstance(chunk, np.ndarray) or hasattr(chunk, "tobytes"): - return chunk.tobytes() - return chunk - - -def _is_audio_done(item: Any) -> bool: - payload = _audio_payload(item) - return isinstance(payload, bytes) and payload == AUDIO_RESPONSE_DONE - - -def _is_pipeline_end(item: Any) -> bool: - payload = _audio_payload(item) - return isinstance(payload, bytes) and payload == PIPELINE_END - - -def _generation_is_discardable(unit: PipelineUnit, generation: int | None) -> bool: - """Whether output tagged with *generation* should be dropped. - - A generation is discardable if it has been superseded (``is_stale``) or if the - cancel scope is in its post-cancel discard window and this is not the current - live generation. Shared by audio and assistant-text so the two paths stay in - lockstep: dropping text whenever ``discarding`` is set (without this generation - check) silently swallows the transcript of a fresh response when ``discarding`` - lingers — e.g. a superseded speculative turn whose TTS never emitted an - AUDIO_RESPONSE_DONE sentinel, so response_done() never cleared the flag. - """ - if generation is not None and unit.cancel_scope.is_stale(generation): - return True - if unit.cancel_scope.discarding and generation != unit.cancel_scope.generation: - return True - return False - - -def _should_discard_audio(unit: PipelineUnit, item: Any) -> bool: - return _generation_is_discardable(unit, _audio_generation(item)) - - -def _safe_unregister(unit: PipelineUnit, session_id: str) -> None: - try: - unit.service.unregister(session_id) - except Exception: - logger.exception(f"Pipeline {unit.index}: unregister failed for session {session_id}") - - -async def _release_unit_after_drain(unit: PipelineUnit, session: Any, session_id: str) -> None: - """Wait for SESSION_END to propagate, then release the unit. - - Runs in its own asyncio task so the route handler's finally block can return - immediately. The unit stays unavailable for new claims (unit.session != None) - until SESSION_END travels all the way through the handler chain back to - output_queue — observed by the send loop, which sets session.drained. - - Past SESSION_END_QUARANTINE_TIMEOUT_S (a wedged or dead handler thread) the - unit is quarantined, NOT released: still-running handlers could emit the old - session's output (transcripts carry no session identity) into whichever - session claimed the unit next, and a dead handler would make the unit accept - clients it can never serve. The session is unregistered right away so late - output can't mutate or bill the closed conversation; the unit itself only - returns to the pool if SESSION_END eventually drains, proving the chain is - clean. Operators can spot quarantined units in `/v1/pool` (state "stuck"). - """ - elapsed = 0.0 - warned = False - try: - while not session.drained.is_set(): - await asyncio.sleep(0.05) - elapsed += 0.05 - if not warned and elapsed >= SESSION_END_DRAIN_TIMEOUT_S: - logger.warning( - f"Pipeline {unit.index}: SESSION_END not drained after {elapsed:.1f}s — " - f"unit will remain unavailable until handlers finish (session {session_id})" - ) - warned = True - if session.quarantined_at is None and elapsed >= SESSION_END_QUARANTINE_TIMEOUT_S: - session.quarantined_at = time.monotonic() - _safe_unregister(unit, session_id) - logger.error( - f"Pipeline {unit.index}: SESSION_END still not drained after {elapsed:.0f}s — " - f"quarantining unit until the handler chain drains (session {session_id})" - ) - finally: - # Runs when the drain completed (chain proven clean) or the task is - # cancelled at shutdown. Release unconditionally: even if unregister - # raises, the unit must not stay claimed forever. - try: - _safe_unregister(unit, session_id) - finally: - unit.session = None - recovered = " after quarantine" if session.quarantined_at is not None else "" - logger.info(f"Pipeline {unit.index} released{recovered} (session {session_id} ended)") - - -# Strong references to in-flight drain-and-release tasks (asyncio only -# holds tasks weakly); each task removes itself on completion. -_release_tasks: set[asyncio.Task[None]] = set() - - -def _release_session(unit: PipelineUnit, session_id: str) -> None: - """Start the release of a unit after its client disconnected. - - Shared by the WebSocket route's finally block and the WebRTC session's - close callback. Marks the session as released, resets the unit, enqueues - SESSION_END, and spawns the drain-and-release task — the unit stays - claimed until SESSION_END propagates back to output_queue. - """ - old_session = unit.session - if old_session is None: - # Already released (e.g. duplicate close callbacks racing). - return - old_session.released_at = time.monotonic() - _clean_unit(unit) - # Tag SESSION_END with this session's id so that, after a force - # release, a late arrival can't satisfy the next session's drain. - unit.input_queue.put(PipelineControlMessage(SESSION_END.kind, session_id=session_id)) - task = asyncio.create_task(_release_unit_after_drain(unit, old_session, session_id)) - _release_tasks.add(task) - task.add_done_callback(_release_tasks.discard) - - -async def _dispatch_client_event( - unit: PipelineUnit, - session_id: str, - raw: dict[str, Any], - transport: SessionTransport, - *, - transport_kind: str = "websocket", -) -> None: - """Parse and apply one client event, replying over *transport*. - - Shared by both transports; ``transport_kind`` gates the events whose - validity depends on how audio travels: ``input_audio_buffer.append`` is - WebSocket-only (WebRTC audio arrives on the media track), and - ``output_audio_buffer.clear`` is WebRTC-only (over WebSocket the unplayed - audio sits client-side). - """ - service = unit.service - event = service.parse_client_event(raw) - if event is None: - await transport.send_events( - [service.make_error(f"Unknown or invalid event: {raw.get('type')}", "unknown_or_invalid_event")] - ) - return - - if isinstance(event, InputAudioBufferAppendEvent): - if transport_kind == "webrtc": - await transport.send_events( - [ - service.make_error( - "In WebRTC mode audio arrives via the media track; input_audio_buffer.append is not supported.", - "invalid_event_for_transport", - ) - ] - ) - return - chunks = service.handle_audio_append(session_id, event) - rt_cfg = service._state(session_id).runtime_config - for chunk in chunks: - unit.input_queue.put((chunk, rt_cfg)) - - elif isinstance(event, InputAudioBufferCommitEvent): - err = service.handle_audio_commit(session_id) - if err: - await transport.send_events([err]) - - elif isinstance(event, OutputAudioBufferClearEvent): - if transport_kind != "webrtc": - await transport.send_events( - [ - service.make_error( - "output_audio_buffer.clear is only supported on the WebRTC transport.", - "invalid_event_for_transport", - ) - ] - ) - return - _flush_queue(unit.output_queue, preserve=_keep_audio_sentinel) - transport.discard_pending_audio() - - elif isinstance(event, SessionUpdateEvent): - err = service.handle_session_update(session_id, event) - if err: - await transport.send_events([err]) - - elif isinstance(event, ConversationItemCreateEvent): - events = service.handle_conversation_item_create(session_id, event) - if events: - await transport.send_events(events) - - elif isinstance(event, ResponseCreateEvent): - result = service.handle_response_create(session_id, event) - if result: - if result.type != "error": - unit.cancel_scope.new_response() - await transport.send_events([result]) - - elif isinstance(event, ResponseCancelEvent): - was_active = service._state(session_id).in_response - if was_active: - unit.cancel_scope.cancel() - _flush_queue(unit.output_queue, preserve=_keep_audio_sentinel) - _flush_queue(unit.text_output_queue, preserve=_keep_user_text_event) - transport.discard_pending_audio() - events = service.handle_response_cancel(session_id) - if events: - await transport.send_events(events) - unit.response_playing.clear() - - -def create_app(pool: list[PipelineUnit], stop_event: ThreadingEvent) -> FastAPI: - @asynccontextmanager - async def lifespan(app: FastAPI) -> AsyncIterator[None]: - # One send loop per pipeline unit; each polls its own queues and forwards - # to the websocket currently attached via unit.session. - send_tasks = [asyncio.create_task(_send_loop_for(unit)) for unit in pool] - yield - for task in send_tasks: - task.cancel() - for task in send_tasks: - try: - await task - except asyncio.CancelledError: - pass - for unit in pool: - sess = unit.session - if sess is not None and sess.transport is not None: - try: - await sess.transport.close() - except Exception: - pass - - app = FastAPI(lifespan=lifespan) - - def _claim_unit(transport: SessionTransport | None) -> PipelineUnit | None: - """Atomically (between asyncio yield points) reserve the first idle unit. - - Creates a placeholder SessionState that the caller fills in with the - session_id after RealtimeService.register(). The WebRTC route claims - with transport=None and attaches the session object once constructed. - """ - for unit in pool: - if unit.session is None: - unit.session = SessionState(transport=transport) - return unit - return None - - @app.websocket("/v1/realtime") - async def realtime_endpoint(ws: WebSocket) -> None: - await ws.accept() - - transport = WebSocketTransport(ws) - unit = _claim_unit(transport) - if unit is None: - logger.warning(f"Rejected connection: all {len(pool)} pipeline slots in use") - # Stateless error event — rejection is not chargeable to any unit's usage metrics. - await send_ws_event( - ws, - build_error_event( - f"All {len(pool)} session slots are in use. Disconnect an existing client first.", - error_type="session_limit_reached", - ), - ) - await ws.close(code=1008, reason="All session slots are in use") - return - - pipeline_log_ctx.set(unit.index) - # _claim_unit guarantees unit.session is not None for the returned unit. - assert unit.session is not None - # Everything after the claim runs inside try so the finally below always - # releases the unit, even if session setup fails. - session_id = "" - try: - session_id = unit.service.register() - unit.session.session_id = session_id - logger.info(f"Client connected to pipeline {unit.index} (session {session_id})") - - # Defensive: drain edge queues and reset events so stale data from a - # previous session that survived SESSION_END propagation doesn't leak. - _clean_unit(unit) - - await send_ws_event(ws, unit.service.build_session_created(session_id)) - - while not stop_event.is_set(): - try: - raw = await asyncio.wait_for(ws.receive_json(), timeout=0.1) - except asyncio.TimeoutError: - continue - - await _dispatch_client_event(unit, session_id, raw, transport) - - except WebSocketDisconnect: - logger.info(f"Client {session_id} disconnected from pipeline {unit.index}") - except Exception as e: - logger.error(f"Client {session_id} on pipeline {unit.index} error: {type(e).__name__}: {e}", exc_info=True) - finally: - # Hold the session reference: the send loop's snapshot will still resolve - # to this object until we clear unit.session, so any handler output that - # arrives during the drain window is sent to the now-closed ws (silently - # dropped) instead of leaking to whichever client claims this unit next. - # _release_session spawns the drain-and-release as a separate task so - # this finally returns immediately. Awaiting here is unreliable: after - # WebSocketDisconnect propagates, subsequent awaits in the same task - # can be skipped/cancelled by Starlette's runner and never resume. - _release_session(unit, session_id) - - @app.get("/v1/usage") - async def usage_endpoint() -> dict[str, Any]: - # Aggregate usage across the pool. Numeric fields sum; dict fields (e.g. - # errors_by_type) merge with numeric leaves summed too, so per-unit error - # counts don't get dropped by the first-unit's value. - def _merge(into: dict[str, Any], src: dict[str, Any]) -> None: - for k, v in src.items(): - if isinstance(v, (int, float)): - into[k] = into.get(k, 0) + v - elif isinstance(v, dict): - sub = into.setdefault(k, {}) - if isinstance(sub, dict): - _merge(sub, v) - else: - into.setdefault(k, v) - - total: dict[str, Any] = {} - for unit in pool: - _merge(total, unit.service.get_usage()) - return total - - @app.get("/v1/pool") - async def pool_endpoint() -> dict[str, Any]: - now = time.monotonic() - - def _state(u: PipelineUnit) -> dict[str, Any]: - s = u.session - if s is None: - return {"index": u.index, "state": "idle", "session_id": None} - if s.released_at is None: - return {"index": u.index, "state": "active", "session_id": s.session_id} - # Drain wait gave up (quarantine timeout): the unit stays occupied - # until SESSION_END actually drains — possibly forever if a handler - # thread died. Surfaced distinctly so operators can act on it. - if s.quarantined_at is not None: - return { - "index": u.index, - "state": "stuck", - "session_id": s.session_id, - "draining_for_s": round(now - s.released_at, 2), - "stuck_for_s": round(now - s.quarantined_at, 2), - } - # released by client but SESSION_END hasn't drained yet → unit - # is still occupied; surface elapsed time so operators can spot - # stuck handlers. - return { - "index": u.index, - "state": "draining", - "session_id": s.session_id, - "draining_for_s": round(now - s.released_at, 2), - } - - return { - "size": len(pool), - "in_use": sum(1 for u in pool if u.session is not None), - "units": [_state(u) for u in pool], - } - - @app.post("/v1/realtime/calls") - async def webrtc_calls_endpoint(request: Request) -> Response: - """WebRTC SDP handshake (OpenAI GA Realtime 'calls' endpoint). - - The client POSTs an SDP offer with Content-Type: application/sdp and - receives an SDP answer. Audio then flows over WebRTC media tracks; - events flow over the 'oai-events' data channel using the same JSON - protocol as the WebSocket transport. - """ - if not WEBRTC_AVAILABLE: - return Response( - content="WebRTC support requires the 'webrtc' extra: pip install 'speech-to-speech[webrtc]'", - status_code=501, - media_type="text/plain", - ) - - if "application/sdp" not in request.headers.get("content-type", ""): - return Response( - content="Content-Type must be application/sdp", - status_code=415, - media_type="text/plain", - ) - offer_sdp = (await request.body()).decode("utf-8") - - # Claim with a placeholder transport; the send loop tolerates a - # transport-less snapshot until the session object below is attached. - unit = _claim_unit(None) - if unit is None: - logger.warning(f"Rejected WebRTC offer: all {len(pool)} pipeline slots in use") - return Response( - content=build_error_event( - f"All {len(pool)} session slots are in use. Disconnect an existing client first.", - error_type="session_limit_reached", - ).model_dump_json(), - status_code=503, - media_type="application/json", - ) - - pipeline_log_ctx.set(unit.index) - try: - session_id = unit.service.register() - assert unit.session is not None - unit.session.session_id = session_id - logger.info(f"WebRTC client claiming pipeline {unit.index} (session {session_id})") - - # Defensive: drain edge queues and reset events so stale data from a - # previous session that survived SESSION_END propagation doesn't leak. - _clean_unit(unit) - except Exception as e: # noqa: BLE001 - logger.error(f"WebRTC call setup failed (pipeline {unit.index}): {type(e).__name__}: {e}") - # No transport or drain task exists yet, so undoing the claim - # directly is the whole release. - unit.session = None - return Response(content="WebRTC session setup failed", status_code=500, media_type="text/plain") - - released = False - - def _on_closed() -> None: - # close() is idempotent but can be reached from several aiortc - # callbacks; release the unit exactly once. - nonlocal released - if released: - return - released = True - logger.info(f"WebRTC client {session_id} disconnected from pipeline {unit.index}") - _release_session(unit, session_id) - - async def _on_client_event(raw: dict[str, Any]) -> None: - assert session is not None # callbacks only fire after setup() - await _dispatch_client_event(unit, session_id, raw, session, transport_kind="webrtc") - - def _on_audio(pcm: bytes) -> None: - chunks = unit.service.append_pcm(session_id, pcm, PIPELINE_SAMPLE_RATE) - if not chunks: - return - rt_cfg = unit.service._state(session_id).runtime_config - for chunk in chunks: - unit.input_queue.put((chunk, rt_cfg)) - - async def _on_open() -> None: - assert session is not None # callbacks only fire after setup() - await session.send_events([unit.service.build_session_created(session_id)]) - logger.info(f"WebRTC session.created sent (session {session_id})") - - # Any failure between the claim above and a successful negotiate() - # must release the unit, or it stays occupied forever with no peer - # attached — the connect watchdog only exists once negotiate() ran. - session = None - try: - config = rtc_configuration_from_env() - pc = RTCPeerConnection(configuration=config) if config is not None else RTCPeerConnection() - session = WebRTCSession( - pc, - on_client_event=_on_client_event, - on_audio=_on_audio, - on_open=_on_open, - on_closed=_on_closed, - ) - session.setup() - unit.session.transport = session - except Exception as e: # noqa: BLE001 - logger.error(f"WebRTC session setup failed (session {session_id}): {type(e).__name__}: {e}") - if session is not None: - await session.close() # fires _on_closed → _release_session - else: - _on_closed() - return Response(content="WebRTC session setup failed", status_code=500, media_type="text/plain") - - try: - answer_sdp = await session.negotiate(offer_sdp) - except Exception as e: # noqa: BLE001 - logger.error(f"WebRTC negotiation failed (session {session_id}): {type(e).__name__}: {e}") - await session.close() - return Response(content="Invalid SDP offer", status_code=400, media_type="text/plain") - - logger.info(f"WebRTC SDP answer returned (session {session_id})") - return Response( - content=answer_sdp, - status_code=201, - media_type="application/sdp", - headers={"Location": f"/v1/realtime/calls/{session_id}"}, - ) - - @app.delete("/v1/realtime/calls/{call_id}") - async def webrtc_hangup_endpoint(call_id: str) -> Response: - """Hang up a WebRTC call — the Location URL advertised by the POST above.""" - for unit in pool: - session = unit.session - if ( - session is None - or session.session_id != call_id - or session.released_at is not None - or session.transport is None - or session.transport.kind != "webrtc" - ): - continue - logger.info(f"WebRTC call {call_id} hung up via DELETE (pipeline {unit.index})") - # close() fires the session's on_closed callback, which releases - # the unit exactly once (idempotent with aiortc's own callbacks). - await session.transport.close() - return Response(status_code=200) - return Response(content="Unknown call", status_code=404, media_type="text/plain") - - async def _send_loop_for(unit: PipelineUnit) -> None: - """Per-pipeline send loop. Polls this unit's output queues and forwards - to the transport currently attached via unit.session. - - Per-session scratch (pending_output_item) lives on SessionState, so it - disappears together with the transport when the session is released — - no stale sentinel can leak into the next claim. - """ - pipeline_log_ctx.set(unit.index) - while not stop_event.is_set(): - try: - # Snapshot the session once per iteration; if the route releases the - # unit mid-iteration, we continue against the prior snapshot which is - # consistent (its transport is still valid until close() returns). - session = unit.session - transport = session.transport if session is not None else None - session_id = session.session_id if session is not None else None - - # Text events first (speech_started cancels active response). - try: - text_msg = unit.text_output_queue.get_nowait() - is_speech_start = isinstance(text_msg, SpeechStartedEvent) - - was_in_response = False - was_response_pending = False - if is_speech_start and session_id: - st = unit.service._state(session_id) - was_in_response = st.in_response - was_response_pending = st.response_pending - - if isinstance(text_msg, AssistantTextEvent) and _generation_is_discardable( - unit, text_msg.cancel_generation - ): - pass - elif transport is not None and isinstance(text_msg, PipelineEvent) and session_id: - events = unit.service.dispatch_pipeline_event(session_id, text_msg) - if events: - await transport.send_events(events) - - if is_speech_start and session_id: - active_cfg = unit.service._state(session_id).runtime_config - interrupt_enabled = text_msg.interrupt_response and ( - active_cfg is None or active_cfg.interrupt_response_enabled - ) - if interrupt_enabled and transport is not None: - # Flush even when no response is active: the WebRTC - # track can still hold unplayed audio from a response - # whose done-sentinel was already observed — - # finish_response() runs on the sentinel, not when - # playback completes. No-op over WebSocket. - transport.discard_pending_audio() - if was_in_response or was_response_pending: - if interrupt_enabled: - unit.cancel_scope.cancel() - unit.service._state(session_id).response_pending = False - _flush_queue(unit.output_queue, preserve=_keep_audio_sentinel) - _flush_queue(unit.text_output_queue, preserve=_keep_user_text_event) - if unit.response_playing.is_set(): - unit.response_playing.clear() - logger.info( - "Pipeline %d: speech during %s: cancelled, queue flushed", - unit.index, - "response" if was_in_response else "pending response", - ) - else: - logger.info( - f"Pipeline {unit.index}: speech during response: interrupt_response disabled, ignoring" - ) - except Empty: - pass - - try: - if session is not None and session.pending_output_item is not None: - audio_chunk = session.pending_output_item - session.pending_output_item = None - else: - audio_chunk = unit.output_queue.get_nowait() - - if _is_pipeline_end(audio_chunk): - await _drain_pending_response_events(transport, unit, session_id) - if transport is not None and session_id: - await transport.send_events(unit.service.finish_response(session_id)) - break - - if _is_audio_done(audio_chunk): - audio_generation = _audio_generation(audio_chunk) - if audio_generation is not None and unit.cancel_scope.is_stale(audio_generation): - if session_id: - unit.service._state(session_id).response_pending = False - unit.cancel_scope.response_done(audio_generation) - unit.should_listen.set() - logger.info(f"Pipeline {unit.index}: stale response complete, listening re-enabled") - continue - await _drain_pending_response_events(transport, unit, session_id) - if transport is not None and session_id: - await transport.send_events(unit.service.finish_response(session_id)) - if session_id: - unit.service._state(session_id).response_pending = False - unit.response_playing.clear() - unit.cancel_scope.response_done(audio_generation) - unit.should_listen.set() - logger.info(f"Pipeline {unit.index}: response complete, listening re-enabled") - continue - - # SESSION_END travels from input_queue through every handler to - # output_queue. Observing it here means the chain has fully reset; - # signal the release path so it can clear unit.session. A tag from - # another session means the emitting session was force-released — - # its late SESSION_END must not stand in for this session's drain. - if is_control_message(audio_chunk, SESSION_END.kind): - chunk_session_id = getattr(audio_chunk, "session_id", None) - if session is not None and chunk_session_id in (None, session.session_id): - session.drained.set() - logger.debug(f"Pipeline {unit.index}: SESSION_END drained") - continue - - if is_control_message(audio_chunk): - continue - - if _should_discard_audio(unit, audio_chunk): - continue - - audio_chunk = _to_audio_bytes(audio_chunk) - - audio_batch = bytearray(audio_chunk) - while len(audio_batch) < MAX_AUDIO_BATCH_BYTES: - try: - next_chunk = unit.output_queue.get_nowait() - except Empty: - break - - if ( - _is_pipeline_end(next_chunk) - or _is_audio_done(next_chunk) - or is_control_message(next_chunk, SESSION_END.kind) - ): - # Only stash if we still have a session; otherwise drop it. - if session is not None: - session.pending_output_item = next_chunk - break - - if _should_discard_audio(unit, next_chunk): - continue - - next_audio = _to_audio_bytes(next_chunk) - if len(audio_batch) + len(next_audio) > MAX_AUDIO_BATCH_BYTES: - if session is not None: - session.pending_output_item = next_chunk - break - audio_batch.extend(next_audio) - - if not unit.response_playing.is_set(): - unit.response_playing.set() - unit.should_listen.set() - - if transport is not None and session_id: - await transport.send_audio_chunk(unit.service, session_id, bytes(audio_batch)) - except Empty: - pass - - await asyncio.sleep(0.01) - - except asyncio.CancelledError: - break - except Exception as e: - logger.error(f"Pipeline {unit.index} send loop error: {e}") - await asyncio.sleep(0.1) - - return app diff --git a/runtime/src/speech_to_speech/arguments_classes/__init__.py b/runtime/src/speech_to_speech/arguments_classes/__init__.py deleted file mode 100644 index 8b137891791fe96927ad78e64b0aad7bded08bdc..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/arguments_classes/__init__.py +++ /dev/null @@ -1 +0,0 @@ - diff --git a/runtime/src/speech_to_speech/arguments_classes/chat_completions_language_model_arguments.py b/runtime/src/speech_to_speech/arguments_classes/chat_completions_language_model_arguments.py deleted file mode 100644 index ad57942ec79b38fbd40a34ee85997b69e3482d17..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/arguments_classes/chat_completions_language_model_arguments.py +++ /dev/null @@ -1,28 +0,0 @@ -from dataclasses import dataclass, field -from typing import Optional - -from speech_to_speech.arguments_classes.responses_api_language_model_arguments import ( - ResponsesApiLanguageModelHandlerArguments, -) - - -@dataclass -class ChatCompletionsLanguageModelHandlerArguments(ResponsesApiLanguageModelHandlerArguments): - """Arguments for the ``chat-completions`` LLM backend. - - Inherits the OpenAI-compatible connection fields from the Responses-API - arguments (``responses_api_base_url`` / ``responses_api_api_key`` / - ``responses_api_stream`` / ``responses_api_disable_thinking``) so the same - CLI flags and launcher env vars drive both backends, and adds the - Chat-Completions-only ``reasoning_effort`` knob. - """ - - responses_api_reasoning_effort: Optional[str] = field( - default=None, - metadata={ - "help": "Provider-specific reasoning level sent as extra_body={'reasoning_effort': } on the " - "Chat Completions request. Use to disable reasoning on providers where " - "chat_template_kwargs.enable_thinking is ignored (e.g. 'none' / 'low'). When unset, falls back to " - "the disable_thinking behaviour (chat_template_kwargs.enable_thinking=false). Default is None." - }, - ) diff --git a/runtime/src/speech_to_speech/arguments_classes/chat_tts_arguments.py b/runtime/src/speech_to_speech/arguments_classes/chat_tts_arguments.py deleted file mode 100644 index 9390c1d4b440144bedad37fec4a5f38b904d7370..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/arguments_classes/chat_tts_arguments.py +++ /dev/null @@ -1,19 +0,0 @@ -from dataclasses import dataclass, field - - -@dataclass -class ChatTTSHandlerArguments: - chat_tts_stream: bool = field( - default=True, - metadata={"help": "The tts mode is stream Default is 'stream'."}, - ) - chat_tts_device: str = field( - default="cuda", - metadata={"help": "The device to be used for speech synthesis. Default is 'cuda'."}, - ) - chat_tts_chunk_size: int = field( - default=512, - metadata={ - "help": "Sets the size of the audio data chunk processed per cycle, balancing playback latency and CPU load.. Default is 512。." - }, - ) diff --git a/runtime/src/speech_to_speech/arguments_classes/facebookmms_tts_arguments.py b/runtime/src/speech_to_speech/arguments_classes/facebookmms_tts_arguments.py deleted file mode 100644 index 86f4977f00739a831e495df77067d428ea597cba..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/arguments_classes/facebookmms_tts_arguments.py +++ /dev/null @@ -1,21 +0,0 @@ -from dataclasses import dataclass, field - - -@dataclass -class FacebookMMSTTSHandlerArguments: - facebook_mms_model_name: str = field( - default="facebook/mms-tts-eng", - metadata={"help": "The model name to use. Default is 'facebook/mms-tts-eng'."}, - ) - tts_language: str = field( - default="en", - metadata={"help": "The language code for the TTS model. Default is 'en' for English."}, - ) - facebook_mms_device: str = field( - default="cuda", - metadata={"help": "The device to use for the TTS model. Default is 'cuda'."}, - ) - facebook_mms_torch_dtype: str = field( - default="float32", - metadata={"help": "The torch data type to use for the TTS model. Default is 'float32'."}, - ) diff --git a/runtime/src/speech_to_speech/arguments_classes/faster_whisper_stt_arguments.py b/runtime/src/speech_to_speech/arguments_classes/faster_whisper_stt_arguments.py deleted file mode 100644 index 5aa6286e984011538fb667e3e8c82d90acf5f022..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/arguments_classes/faster_whisper_stt_arguments.py +++ /dev/null @@ -1,54 +0,0 @@ -from dataclasses import dataclass, field - - -@dataclass -class FasterWhisperSTTHandlerArguments: - faster_whisper_stt_model_name: str = field( - default="tiny.en", - metadata={ - "help": ( - "The pretrained Faster Whisper model to use.\n" - " One of ('tiny', 'tiny.en', 'base', 'base.en', 'small', 'small.en', " - "'distil-small.en', 'medium', 'medium.en', 'distil-medium.en', 'large-v1', 'large-v2', " - "'large-v3', 'large', 'distil-large-v2', 'distil-large-v3').\n" - " Default is 'small'." - ) - }, - ) - faster_whisper_stt_device: str = field( - default="auto", - metadata={ - "help": """The device type on which the model will run. - One of ('cpu', 'cuda', 'auto'). - Default is 'auto'.""" - }, - ) - faster_whisper_stt_compute_type: str = field( - default="auto", - metadata={ - "help": """The data type to use for computation. - One of ('default', 'auto', 'int8', 'int8_float32', 'int8_float16', 'int8_bfloat16', 'int16', 'float16', 'float32', 'bfloat16') - Default is 'auto'. - Refer to 'https://opennmt.net/CTranslate2/quantization.html#quantize-on-model-loading'""" - }, - ) - faster_whisper_stt_gen_max_new_tokens: int = field( - default=128, - metadata={"help": "The maximum number of new tokens to generate. Default is 128."}, - ) - faster_whisper_stt_gen_beam_size: int = field( - default=1, - metadata={"help": "The number of beams for beam search. Default is 1, implying greedy decoding."}, - ) - faster_whisper_stt_gen_return_timestamps: bool = field( - default=False, - metadata={"help": "Whether to return timestamps with transcriptions. Default is False."}, - ) - faster_whisper_stt_gen_task: str = field( - default="transcribe", - metadata={"help": "The task to perform, typically 'transcribe' for transcription. Default is 'transcribe'."}, - ) - faster_whisper_stt_gen_language: str = field( - default="en", - metadata={"help": "The language of the speech to transcribe. Default is 'en' for English."}, - ) diff --git a/runtime/src/speech_to_speech/arguments_classes/kokoro_tts_arguments.py b/runtime/src/speech_to_speech/arguments_classes/kokoro_tts_arguments.py deleted file mode 100644 index 78f95cb833869ccb3bcd796e7cdff96c419135e6..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/arguments_classes/kokoro_tts_arguments.py +++ /dev/null @@ -1,38 +0,0 @@ -from dataclasses import dataclass, field -from typing import Optional - - -@dataclass -class KokoroTTSHandlerArguments: - kokoro_model_name: Optional[str] = field( - default=None, - metadata={ - "help": "The Kokoro TTS model to use. Auto-selects based on device: 'mlx-community/Kokoro-82M-bf16' for MPS, 'hexgrad/Kokoro-82M' for CUDA/CPU." - }, - ) - kokoro_device: str = field( - default="auto", - metadata={ - "help": "The device to run Kokoro TTS on. Options: 'auto', 'cuda', 'cpu', 'mps'. Default is 'auto' (MPS on Mac, CUDA on GPU systems)." - }, - ) - kokoro_voice: str = field( - default="bm_fable", - metadata={ - "help": "The voice to use for synthesis. See VOICES.md in the Kokoro repo for options. Default is 'bm_fable' (British male)." - }, - ) - kokoro_lang_code: str = field( - default="b", - metadata={ - "help": "Language code: 'a' for American English, 'b' for British English, 'j' for Japanese, etc. Default is 'b'." - }, - ) - kokoro_speed: float = field( - default=1.0, - metadata={"help": "Speech speed multiplier. Values > 1.0 speed up, < 1.0 slow down. Default is 1.0."}, - ) - kokoro_blocksize: int = field( - default=512, - metadata={"help": "The audio chunk size in samples for streaming output. Default is 512."}, - ) diff --git a/runtime/src/speech_to_speech/arguments_classes/language_model_arguments.py b/runtime/src/speech_to_speech/arguments_classes/language_model_arguments.py deleted file mode 100644 index d076b4f57efc9c3e17d3bd8942f006aa6acd8ae9..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/arguments_classes/language_model_arguments.py +++ /dev/null @@ -1,43 +0,0 @@ -from dataclasses import dataclass, field - -from speech_to_speech.arguments_classes.language_model_base_arguments import LanguageModelBaseArguments - - -@dataclass -class LanguageModelHandlerArguments(LanguageModelBaseArguments): - llm_device: str = field( - default="cuda", - metadata={"help": "The device type on which the model will run. Default is 'cuda' for GPU acceleration."}, - ) - llm_torch_dtype: str = field( - default="float16", - metadata={ - "help": "The PyTorch data type for the model and input tensors. One of `float32` (full-precision), `float16` or `bfloat16` (both half-precision)." - }, - ) - llm_gen_max_new_tokens: int = field( - default=1024, - metadata={"help": "Maximum number of new tokens to generate in a single completion. Default is 1024."}, - ) - llm_gen_min_new_tokens: int = field( - default=0, - metadata={"help": "Minimum number of new tokens to generate in a single completion. Default is 0."}, - ) - llm_gen_temperature: float = field( - default=0.0, - metadata={ - "help": "Controls the randomness of the output. Set to 0.0 for deterministic (repeatable) outputs. Default is 0.0." - }, - ) - llm_gen_do_sample: bool = field( - default=False, - metadata={"help": "Whether to use sampling; set this to False for deterministic outputs. Default is False."}, - ) - llm_is_vlm: bool = field( - default=False, - metadata={ - "help": "Set to True when using a Vision Language Model (VLM) that accepts image inputs. " - "Loads AutoProcessor + AutoModelForImageTextToText instead of the default text-only model. " - "Default is False." - }, - ) diff --git a/runtime/src/speech_to_speech/arguments_classes/language_model_base_arguments.py b/runtime/src/speech_to_speech/arguments_classes/language_model_base_arguments.py deleted file mode 100644 index 6b23e5e686ccbdbecb57357222aec890711a3013..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/arguments_classes/language_model_base_arguments.py +++ /dev/null @@ -1,46 +0,0 @@ -from dataclasses import dataclass, field - - -@dataclass -class LanguageModelBaseArguments: - model_name: str = field( - default="Qwen/Qwen3-4B-Instruct-2507", - metadata={"help": "The pretrained language model to use."}, - ) - user_role: str = field( - default="user", - metadata={"help": "Role assigned to the user in the chat context. Default is 'user'."}, - ) - init_chat_role: str = field( - default="system", - metadata={"help": "Initial role for setting up the chat context. Default is 'system'."}, - ) - init_chat_prompt: str = field( - default="You are a helpful and friendly AI assistant. You are polite, respectful, and aim to provide concise responses of less than 20 words.", - metadata={"help": "The initial chat prompt to establish context for the language model."}, - ) - chat_size: int = field( - default=30, - metadata={"help": "Number of interactions assistant-user to keep for the chat."}, - ) - stream_batch_sentences: int = field( - default=3, - metadata={ - "help": "Number of sentences to accumulate before yielding a batch during streaming. " - "Set to 1 for sentence-by-sentence streaming. Default is 3." - }, - ) - enable_lang_prompt: bool = field( - default=False, - metadata={ - "help": "When True, append a user message instructing the model to reply in the detected/selected " - "language (e.g. 'Please reply to my message in French.'). Default is False." - }, - ) - compact_history: bool = field( - default=True, - metadata={ - "help": "When True, summarize older turns in the background once the chat exceeds chat_size, " - "instead of synchronously evicting them. Adds an extra LLM call per compaction. Default is True." - }, - ) diff --git a/runtime/src/speech_to_speech/arguments_classes/mlx_audio_whisper_arguments.py b/runtime/src/speech_to_speech/arguments_classes/mlx_audio_whisper_arguments.py deleted file mode 100644 index ad64ebc484279b6f42156ceeffea9f73355e0293..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/arguments_classes/mlx_audio_whisper_arguments.py +++ /dev/null @@ -1,15 +0,0 @@ -from dataclasses import dataclass, field - - -@dataclass -class MLXAudioWhisperSTTHandlerArguments: - mlx_audio_whisper_model_name: str = field( - default="mlx-community/whisper-large-v3-turbo", - metadata={ - "help": "The pretrained MLX Audio Whisper model to use. Default is 'mlx-community/whisper-large-v3-turbo'." - }, - ) - mlx_audio_whisper_gen_kwargs: dict = field( - default_factory=dict, - metadata={"help": "Additional generation kwargs to pass to the model. Default is an empty dict."}, - ) diff --git a/runtime/src/speech_to_speech/arguments_classes/module_arguments.py b/runtime/src/speech_to_speech/arguments_classes/module_arguments.py deleted file mode 100644 index f9b3e5e56ac18714dc9d85b5202ff508a4f0c18a..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/arguments_classes/module_arguments.py +++ /dev/null @@ -1,84 +0,0 @@ -from dataclasses import dataclass, field -from typing import Literal, Optional - - -@dataclass -class ModuleArguments: - device: Optional[str] = field( - default=None, - metadata={"help": "If specified, overrides the device for all handlers."}, - ) - mode: Optional[Literal["local", "socket", "websocket", "realtime"]] = field( - default="realtime", - metadata={ - "help": "The mode to run the pipeline in. Either 'local', 'socket', 'websocket', or 'realtime'. Default is 'realtime'." - }, - ) - local_mac_optimal_settings: bool = field( - default=False, - metadata={ - "help": "If specified, sets the optimal settings for Mac OS. Sets Parakeet TDT for STT, MLX LM for language model, and Qwen3-TTS for TTS, with MPS device and local mode." - }, - ) - stt: Optional[ - Literal[ - "none", - "whisper", - "whisper-mlx", - "mlx-audio-whisper", - "faster-whisper", - "parakeet-tdt", - "paraformer", - ] - ] = field( - default="parakeet-tdt", - metadata={ - "help": "The STT to use. Use 'none' to send VAD audio directly to an audio-input LLM. " - "This requires --llm_backend chat-completions and an explicitly selected audio-capable " - "API model with --model_name. " - "Otherwise choose 'whisper', 'whisper-mlx', 'mlx-audio-whisper', 'faster-whisper', " - "'parakeet-tdt', or 'paraformer'. Default is 'parakeet-tdt'." - }, - ) - llm_backend: Optional[Literal["transformers", "mlx-lm", "responses-api", "chat-completions"]] = field( - default="responses-api", - metadata={ - "help": "The LLM backend to use. Either 'transformers', 'mlx-lm', 'responses-api', or " - "'chat-completions' (OpenAI-compatible /v1/chat/completions). Default is 'responses-api'." - }, - ) - tts: Optional[Literal["chatTTS", "facebookMMS", "pocket", "kokoro", "qwen3"]] = field( - default="qwen3", - metadata={ - "help": "The TTS to use. Either 'chatTTS', 'facebookMMS', 'pocket', 'kokoro', or 'qwen3'. Default is 'qwen3'." - }, - ) - log_level: str = field( - default="info", - metadata={"help": "Provide logging level. Example --log_level debug, default=info."}, - ) - enable_live_transcription: bool = field( - default=True, - metadata={ - "help": "Enable live transcription display while user is speaking (works with parakeet-tdt). Default is true." - }, - ) - live_transcription_update_interval: float = field( - default=0.5, - metadata={"help": "Update interval for live transcription in seconds (default: 0.5s = 500ms)"}, - ) - live_transcription_min_silence_ms: int = field( - default=500, - metadata={ - "help": "Minimum silence duration (ms) before ending speech when live transcription is enabled (default: 500ms)" - }, - ) - num_pipelines: int = field( - default=1, - metadata={ - "help": "Number of isolated realtime pipelines in the pool. One uvicorn server listens on " - "--ws_port and routes each incoming websocket to the next free pipeline (each has its own " - "VAD/STT/LM/TTS handlers and conversation state). Max concurrent websocket sessions equals " - "num_pipelines; further connections are rejected. Only valid for --mode realtime. Default is 1." - }, - ) diff --git a/runtime/src/speech_to_speech/arguments_classes/paraformer_stt_arguments.py b/runtime/src/speech_to_speech/arguments_classes/paraformer_stt_arguments.py deleted file mode 100644 index 0fdfe4faac184ee7d614ed01a66d4760d8863575..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/arguments_classes/paraformer_stt_arguments.py +++ /dev/null @@ -1,15 +0,0 @@ -from dataclasses import dataclass, field - - -@dataclass -class ParaformerSTTHandlerArguments: - paraformer_stt_model_name: str = field( - default="paraformer-zh", - metadata={ - "help": "The pretrained model to use. Default is 'paraformer-zh'. Can be choose from https://github.com/modelscope/FunASR" - }, - ) - paraformer_stt_device: str = field( - default="cuda", - metadata={"help": "The device type on which the model will run. Default is 'cuda' for GPU acceleration."}, - ) diff --git a/runtime/src/speech_to_speech/arguments_classes/parakeet_tdt_arguments.py b/runtime/src/speech_to_speech/arguments_classes/parakeet_tdt_arguments.py deleted file mode 100644 index 13e367fcf595e27e5a90931839c9d32c084e2b96..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/arguments_classes/parakeet_tdt_arguments.py +++ /dev/null @@ -1,39 +0,0 @@ -from dataclasses import dataclass, field -from typing import Optional - - -@dataclass -class ParakeetTDTSTTHandlerArguments: - """ - Arguments for the Parakeet TDT Speech-to-Text handler. - - Parakeet TDT 0.6B v3 is a 600M parameter multilingual ASR model from NVIDIA. - - On MPS (Apple Silicon): Uses mlx-audio with mlx-community/parakeet-tdt-0.6b-v3 - - On CUDA/CPU: Uses nano-parakeet (pure PyTorch) with nvidia/parakeet-tdt-0.6b-v3 - """ - - parakeet_tdt_model_name: Optional[str] = field( - default=None, - metadata={ - "help": "The Parakeet TDT model to use. Defaults to 'mlx-community/parakeet-tdt-0.6b-v3' " - "for MPS or 'nvidia/parakeet-tdt-0.6b-v3' for CUDA/CPU." - }, - ) - parakeet_tdt_device: str = field( - default="auto", - metadata={ - "help": "Device to run the model on. 'auto' will use MPS on macOS and CUDA otherwise. " - "Options: 'auto', 'cuda', 'mps', 'cpu'. Default is 'auto'." - }, - ) - parakeet_tdt_compute_type: str = field( - default="float16", - metadata={"help": "Compute type for the model. Options: 'float16', 'float32'. Default is 'float16'."}, - ) - parakeet_tdt_language: Optional[str] = field( - default=None, - metadata={ - "help": "Target language code for transcription. If not specified, the model will " - "auto-detect the language. Supports 25 European languages." - }, - ) diff --git a/runtime/src/speech_to_speech/arguments_classes/pocket_tts_arguments.py b/runtime/src/speech_to_speech/arguments_classes/pocket_tts_arguments.py deleted file mode 100644 index fcd0bcd34a23c2a8d97e01dbef445bc4ec1c8986..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/arguments_classes/pocket_tts_arguments.py +++ /dev/null @@ -1,31 +0,0 @@ -from dataclasses import dataclass, field - - -@dataclass -class PocketTTSHandlerArguments: - pocket_tts_device: str = field( - default="cpu", - metadata={ - "help": "The device type on which the Pocket TTS model will run. Options: 'cpu', 'cuda', 'mps'. Default is 'cpu'." - }, - ) - pocket_tts_voice: str = field( - default="jean", - metadata={ - "help": "Voice to use for Pocket TTS. Can be a preset name ('alba', 'marius', 'javert', 'jean', 'fantine', 'cosette', 'eponine', 'azelma'), a local audio file path, or a HuggingFace path like 'hf://kyutai/tts-voices/...'. Default is 'jean'." - }, - ) - pocket_tts_sample_rate: int = field( - default=16000, - metadata={ - "help": "Output sample rate in Hz for Pocket TTS. Pocket TTS uses 24kHz internally but will be resampled to this rate to match the audio output. Default is 16000 to match the pipeline's audio streamer." - }, - ) - pocket_tts_blocksize: int = field( - default=512, - metadata={"help": "Size of audio blocks to yield for streaming in Pocket TTS. Default is 512."}, - ) - pocket_tts_max_tokens: int = field( - default=50, - metadata={"help": "Maximum number of tokens to generate per sentence in Pocket TTS. Default is 50."}, - ) diff --git a/runtime/src/speech_to_speech/arguments_classes/qwen3_tts_arguments.py b/runtime/src/speech_to_speech/arguments_classes/qwen3_tts_arguments.py deleted file mode 100644 index b857d17e783031a2541408a7463a6593ece0c7be..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/arguments_classes/qwen3_tts_arguments.py +++ /dev/null @@ -1,100 +0,0 @@ -from dataclasses import dataclass, field -from typing import Literal, Optional - - -@dataclass -class Qwen3TTSHandlerArguments: - qwen3_tts_model_name: str = field( - default="Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice", - metadata={ - "help": "The Qwen3-TTS model to use (HuggingFace Hub ID or local path). Default is 'Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice'. On Apple Silicon, Qwen/* model IDs are auto-mapped to the corresponding mlx-community/* model when possible, defaulting to the 6bit MLX variant unless the model name already pins a specific suffix." - }, - ) - qwen3_tts_device: str = field( - default="cuda", - metadata={ - "help": "Preferred device for Qwen3-TTS. Options: 'cuda', 'cpu', 'mps', 'auto'. Default is 'cuda'. On Apple Silicon the mlx-audio backend is selected automatically." - }, - ) - qwen3_tts_dtype: str = field( - default="auto", - metadata={ - "help": "Data type for inference. Options: 'auto', 'float16', 'bfloat16', 'float32'. Default is 'auto'." - }, - ) - qwen3_tts_attn_implementation: str = field( - default="eager", - metadata={ - "help": "Attention implementation. Options: 'eager', 'flash_attention_2', 'sdpa'. Use 'eager' on Jetson. Default is 'eager'." - }, - ) - qwen3_tts_backend: Literal["ggml", "torch"] = field( - default="ggml", - metadata={ - "help": "faster-qwen3-tts backend on non-macOS platforms. Options: 'ggml' or 'torch'. Default is 'ggml'. On Apple Silicon, mlx-audio is selected automatically and this option is ignored." - }, - ) - qwen3_tts_ref_audio: Optional[str] = field( - default=None, - metadata={ - "help": "Optional path to reference audio file for voice cloning. Leave unset when using a CustomVoice model." - }, - ) - qwen3_tts_ref_text: str = field( - default="I'm confused why some people have super short timelines, yet at the same time are bullish on scaling up reinforcement learning atop LLMs. If we're actually close to a human-like learner, then this whole approach of training on verifiable outcomes.", - metadata={"help": "Transcription of the reference audio for voice cloning."}, - ) - qwen3_tts_speaker: Optional[str] = field( - default="Aiden", - metadata={ - "help": "Speaker name for CustomVoice models. Default is 'Aiden'. If not provided, the first supported speaker is used when available." - }, - ) - qwen3_tts_instruct: Optional[str] = field( - default=None, - metadata={"help": "Instruction text for VoiceDesign models (optional, required for voice design)."}, - ) - qwen3_tts_xvec_only: bool = field( - default=False, - metadata={ - "help": "Use x-vector only voice cloning mode (recommended for cleaner starts and language switching). Default is False." - }, - ) - qwen3_tts_parity_mode: bool = field( - default=False, - metadata={"help": "Disable CUDA-graph streaming path and use parity mode for stability. Default is False."}, - ) - qwen3_tts_non_streaming_mode: Optional[bool] = field( - default=True, - metadata={ - "help": "Optional override for Qwen3-TTS text prefill behavior. Default is true, which pre-fills the full target text before decode on faster-qwen3-tts. Currently ignored on Apple Silicon because mlx-audio does not expose this yet." - }, - ) - qwen3_tts_mlx_quantization: Optional[str] = field( - default="6bit", - metadata={ - "help": "Optional MLX quantization override on Apple Silicon. Supported values: 'bf16', '4bit', '6bit', '8bit'. Default is '6bit'." - }, - ) - qwen3_tts_language: str = field( - default="auto", - metadata={"help": "Target language for synthesis. Default is 'auto'."}, - ) - qwen3_tts_streaming_chunk_size: Optional[int] = field( - default=None, - metadata={ - "help": "Codec steps per streaming chunk. If unset, the handler uses a backend-specific default: 8 on faster-qwen3-tts and 4 on mlx-audio." - }, - ) - qwen3_tts_max_new_tokens: int = field( - default=1536, - metadata={ - "help": "Upper cap for Qwen3-TTS codec tokens. The handler estimates a per-utterance budget from the text and clamps it to this ceiling (~12 tokens per second of audio). Raise this above 1536 if you want to allow longer utterances." - }, - ) - qwen3_tts_blocksize: int = field( - default=512, - metadata={ - "help": "Audio chunk size in samples for streaming output. Must match LocalAudioStreamer blocksize. Default is 512." - }, - ) diff --git a/runtime/src/speech_to_speech/arguments_classes/responses_api_language_model_arguments.py b/runtime/src/speech_to_speech/arguments_classes/responses_api_language_model_arguments.py deleted file mode 100644 index 7e9df371570bb31c88fb4e2cb26b37d4aead829d..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/arguments_classes/responses_api_language_model_arguments.py +++ /dev/null @@ -1,59 +0,0 @@ -from dataclasses import dataclass, field -from typing import Literal, Optional - -from speech_to_speech.arguments_classes.language_model_base_arguments import LanguageModelBaseArguments - - -@dataclass -class ResponsesApiLanguageModelHandlerArguments(LanguageModelBaseArguments): - model_name: str = field( - default="gpt-5.4-mini", - metadata={ - "help": "The model to use with the OpenAI-compatible API. Default is 'gpt-5.4-mini', " - "which is not audio-capable; --stt none requires an explicitly selected audio-input model." - }, - ) - responses_api_api_key: Optional[str] = field( - default=None, - metadata={"help": "API key used to authenticate access to the OpenAI-compatible API. Default is None."}, - ) - responses_api_base_url: Optional[str] = field( - default=None, - metadata={"help": "Base URL for the OpenAI-compatible API endpoint. Default is None (uses OpenAI)."}, - ) - responses_api_stream: bool = field( - default=True, - metadata={ - "help": "The stream parameter typically indicates whether data should be transmitted in a continuous flow rather" - " than in a single, complete response, often used for handling large or real-time data.Default is True" - }, - ) - responses_api_disable_thinking: bool = field( - default=True, - metadata={ - "help": "Disable provider-side thinking/reasoning when supported by the OpenAI-compatible backend. " - "For Together Qwen3.5 models this sends chat_template_kwargs.enable_thinking=false." - }, - ) - responses_api_audio_max_tokens: int = field( - default=256, - metadata={"help": "Maximum chat completion tokens for audio-input LLM requests. Default is 256."}, - ) - responses_api_audio_temperature: float = field( - default=0.0, - metadata={"help": "Chat completion temperature for audio-input LLM requests. Default is 0.0."}, - ) - responses_api_audio_content_type: Literal["input_audio", "audio_url"] = field( - default="input_audio", - metadata={ - "help": "Chat Completions content representation for audio: 'input_audio' embeds WAV base64 directly; " - "'audio_url' sends a base64 data URL. Default is 'input_audio'." - }, - ) - responses_api_audio_history_turns: int = field( - default=1, - metadata={ - "help": "Number of recent completed audio user turns retained in history. Older audio is replaced with " - "a role-preserving placeholder. Set to 0 to replace every completed audio turn. Default is 1." - }, - ) diff --git a/runtime/src/speech_to_speech/arguments_classes/socket_receiver_arguments.py b/runtime/src/speech_to_speech/arguments_classes/socket_receiver_arguments.py deleted file mode 100644 index e69655c1b7a8dba24ee448179b550aa11a163a89..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/arguments_classes/socket_receiver_arguments.py +++ /dev/null @@ -1,20 +0,0 @@ -from dataclasses import dataclass, field - - -@dataclass -class SocketReceiverArguments: - recv_host: str = field( - default="localhost", - metadata={ - "help": "The host IP ddress for the socket connection. Default is '0.0.0.0' which binds to all " - "available interfaces on the host machine." - }, - ) - recv_port: int = field( - default=12345, - metadata={"help": "The port number on which the socket server listens. Default is 12346."}, - ) - chunk_size: int = field( - default=1024, - metadata={"help": "The size of each data chunk to be sent or received over the socket. Default is 1024 bytes."}, - ) diff --git a/runtime/src/speech_to_speech/arguments_classes/socket_sender_arguments.py b/runtime/src/speech_to_speech/arguments_classes/socket_sender_arguments.py deleted file mode 100644 index e0011d41ea018410babee34bbe2aa9311619609c..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/arguments_classes/socket_sender_arguments.py +++ /dev/null @@ -1,16 +0,0 @@ -from dataclasses import dataclass, field - - -@dataclass -class SocketSenderArguments: - send_host: str = field( - default="localhost", - metadata={ - "help": "The host IP address for the socket connection. Default is '0.0.0.0' which binds to all " - "available interfaces on the host machine." - }, - ) - send_port: int = field( - default=12346, - metadata={"help": "The port number on which the socket server listens. Default is 12346."}, - ) diff --git a/runtime/src/speech_to_speech/arguments_classes/vad_arguments.py b/runtime/src/speech_to_speech/arguments_classes/vad_arguments.py deleted file mode 100644 index cd36378a422b2c636b2939a0c20b990e0f23265d..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/arguments_classes/vad_arguments.py +++ /dev/null @@ -1,81 +0,0 @@ -from dataclasses import dataclass, field - - -@dataclass -class VADHandlerArguments: - thresh: float = field( - default=0.6, - metadata={ - "help": "The threshold value for voice activity detection (VAD). Values typically range from 0 to 1, with higher values requiring higher confidence in speech detection." - }, - ) - sample_rate: int = field( - default=16000, - metadata={ - "help": "The sample rate of the audio in Hertz. Default is 16000 Hz, which is a common setting for voice audio." - }, - ) - min_silence_ms: int = field( - default=64, - metadata={ - "help": "Minimum length of silence intervals to be used for segmenting speech. Measured in milliseconds. Default is 64 ms." - }, - ) - min_speech_ms: int = field( - default=384, - metadata={ - "help": "Minimum length of speech segments to be considered valid speech. Measured in milliseconds. Default is 384 ms." - }, - ) - min_speech_continuation_ms: int = field( - default=192, - metadata={ - "help": "Hysteresis threshold (ms of active speech) for accepting speech that continues a reopenable turn (soft-ended, uncommitted, within the reopen window). Set to 0 to disable the split and use min_speech_ms. Clamped to [100, min_speech_ms]. New turns and barge-ins always require min_speech_ms. Default and recommended: 192 with min_speech_ms 384." - }, - ) - max_speech_ms: float = field( - default=float("inf"), - metadata={ - "help": "Maximum length of continuous speech before forcing a split. Default is infinite, allowing for uninterrupted speech segments." - }, - ) - speech_pad_ms: int = field( - default=500, - metadata={ - "help": "Amount of audio retained before VAD triggers and prepended to detected speech segments. Once speech is detected, audio continues to be kept until VAD declares the segment done. Measured in milliseconds. Default is 500 ms." - }, - ) - audio_enhancement: bool = field( - default=False, - metadata={ - "help": "improves sound quality by applying techniques like noise reduction, equalization, and echo cancellation. Default is False." - }, - ) - enable_realtime_transcription: bool = field( - default=False, - metadata={"help": "Enable progressive audio release for live transcription during speech. Default is False."}, - ) - realtime_processing_pause: float = field( - default=0.5, - metadata={ - "help": "Interval (in seconds) for releasing progressive audio chunks during speech. Default is 0.5s." - }, - ) - speculative_reopen_ms: int = field( - default=1000, - metadata={ - "help": "In realtime mode, keep a soft-ended turn reopenable for this many milliseconds unless a response commits it." - }, - ) - unanswered_reopen_ms: int = field( - default=7000, - metadata={ - "help": "Sanity cap (ms) for reopening a soft-ended speculative turn that has not yet been answered by any assistant output. While a turn is uncommitted, resumed speech within this window reopens the same turn instead of starting a new one. Has no effect below speculative_reopen_ms." - }, - ) - short_segment_merge_ms: int = field( - default=0, - metadata={ - "help": "When greater than 0, adjacent VAD segments below min_speech_ms are held and stitched for this many milliseconds before being discarded. Fragments shorter than 100 ms of active speech are never held. Useful with very low min_silence_ms values." - }, - ) diff --git a/runtime/src/speech_to_speech/arguments_classes/websocket_streamer_arguments.py b/runtime/src/speech_to_speech/arguments_classes/websocket_streamer_arguments.py deleted file mode 100644 index 09a57b5adee92ba9cf5a758f735aa93dc41d57d3..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/arguments_classes/websocket_streamer_arguments.py +++ /dev/null @@ -1,16 +0,0 @@ -from dataclasses import dataclass, field - - -@dataclass -class WebSocketStreamerArguments: - ws_host: str = field( - default="0.0.0.0", - metadata={ - "help": "The host IP address for the WebSocket server. Default is '0.0.0.0' which binds to all " - "available interfaces on the host machine." - }, - ) - ws_port: int = field( - default=8765, - metadata={"help": "The port number on which the WebSocket server listens. Default is 8765."}, - ) diff --git a/runtime/src/speech_to_speech/arguments_classes/whisper_stt_arguments.py b/runtime/src/speech_to_speech/arguments_classes/whisper_stt_arguments.py deleted file mode 100644 index 617ae04266f95063c675613ff134ec96b2ef8a3e..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/arguments_classes/whisper_stt_arguments.py +++ /dev/null @@ -1,52 +0,0 @@ -from dataclasses import dataclass, field -from typing import Optional - - -@dataclass -class WhisperSTTHandlerArguments: - stt_model_name: str = field( - default="distil-whisper/distil-large-v3", - metadata={"help": "The pretrained Whisper model to use. Default is 'distil-whisper/distil-large-v3'."}, - ) - stt_device: str = field( - default="cuda", - metadata={"help": "The device type on which the model will run. Default is 'cuda' for GPU acceleration."}, - ) - stt_torch_dtype: str = field( - default="float16", - metadata={ - "help": "The PyTorch data type for the model and input tensors. One of `float32` (full-precision), `float16` or `bfloat16` (both half-precision)." - }, - ) - stt_compile_mode: Optional[str] = field( - default=None, - metadata={ - "help": "Compile mode for torch compile. Either 'default', 'reduce-overhead' and 'max-autotune'. Default is None (no compilation)" - }, - ) - stt_gen_max_new_tokens: int = field( - default=128, - metadata={"help": "The maximum number of new tokens to generate. Default is 128."}, - ) - stt_gen_num_beams: int = field( - default=1, - metadata={"help": "The number of beams for beam search. Default is 1, implying greedy decoding."}, - ) - stt_gen_return_timestamps: bool = field( - default=False, - metadata={"help": "Whether to return timestamps with transcriptions. Default is False."}, - ) - stt_gen_task: str = field( - default="transcribe", - metadata={"help": "The task to perform, typically 'transcribe' for transcription. Default is 'transcribe'."}, - ) - language: Optional[str] = field( - default="en", - metadata={ - "help": """The language for the conversation. - Choose between 'en' (english), 'fr' (french), 'es' (spanish), - 'zh' (chinese), 'ko' (korean), 'ja' (japanese), 'hi' (hindi) or 'None'. - If using 'auto', the language is automatically detected and can - change during the conversation. Default is 'en'.""" - }, - ) diff --git a/runtime/src/speech_to_speech/baseHandler.py b/runtime/src/speech_to_speech/baseHandler.py deleted file mode 100644 index f9658522643cefee15bec822795484333775850c..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/baseHandler.py +++ /dev/null @@ -1,161 +0,0 @@ -from __future__ import annotations - -# ruff: noqa: I001 - -import logging -from queue import Empty, Queue -from threading import Event -from time import perf_counter -from typing import Any, Generic, Iterator, TypeVar, cast - -import numpy as np - -from speech_to_speech.pipeline.control import PipelineControlMessage, is_control_message, SESSION_END -from speech_to_speech.pipeline.log_context import pipeline_log_ctx -from speech_to_speech.pipeline.messages import PIPELINE_END, AudioOutput, EndOfResponse - -logger = logging.getLogger(__name__) - -InT = TypeVar("InT") -OutT = TypeVar("OutT") - - -class BaseHandler(Generic[InT, OutT]): - """ - Base class for pipeline parts. Each part of the pipeline has an input and an output queue. - The `setup` method along with `setup_args` and `setup_kwargs` can be used to address the specific requirements of the implemented pipeline part. - To stop a handler properly, set the stop_event and, to avoid queue deadlocks, place PIPELINE_END in the input queue. - Objects placed in the input queue will be processed by the `process` method, and the yielded results will be placed in the output queue. - `SESSION_END` is a soft control message used to reset per-session state without stopping the handler thread. - The cleanup method handles stopping the handler, and PIPELINE_END is placed in the output queue. - """ - - def __init__( - self, - stop_event: Event, - queue_in: Queue[InT | PipelineControlMessage | bytes], - queue_out: Queue[OutT | PipelineControlMessage | bytes], - setup_args: tuple[Any, ...] = (), - setup_kwargs: dict[str, Any] = {}, - ) -> None: - self.stop_event = stop_event - self.queue_in = queue_in - self.queue_out = queue_out - self.pipeline_index: int | None = None - self.setup(*setup_args, **setup_kwargs) - self._times: list[float] = [] - - def setup(self, *arg: Any, **kwargs: Any) -> None: - pass - - def process(self, input: InT) -> Iterator[OutT]: - raise NotImplementedError - - def should_process_input(self, item: InT) -> bool: - cancel_scope = getattr(self, "cancel_scope", None) - cancel_generation = getattr(item, "cancel_generation", None) - if ( - cancel_scope is not None - and cancel_generation is not None - and not isinstance(item, EndOfResponse) - and cancel_scope.is_stale(cancel_generation) - ): - logger.debug( - "%s: dropping stale input for cancel generation %s", - self.__class__.__name__, - cancel_generation, - ) - return False - return True - - def should_emit_output(self, output: OutT) -> bool: - return True - - def before_emit_output(self, output: OutT) -> None: - pass - - def output_for_queue(self, output: OutT, source_input: InT) -> OutT | AudioOutput: - cancel_generation = getattr(source_input, "cancel_generation", None) - if cancel_generation is not None and (isinstance(output, bytes) or hasattr(output, "tobytes")): - audio = cast(bytes | np.ndarray, output) - return AudioOutput(audio=audio, cancel_generation=cancel_generation) - return output - - def run(self) -> None: - if self.pipeline_index is not None: - pipeline_log_ctx.set(self.pipeline_index) - logger.debug(f"{self.__class__.__name__}: Handler thread started") - while not self.stop_event.is_set(): - try: - # Use timeout to check stop_event periodically - item = self.queue_in.get(timeout=0.1) - except Empty: - continue - - if isinstance(item, PipelineControlMessage) and is_control_message(item, SESSION_END.kind): - logger.debug(f"{self.__class__.__name__}: session end received") - try: - self.on_session_end() - except Exception as e: - logger.error( - f"{self.__class__.__name__}: Error in on_session_end(): {type(e).__name__}: {e}", - exc_info=True, - ) - self.queue_out.put(item) - continue - - if isinstance(item, bytes) and item == PIPELINE_END: - # sentinel signal to avoid queue deadlock - logger.debug("Stopping thread") - break - - if isinstance(item, PipelineControlMessage): - logger.warning("%s: unexpected control message kind: %s", self.__class__.__name__, item.kind) - continue - - typed_item = cast(InT, item) - if not self.should_process_input(typed_item): - continue - - start_time = perf_counter() - try: - for output in self.process(typed_item): - if not self.should_emit_output(output): - start_time = perf_counter() - continue - self._times.append(perf_counter() - start_time) - if self.should_log_timing(output): - logger.log(self.timing_log_level, "%s: %.3f s", self.__class__.__name__, self.last_time) - self.before_emit_output(output) - queued_output = cast( - OutT | PipelineControlMessage | bytes, - self.output_for_queue(output, typed_item), - ) - self.queue_out.put(queued_output) - start_time = perf_counter() - except Exception as e: - logger.error(f"{self.__class__.__name__}: Error in process(): {type(e).__name__}: {e}", exc_info=True) - - self.cleanup() - self.queue_out.put(PIPELINE_END) - - @property - def last_time(self) -> float: - return self._times[-1] - - @property - def min_time_to_debug(self) -> float: - return 0.001 - - @property - def timing_log_level(self) -> int: - return logging.DEBUG - - def should_log_timing(self, output: OutT) -> bool: - return self.last_time > self.min_time_to_debug - - def cleanup(self) -> None: - pass - - def on_session_end(self) -> None: - pass diff --git a/runtime/src/speech_to_speech/connections/__init__.py b/runtime/src/speech_to_speech/connections/__init__.py deleted file mode 100644 index 8b137891791fe96927ad78e64b0aad7bded08bdc..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/connections/__init__.py +++ /dev/null @@ -1 +0,0 @@ - diff --git a/runtime/src/speech_to_speech/connections/local_audio_streamer.py b/runtime/src/speech_to_speech/connections/local_audio_streamer.py deleted file mode 100644 index 2ed283703ebcc2d0e684481872df447a8609072a..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/connections/local_audio_streamer.py +++ /dev/null @@ -1,71 +0,0 @@ -import logging -import threading -import time -from queue import Queue - -import numpy as np -import sounddevice as sd - -from speech_to_speech.pipeline.messages import AUDIO_RESPONSE_DONE -from speech_to_speech.pipeline.queue_types import AudioInItem, AudioOutItem - -logger = logging.getLogger(__name__) - - -class LocalAudioStreamer: - def __init__( - self, - input_queue: Queue[AudioInItem], - output_queue: Queue[AudioOutItem], - should_listen: threading.Event, - list_play_chunk_size: int = 512, - ) -> None: - self.list_play_chunk_size = list_play_chunk_size - - self.stop_event = threading.Event() - self.input_queue = input_queue - self.output_queue = output_queue - self.should_listen = should_listen - - def run(self) -> None: - # Pre-generate a static dither buffer (±1 LSB, -96 dB) to keep the - # audio sink active without calling numpy inside the real-time callback. - dither = np.random.randint(-1, 2, size=(self.list_play_chunk_size, 1), dtype=np.int16) - - def callback(indata: np.ndarray, outdata: np.ndarray, frames: int, time: float, status: str) -> None: - # During shutdown, just output silence - if self.stop_event.is_set(): - outdata[:] = 0 * outdata - return - - if self.output_queue.empty(): - pcm = np.ascontiguousarray(indata, dtype=np.int16) - self.input_queue.put(pcm.tobytes()) - outdata[:] = dither - else: - try: - audio_chunk = self.output_queue.get_nowait() - if isinstance(audio_chunk, np.ndarray): - outdata[:] = audio_chunk[:, np.newaxis] - elif audio_chunk == AUDIO_RESPONSE_DONE: - self.should_listen.set() - logger.debug("Response complete, listening re-enabled") - outdata[:] = 0 * outdata - else: - outdata[:] = 0 * outdata - except Exception: - outdata[:] = 0 * outdata - - logger.debug("Available devices:") - logger.debug(sd.query_devices()) - with sd.Stream( - samplerate=16000, - dtype="int16", - channels=1, - callback=callback, - blocksize=self.list_play_chunk_size, - ): - logger.info("Starting local audio stream") - while not self.stop_event.is_set(): - time.sleep(0.001) - print("Stopping recording") diff --git a/runtime/src/speech_to_speech/connections/socket_receiver.py b/runtime/src/speech_to_speech/connections/socket_receiver.py deleted file mode 100644 index 21d7563a7e822062bacd84f75a65da72d9fa8329..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/connections/socket_receiver.py +++ /dev/null @@ -1,85 +0,0 @@ -import logging -import socket -import time -from queue import Queue -from threading import Event - -from rich.console import Console - -from speech_to_speech.pipeline.messages import PIPELINE_END -from speech_to_speech.pipeline.queue_types import AudioInItem - -logger = logging.getLogger(__name__) - -console = Console() - -# If should_listen stays cleared for longer than this, something in the -# pipeline probably failed (LLM exception, TTS crash, etc.). Re-enable -# listening so the user isn't permanently locked out. -SHOULD_LISTEN_TIMEOUT_S = 30.0 - - -class SocketReceiver: - """ - Handles reception of the audio packets from the client. - """ - - def __init__( - self, - stop_event: Event, - queue_out: Queue[AudioInItem], - should_listen: Event, - host: str = "0.0.0.0", - port: int = 12345, - chunk_size: int = 1024, - ) -> None: - self.stop_event = stop_event - self.queue_out = queue_out - self.should_listen = should_listen - self.chunk_size = chunk_size - self.host = host - self.port = port - - def receive_full_chunk(self, conn: socket.socket, chunk_size: int) -> bytes | None: - data = b"" - while len(data) < chunk_size: - packet = conn.recv(chunk_size - len(data)) - if not packet: - # connection closed - return None - data += packet - return data - - def run(self) -> None: - self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) - self.socket.bind((self.host, self.port)) - self.socket.listen(1) - logger.info("Receiver waiting to be connected...") - self.conn, _ = self.socket.accept() - logger.info("receiver connected") - - self.should_listen.set() - listen_cleared_at = None - while not self.stop_event.is_set(): - audio_chunk = self.receive_full_chunk(self.conn, self.chunk_size) - if audio_chunk is None: - # connection closed - self.queue_out.put(PIPELINE_END) - break - if self.should_listen.is_set(): - self.queue_out.put(audio_chunk) - listen_cleared_at = None - else: - # Track how long should_listen has been cleared - if listen_cleared_at is None: - listen_cleared_at = time.monotonic() - elif time.monotonic() - listen_cleared_at > SHOULD_LISTEN_TIMEOUT_S: - logger.warning( - "should_listen has been cleared for %.0fs — pipeline may be stuck, re-enabling listening", - SHOULD_LISTEN_TIMEOUT_S, - ) - self.should_listen.set() - listen_cleared_at = None - self.conn.close() - logger.info("Receiver closed") diff --git a/runtime/src/speech_to_speech/connections/socket_sender.py b/runtime/src/speech_to_speech/connections/socket_sender.py deleted file mode 100644 index c2b465b98edf3a0d7c8bb9f24c0142992b85738a..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/connections/socket_sender.py +++ /dev/null @@ -1,69 +0,0 @@ -import logging -import socket -from queue import Empty, Queue -from threading import Event - -import numpy as np -from rich.console import Console - -from speech_to_speech.pipeline.control import PipelineControlMessage -from speech_to_speech.pipeline.messages import AUDIO_RESPONSE_DONE, PIPELINE_END -from speech_to_speech.pipeline.queue_types import AudioOutItem - -logger = logging.getLogger(__name__) - -console = Console() - - -class SocketSender: - """ - Handles sending generated audio packets to the clients. - """ - - def __init__( - self, - stop_event: Event, - queue_in: Queue[AudioOutItem], - should_listen: Event, - host: str = "0.0.0.0", - port: int = 12346, - ) -> None: - self.stop_event = stop_event - self.queue_in = queue_in - self.should_listen = should_listen - self.host = host - self.port = port - - def run(self) -> None: - self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) - self.socket.bind((self.host, self.port)) - self.socket.listen(1) - logger.info("Sender waiting to be connected...") - self.conn, _ = self.socket.accept() - logger.info("sender connected") - - while not self.stop_event.is_set(): - try: - audio_chunk = self.queue_in.get(timeout=0.5) - except Empty: - continue - if isinstance(audio_chunk, PipelineControlMessage): - continue - if isinstance(audio_chunk, bytes) and audio_chunk == AUDIO_RESPONSE_DONE: - self.should_listen.set() - continue - payload: bytes - if isinstance(audio_chunk, bytes): - payload = audio_chunk - elif isinstance(audio_chunk, np.ndarray): - payload = audio_chunk.tobytes() - elif hasattr(audio_chunk, "tobytes"): - payload = audio_chunk.tobytes() - else: - continue - self.conn.sendall(payload) - if isinstance(audio_chunk, bytes) and audio_chunk == PIPELINE_END: - break - self.conn.close() - logger.info("Sender closed") diff --git a/runtime/src/speech_to_speech/connections/websocket_streamer.py b/runtime/src/speech_to_speech/connections/websocket_streamer.py deleted file mode 100644 index a126fd958a58917da582b1c4b53e24da174613f9..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/connections/websocket_streamer.py +++ /dev/null @@ -1,238 +0,0 @@ -import asyncio -import json -import logging -from queue import Empty, Queue -from threading import Event -from typing import Any - -import numpy as np -from websockets.asyncio.server import ServerConnection - -from speech_to_speech.pipeline.control import SESSION_END, PipelineControlMessage, is_control_message -from speech_to_speech.pipeline.events import PipelineEvent -from speech_to_speech.pipeline.messages import AUDIO_RESPONSE_DONE, PIPELINE_END -from speech_to_speech.pipeline.queue_types import AudioInItem, AudioOutItem, TextEventItem - -logger = logging.getLogger(__name__) - - -class WebSocketStreamer: - """ - Handles bidirectional audio streaming over WebSocket. - - Receives audio from clients and puts it in the input_queue. - Sends audio from the output_queue to clients. - Sends text messages (transcripts/tools) from text_output_queue to clients. - """ - - def __init__( - self, - stop_event: Event, - input_queue: Queue[AudioInItem], - output_queue: Queue[AudioOutItem], - should_listen: Event, - text_output_queue: Queue[TextEventItem] | None = None, - host: str = "0.0.0.0", - port: int = 8765, - ) -> None: - self.stop_event = stop_event - self.input_queue = input_queue # clients -> VAD - self.output_queue = output_queue # TTS -> clients - self.text_output_queue = text_output_queue # Text messages -> clients - self.should_listen = should_listen - self.host = host - self.port = port - self.clients: set[ServerConnection] = set() - self.loop: asyncio.AbstractEventLoop | None = None - self.server: Any = None - - def run(self) -> None: - """Run the WebSocket server (called from a thread).""" - self.loop = asyncio.new_event_loop() - asyncio.set_event_loop(self.loop) - - try: - self.loop.run_until_complete(self._run_server()) - except Exception as e: - logger.error(f"WebSocket server error: {e}") - finally: - self.loop.close() - - async def _run_server(self) -> None: - """Main async server loop.""" - import websockets - - logger.info(f"WebSocket server starting on ws://{self.host}:{self.port}") - - self.server = await websockets.serve( - self._handle_client, - self.host, - self.port, - ) - - logger.info("WebSocket server ready, waiting for connections...") - - # Start the sender task - sender_task = asyncio.create_task(self._send_loop()) - - # Wait until stop_event is set - while not self.stop_event.is_set(): - await asyncio.sleep(0.1) - - # Cleanup - sender_task.cancel() - try: - await sender_task - except asyncio.CancelledError: - pass - - # Close all clients - for client in list(self.clients): - try: - await client.close() - except Exception: - pass - - self.server.close() - await self.server.wait_closed() - logger.info("WebSocket server closed") - - async def _handle_client(self, websocket: ServerConnection) -> None: - """Handle a single WebSocket client connection.""" - client_id = id(websocket) - logger.info(f"Client {client_id} connected") - self.clients.add(websocket) - recv_buffer = bytearray() - - # Enable listening when first client connects - if len(self.clients) == 1: - # Drain edge queues so stale data from a previous session doesn't - # leak into the new one (SESSION_END may not have flushed everything). - for q in (self.output_queue, self.text_output_queue): - if q is not None: - while not q.empty(): - try: - q.get_nowait() - except Empty: - break - self.should_listen.set() - logger.debug("Listening enabled, edge queues drained (should_listen.set())") - - try: - logger.debug(f"Client {client_id}: Starting message receive loop") - async for message in websocket: - if isinstance(message, bytes): - logger.debug(f"Client {client_id}: Received {len(message)} bytes of audio") - if self.should_listen.is_set(): - # Split into 512-sample (1024 bytes) chunks for VAD. - # Keep a per-client remainder buffer so no samples are dropped - # when WebSocket frame boundaries are not aligned. - chunk_size_bytes = 512 * 2 # 512 samples * 2 bytes per int16 - recv_buffer.extend(message) - num_chunks = 0 - while len(recv_buffer) >= chunk_size_bytes: - chunk = bytes(recv_buffer[:chunk_size_bytes]) - del recv_buffer[:chunk_size_bytes] - self.input_queue.put(chunk) - num_chunks += 1 - logger.debug(f"Client {client_id}: Queued {num_chunks} chunks for processing") - else: - logger.debug(f"Client {client_id}: Skipping audio (should_listen not set)") - - except Exception as e: - logger.error(f"Client {client_id} error: {type(e).__name__}: {e}", exc_info=True) - finally: - self.clients.discard(websocket) - logger.info(f"Client {client_id} disconnected (finally block)") - - if len(self.clients) == 0: - logger.debug("Last WebSocket client disconnected, ending session") - self.input_queue.put(SESSION_END) - - async def _send_loop(self) -> None: - """Send audio and text from queues to all connected clients.""" - # Buffer audio until we have at least 100ms worth (3200 bytes = 1600 samples at 16kHz int16) - MIN_AUDIO_BYTES = 3200 - audio_buffer = bytearray() - - while not self.stop_event.is_set(): - try: - # Check for audio - try: - audio_chunk = self.output_queue.get_nowait() - if isinstance(audio_chunk, bytes) and audio_chunk == PIPELINE_END: - if audio_buffer and self.clients: - data = bytes(audio_buffer) - audio_buffer.clear() - await asyncio.gather( - *[client.send(data) for client in self.clients], - return_exceptions=True, - ) - break - if isinstance(audio_chunk, bytes) and audio_chunk == AUDIO_RESPONSE_DONE: - if audio_buffer and self.clients: - data = bytes(audio_buffer) - audio_buffer.clear() - await asyncio.gather( - *[client.send(data) for client in self.clients], - return_exceptions=True, - ) - self.should_listen.set() - logger.debug("Response complete, listening re-enabled") - continue - if is_control_message(audio_chunk, SESSION_END.kind): - audio_buffer.clear() - continue - - if isinstance(audio_chunk, PipelineControlMessage): - continue - - if self.clients: - chunk_bytes: bytes - if isinstance(audio_chunk, bytes): - chunk_bytes = audio_chunk - elif isinstance(audio_chunk, np.ndarray): - chunk_bytes = audio_chunk.tobytes() - elif hasattr(audio_chunk, "tobytes"): - chunk_bytes = audio_chunk.tobytes() - else: - continue - audio_buffer.extend(chunk_bytes) - - if len(audio_buffer) >= MIN_AUDIO_BYTES: - data = bytes(audio_buffer) - audio_buffer.clear() - logger.debug(f"Sending {len(data)} bytes of audio to {len(self.clients)} client(s)") - await asyncio.gather( - *[client.send(data) for client in self.clients], return_exceptions=True - ) - except Empty: - # Flush any buffered audio when queue is empty - if audio_buffer and self.clients: - data = bytes(audio_buffer) - audio_buffer.clear() - logger.debug(f"Flushing {len(data)} bytes of audio to {len(self.clients)} client(s)") - await asyncio.gather(*[client.send(data) for client in self.clients], return_exceptions=True) - - # Check for text/tool messages - if self.text_output_queue: - try: - text_message = self.text_output_queue.get_nowait() - if self.clients: - if isinstance(text_message, PipelineEvent): - payload = text_message.model_dump() - await asyncio.gather( - *[client.send(json.dumps(payload)) for client in self.clients], - return_exceptions=True, - ) - elif isinstance(text_message, (PipelineControlMessage, bytes)): - continue - except Empty: - pass - - await asyncio.sleep(0.01) - - except asyncio.CancelledError: - break - except Exception as e: - logger.error(f"Send loop error: {e}") diff --git a/runtime/src/speech_to_speech/pipeline/__init__.py b/runtime/src/speech_to_speech/pipeline/__init__.py deleted file mode 100644 index a6224b45130474b071d4173fe339634ed90e1ab5..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/pipeline/__init__.py +++ /dev/null @@ -1,57 +0,0 @@ -from speech_to_speech.pipeline.cancel_scope import CancelScope -from speech_to_speech.pipeline.control import ( - SESSION_END, - ControlKind, - PipelineControlMessage, - is_control_message, -) -from speech_to_speech.pipeline.events import ( - AssistantTextEvent, - AudioInputCompletedEvent, - PartialTranscriptionEvent, - PipelineEvent, - SpeechStartedEvent, - SpeechStoppedEvent, - TokenUsageEvent, - TranscriptionCompletedEvent, -) -from speech_to_speech.pipeline.messages import ( - AUDIO_RESPONSE_DONE, - PIPELINE_END, - EndOfResponse, - GenerateResponseRequest, - LLMResponseChunk, - PartialTranscription, - PipelineMessage, - TokenUsage, - Transcription, - TTSInput, - VADAudio, -) - -__all__ = [ - "AUDIO_RESPONSE_DONE", - "AssistantTextEvent", - "AudioInputCompletedEvent", - "CancelScope", - "ControlKind", - "EndOfResponse", - "GenerateResponseRequest", - "LLMResponseChunk", - "PIPELINE_END", - "PartialTranscription", - "PartialTranscriptionEvent", - "PipelineControlMessage", - "PipelineEvent", - "PipelineMessage", - "SESSION_END", - "SpeechStartedEvent", - "SpeechStoppedEvent", - "TTSInput", - "TokenUsage", - "TokenUsageEvent", - "Transcription", - "TranscriptionCompletedEvent", - "VADAudio", - "is_control_message", -] diff --git a/runtime/src/speech_to_speech/pipeline/cancel_scope.py b/runtime/src/speech_to_speech/pipeline/cancel_scope.py deleted file mode 100644 index d731ac1fc671fb445ce496b155ce903dd86438bf..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/pipeline/cancel_scope.py +++ /dev/null @@ -1,64 +0,0 @@ -class CancelScope: - """Unified cancellation signal for the speech-to-speech pipeline. - - Uses a generation counter so pipeline threads (LLM, TTS) can detect - cancellation without brief-pulse timing games, and an internal - ``discarding`` flag so the async send loop can drop stale output. - - Thread safety: one writer (the asyncio router thread) and multiple - readers (pipeline handler threads). Python's GIL makes int/bool - reads and writes atomic at the bytecode level, so no lock is needed. - """ - - def __init__(self) -> None: - self._gen: int = 0 - self._discarding: bool = False - self._discarded_generation: int | None = None - - @property - def generation(self) -> int: - """Current generation number. Pipeline threads capture this at - the start of each response and compare with ``is_stale``.""" - return self._gen - - def cancel(self) -> None: - """Cancel the current response. - - Increments the generation (so pipeline threads see their captured - generation as stale) and enables the send-loop discard guard. - """ - # prevent overflow... after 4 billion generations, we'll wrap around xD... - self._discarded_generation = self._gen - self._gen = (self._gen + 1) & 0xFFFFFFFF - self._discarding = True - - def response_done(self, generation: int | None = None) -> None: - """Pipeline acknowledged completion. Clears the discard guard.""" - if ( - generation is not None - and self._discarded_generation is not None - and generation not in {self._discarded_generation, self._gen} - ): - return - self._discarding = False - self._discarded_generation = None - - def new_response(self) -> None: - """An explicit ``response.create`` starts a new response. - Clears the discard guard.""" - self._discarding = False - self._discarded_generation = None - - def is_stale(self, gen: int) -> bool: - """Return True if *gen* has been superseded by a ``cancel`` call.""" - return gen != self._gen - - @property - def discarding(self) -> bool: - """Whether the send loop should silently drop stale output.""" - return self._discarding - - def reset(self) -> None: - """Clear discard state (e.g. on new session connect).""" - self._discarding = False - self._discarded_generation = None diff --git a/runtime/src/speech_to_speech/pipeline/control.py b/runtime/src/speech_to_speech/pipeline/control.py deleted file mode 100644 index ce8ed3b408925486eb9b833f9c4dcdc688d0906c..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/pipeline/control.py +++ /dev/null @@ -1,26 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass -from enum import Enum - - -class ControlKind(str, Enum): - """Strongly-typed kinds for :class:`PipelineControlMessage`.""" - - SESSION_END = "session_end" - - -@dataclass(frozen=True) -class PipelineControlMessage: - kind: ControlKind - # Session that enqueued the message, when known. Lets the pooled realtime - # send loop ignore a SESSION_END from a force-released session so it can't - # satisfy the drain wait of the session that claimed the unit afterwards. - session_id: str | None = None - - -SESSION_END = PipelineControlMessage(ControlKind.SESSION_END) - - -def is_control_message(message: object, kind: ControlKind | None = None) -> bool: - return isinstance(message, PipelineControlMessage) and (kind is None or message.kind == kind) diff --git a/runtime/src/speech_to_speech/pipeline/events.py b/runtime/src/speech_to_speech/pipeline/events.py deleted file mode 100644 index d380a83329b87589100a6b84635b05646e27282e..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/pipeline/events.py +++ /dev/null @@ -1,117 +0,0 @@ -"""Typed events flowing on ``text_output_queue``. - -These are internal pipeline events produced by VAD, TranscriptionNotifier, and -LMOutputProcessor, consumed by the realtime ``WebSocketRouter`` send-loop and -``RealtimeService.dispatch_pipeline_event``. They replace the raw ``dict`` -literals that were previously put on the queue. -""" - -from __future__ import annotations - -from typing import Literal, Optional - -import numpy as np -from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall -from pydantic import BaseModel, ConfigDict, Field - - -class PipelineEvent(BaseModel): - """Base for all text_output_queue events. - - The ``type`` field mirrors the former dict ``"type"`` key and acts as a - Pydantic discriminator. - """ - - type: str - - -# ── VAD events ──────────────────────────────────────────────────────── - - -class SpeechStartedEvent(PipelineEvent): - type: Literal["speech_started"] = "speech_started" - audio_start_ms: int = 0 - turn_id: str | None = None - turn_revision: int | None = None - reopened: bool = False - interrupt_response: bool = Field(default=True, exclude=True) - - -class SpeechStoppedEvent(PipelineEvent): - type: Literal["speech_stopped"] = "speech_stopped" - duration_s: float = 0.0 - audio_end_ms: int = 0 - turn_id: str | None = None - turn_revision: int | None = None - - -# ── Transcription events (TranscriptionNotifier) ───────────────────── - - -class PartialTranscriptionEvent(PipelineEvent): - type: Literal["partial_transcription"] = "partial_transcription" - delta: str - turn_id: str | None = None - turn_revision: int | None = None - - -class TranscriptionCompletedEvent(PipelineEvent): - type: Literal["transcription_completed"] = "transcription_completed" - transcript: str - language_code: Optional[str] = None - turn_id: str | None = None - turn_revision: int | None = None - speech_stopped_at_s: float | None = Field(default=None, exclude=True) - - -# ── Direct audio events (AudioInputNotifier) ───────────────────────── - - -class AudioInputCompletedEvent(PipelineEvent): - """Final VAD audio awaiting realtime response bookkeeping.""" - - model_config = ConfigDict(arbitrary_types_allowed=True) - - type: Literal["audio_input_completed"] = "audio_input_completed" - audio: np.ndarray = Field(exclude=True) - audio_sample_rate: int = 16000 - audio_duration_s: float = 0.0 - turn_id: str | None = None - turn_revision: int | None = None - speech_stopped_at_s: float | None = Field(default=None, exclude=True) - - -# ── LLM output events (LMOutputProcessor) ──────────────────────────── - - -class AssistantTextEvent(PipelineEvent): - type: Literal["assistant_text"] = "assistant_text" - text: str - tools: list[ResponseFunctionToolCall] = Field(default_factory=list) - turn_id: str | None = None - turn_revision: int | None = None - # Response generation that produced this text, mirroring AudioOutput. Lets the - # send loop discard stale assistant text by the same generation-aware rule as - # audio, instead of blanket-dropping while cancel_scope.discarding is set. - cancel_generation: int | None = None - - -class TokenUsageEvent(PipelineEvent): - type: Literal["token_usage"] = "token_usage" - input_tokens: int = 0 - output_tokens: int = 0 - turn_id: str | None = None - turn_revision: int | None = None - - -class ResponseFailedEvent(PipelineEvent): - """Signals that a response could not be generated (e.g. invalid out-of-band input). - - Dispatched to the service so it can close the in-progress response with - ``status="failed"`` instead of the usual ``completed``. - """ - - type: Literal["response_failed"] = "response_failed" - message: str = "" - turn_id: str | None = None - turn_revision: int | None = None diff --git a/runtime/src/speech_to_speech/pipeline/handler_types.py b/runtime/src/speech_to_speech/pipeline/handler_types.py deleted file mode 100644 index 8462a85510bd45aac0a19ad1f889348308026655..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/pipeline/handler_types.py +++ /dev/null @@ -1,44 +0,0 @@ -"""Convenience `TypeAlias` definitions for pipeline handler generics. - -These aliases keep handler declarations readable (e.g. `BaseHandler[STTIn, STTOut]`) -without hiding what actually flows between stages (see `pipeline/messages.py` and -`pipeline/queue_types.py` for the full story, including control + sentinel items). -""" - -from __future__ import annotations - -# ruff: noqa: I001 - -from typing import TypeAlias - -import numpy as np - -from speech_to_speech.pipeline.messages import ( - AudioOutput, - EndOfResponse, - GenerateResponseRequest, - LLMResponseChunk, - PartialTranscription, - TTSInput, - TokenUsage, - Transcription, - VADAudio, -) - -from speech_to_speech.api.openai_realtime.runtime_config import RuntimeConfig - -# ── VAD stage ───────────────────────────────────────────────────────── -VADIn: TypeAlias = bytes | tuple[bytes, RuntimeConfig] -VADOut: TypeAlias = VADAudio - -# ── STT stage ───────────────────────────────────────────────────────── -STTIn: TypeAlias = VADAudio -STTOut: TypeAlias = PartialTranscription | Transcription - -# ── LLM stage ───────────────────────────────────────────────────────── -LLMIn: TypeAlias = GenerateResponseRequest -LLMOut: TypeAlias = LLMResponseChunk | TokenUsage | EndOfResponse - -# ── TTS stage ───────────────────────────────────────────────────────── -TTSIn: TypeAlias = TTSInput | EndOfResponse -TTSOut: TypeAlias = bytes | np.ndarray | AudioOutput diff --git a/runtime/src/speech_to_speech/pipeline/log_context.py b/runtime/src/speech_to_speech/pipeline/log_context.py deleted file mode 100644 index f9c5501689b986c71be18315c23646fd93b951a4..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/pipeline/log_context.py +++ /dev/null @@ -1,25 +0,0 @@ -"""Per-pipeline logging context. - -Each isolated pipeline unit sets the contextvar at thread / asyncio-task entry, -and every log record emitted from that thread/task is automatically tagged with -a `[pipeline N] ` prefix via PipelineLogFilter. Non-pipeline-bound code emits -records with an empty prefix. - -Works for both threading (handler threads) and asyncio (websocket route, send loops) -because contextvars are per-thread and per-task. -""" - -import contextvars -import logging -from typing import Optional - -pipeline_log_ctx: contextvars.ContextVar[Optional[int]] = contextvars.ContextVar("pipeline_index", default=None) - - -class PipelineLogFilter(logging.Filter): - """Inject `pipeline_prefix` into every LogRecord based on the current contextvar.""" - - def filter(self, record: logging.LogRecord) -> bool: - idx = pipeline_log_ctx.get() - record.pipeline_prefix = f"[pipeline {idx}] " if idx is not None else "" - return True diff --git a/runtime/src/speech_to_speech/pipeline/messages.py b/runtime/src/speech_to_speech/pipeline/messages.py deleted file mode 100644 index f62a1b0cdd4de175198c46e8313c011bd0d45467..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/pipeline/messages.py +++ /dev/null @@ -1,175 +0,0 @@ -"""Single source of truth for inter-component pipeline messages. - -Typed :class:`PipelineMessage` subclasses replace the ad-hoc tuples that -previously flowed between STT, LLM, LMOutputProcessor and TTS stages. -Binary sentinels carried on the audio/output queue are plain ``bytes`` -constants. -""" - -from __future__ import annotations - -from time import perf_counter -from typing import Final, Literal, Optional, TypeAlias - -import numpy as np -from openai.types.realtime.realtime_response_create_params import RealtimeResponseCreateParams -from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall -from pydantic import BaseModel, ConfigDict, Field - -from speech_to_speech.api.openai_realtime.runtime_config import RuntimeConfig - -# ── Base class ──────────────────────────────────────────────────────── - - -class PipelineMessage(BaseModel): - """Base for all typed pipeline messages. - - The ``tag`` field acts as a Pydantic discriminator so a ``Union`` of - subtypes can be validated from raw dicts when needed. - """ - - model_config = ConfigDict(arbitrary_types_allowed=True) - - tag: str - - -# ── VAD → STT ───────────────────────────────────────────────────────── - - -class VADAudio(PipelineMessage): - """Audio segment from VAD, with optional mode for realtime transcription.""" - - tag: Literal["vad_audio"] = "vad_audio" - audio: np.ndarray - runtime_config: RuntimeConfig | None = None - mode: Literal["progressive", "final"] | None = None - turn_id: str | None = None - turn_revision: int | None = None - created_at_s: float = Field(default_factory=perf_counter) - - -# ── STT → TranscriptionNotifier → LLM ──────────────────────────────── - - -class PartialTranscription(PipelineMessage): - """Live partial transcription (consumed by TranscriptionNotifier, not forwarded to LLM).""" - - tag: Literal["partial_transcription"] = "partial_transcription" - text: str - turn_id: str | None = None - turn_revision: int | None = None - - -class Transcription(PipelineMessage): - """Final transcription result.""" - - tag: Literal["transcription"] = "transcription" - text: str - language_code: Optional[str] = None - turn_id: str | None = None - turn_revision: int | None = None - speech_stopped_at_s: float | None = None - - -# ── LLM → LMOutputProcessor ────────────────────────────────────────── - - -class LLMResponseChunk(PipelineMessage): - """One sentence/chunk of the LLM response.""" - - tag: Literal["llm_response_chunk"] = "llm_response_chunk" - text: str - language_code: Optional[str] = None - tools: list[ResponseFunctionToolCall] = Field(default_factory=list) - runtime_config: RuntimeConfig | None = None - response: RealtimeResponseCreateParams | None = None - turn_id: str | None = None - turn_revision: int | None = None - speech_stopped_at_s: float | None = None - cancel_generation: int | None = None - - -class TokenUsage(PipelineMessage): - """Token count report (side-channel, not forwarded to TTS).""" - - tag: Literal["token_usage"] = "token_usage" - input_tokens: int - output_tokens: int - turn_id: str | None = None - turn_revision: int | None = None - - -class EndOfResponse(PipelineMessage): - """Sentinel marking the end of a response. - - ``error`` is set when generation could not start (e.g. an out-of-band - response whose ``input`` failed validation); the output processor turns it - into a ``response.done(status="failed")`` while still closing the response - normally for pipeline cleanup. - """ - - tag: Literal["end_of_response"] = "end_of_response" - turn_id: str | None = None - turn_revision: int | None = None - cancel_generation: int | None = None - error: str | None = None - - -# ── LMOutputProcessor → TTS ────────────────────────────────────────── - - -class TTSInput(PipelineMessage): - """Text to synthesize with per-response context.""" - - tag: Literal["tts_input"] = "tts_input" - text: str - language_code: Optional[str] = None - runtime_config: RuntimeConfig | None = None - response: RealtimeResponseCreateParams | None = None - turn_id: str | None = None - turn_revision: int | None = None - speech_stopped_at_s: float | None = None - cancel_generation: int | None = None - - -class AudioOutput(PipelineMessage): - """Audio queue item tagged with the response generation that produced it.""" - - tag: Literal["audio_output"] = "audio_output" - audio: bytes | np.ndarray - cancel_generation: int | None = None - - -# ── Realtime service → LLM ──────────────────────────────────────────── - - -class GenerateResponseRequest(PipelineMessage): - """Triggers LLM generation for a realtime session. - - Carries everything the LM handler needs to produce a response so it - never has to reach back into shared objects. ``runtime_config`` - holds the per-connection session config *and* the conversation chat; - ``response`` carries per-response overrides from ``response.create``. - Downstream handlers resolve each attribute by preferring the - per-response value over the session default. - """ - - tag: Literal["generate_response"] = "generate_response" - runtime_config: RuntimeConfig - response: RealtimeResponseCreateParams | None = None - audio: np.ndarray | None = None - audio_sample_rate: int = 16000 - language_code: Optional[str] = None - turn_id: str | None = None - turn_revision: int | None = None - speech_stopped_at_s: float | None = None - - -# ── Binary sentinels (audio/output queue) ───────────────────────────── - -AUDIO_RESPONSE_DONE: Final[bytes] = b"__RESPONSE_DONE__" -PIPELINE_END: Final[bytes] = b"END" - -PipelineEndSentinel: TypeAlias = Literal[b"END"] -AudioResponseDoneSentinel: TypeAlias = Literal[b"__RESPONSE_DONE__"] -SentinelMessage: TypeAlias = PipelineEndSentinel | AudioResponseDoneSentinel diff --git a/runtime/src/speech_to_speech/pipeline/queue_types.py b/runtime/src/speech_to_speech/pipeline/queue_types.py deleted file mode 100644 index 27bc27a038d8e5f1dd0ccc7d14b730f3ddb80f45..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/pipeline/queue_types.py +++ /dev/null @@ -1,46 +0,0 @@ -"""Typed aliases for items flowing through pipeline queues. - -This module centralizes queue payload unions so the rest of the codebase can -reference a small set of stable types instead of repeating large unions. -""" - -from __future__ import annotations - -# ruff: noqa: I001 - -from typing import TypeAlias - -import numpy as np - -from speech_to_speech.pipeline.control import PipelineControlMessage -from speech_to_speech.pipeline.events import PipelineEvent -from speech_to_speech.pipeline.handler_types import LLMIn, LLMOut, STTOut, TTSIn, VADIn, VADOut -from speech_to_speech.pipeline.messages import AudioOutput - -# Use plain ``bytes`` for sentinel values on queues (``PIPELINE_END``, ``AUDIO_RESPONSE_DONE``). -# ``Queue`` is invariant; ``Literal[b"END"]`` is not accepted where ``bytes`` is required. -PipelineInternalItem: TypeAlias = PipelineControlMessage | bytes - -# Audio chunks coming from IO (socket/websocket/mic) into VAD. -AudioInItem: TypeAlias = VADIn | PipelineControlMessage - -# Audio segments flowing from VAD to STT. -VADOutItem: TypeAlias = VADOut | PipelineInternalItem - -# STT output (partial + final) flowing to TranscriptionNotifier. -STTOutItem: TypeAlias = STTOut | PipelineInternalItem - -# Text prompts flowing into LLM. -TextPromptItem: TypeAlias = LLMIn | PipelineInternalItem - -# LLM outputs flowing into LMOutputProcessor. -LMOutItem: TypeAlias = LLMOut | PipelineInternalItem - -# Inputs flowing into TTS. -TTSInItem: TypeAlias = TTSIn | PipelineInternalItem - -# Audio outputs flowing to speakers / client (includes sentinels as bytes). -AudioOutItem: TypeAlias = bytes | np.ndarray | AudioOutput | PipelineControlMessage - -# Side-channel text events sent to websocket/realtime clients. -TextEventItem: TypeAlias = PipelineEvent | PipelineInternalItem diff --git a/runtime/src/speech_to_speech/pipeline/speculative_turns.py b/runtime/src/speech_to_speech/pipeline/speculative_turns.py deleted file mode 100644 index 3f0ad0b597f54ba852befd0fa63dc86a71d15670..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/pipeline/speculative_turns.py +++ /dev/null @@ -1,418 +0,0 @@ -from __future__ import annotations - -import logging -import time -from collections import OrderedDict -from dataclasses import dataclass -from threading import Condition - -logger = logging.getLogger(__name__) - - -@dataclass(frozen=True) -class _PendingReopen: - base_revision: int - candidate_revision: int - - -@dataclass(frozen=True) -class _ReopenGrace: - revision: int - deadline: float - - -class SpeculativeTurnTracker: - """Thread-safe revision tracker for raw-audio speculative turns.""" - - _PENDING_REOPEN_WAIT_TIMEOUT_S = 2.0 - _MAX_TRACKED_TURNS = 2048 - - def __init__(self, max_tracked_turns: int = _MAX_TRACKED_TURNS) -> None: - self._condition = Condition() - self._max_tracked_turns = max_tracked_turns - self._latest_revision: OrderedDict[str, int] = OrderedDict() - self._committed_revision: dict[str, int] = {} - self._pending_reopen: dict[str, _PendingReopen] = {} - self._reopen_grace: dict[str, _ReopenGrace] = {} - - def observe(self, turn_id: str | None, revision: int | None) -> None: - if turn_id is None or revision is None: - return - with self._condition: - current = self._latest_revision.get(turn_id, -1) - if revision > current: - self._latest_revision[turn_id] = revision - self._latest_revision.move_to_end(turn_id) - self._prune_tracked_turns() - logger.debug("Observed speculative turn %s revision %d", turn_id, revision) - self._condition.notify_all() - - def is_latest(self, turn_id: str | None, revision: int | None) -> bool: - if turn_id is None or revision is None: - return True - with self._condition: - return self._latest_revision.get(turn_id, revision) == revision - - def is_latest_after_pending_reopen(self, turn_id: str | None, revision: int | None) -> bool: - if turn_id is None or revision is None: - return True - with self._condition: - self._wait_for_pending_reopen_locked(turn_id, revision, self._PENDING_REOPEN_WAIT_TIMEOUT_S) - return self._latest_revision.get(turn_id, revision) == revision - - def try_is_latest_after_pending_reopen(self, turn_id: str | None, revision: int | None) -> bool | None: - """Non-blocking variant of ``is_latest_after_pending_reopen``. - - Returns ``None`` when a matching reopen candidate is still pending and - the caller should retry after it resolves. - """ - if turn_id is None or revision is None: - return True - with self._condition: - if self._has_pending_reopen_locked(turn_id, revision): - return None - return self._latest_revision.get(turn_id, revision) == revision - - def is_latest_after_reopen_grace(self, turn_id: str | None, revision: int | None) -> bool: - if turn_id is None or revision is None: - return True - with self._condition: - self._wait_for_reopen_gate_locked(turn_id, revision) - return self._latest_revision.get(turn_id, revision) == revision - - def try_is_latest_after_reopen_grace(self, turn_id: str | None, revision: int | None) -> bool | None: - if turn_id is None or revision is None: - return True - with self._condition: - if ( - self._has_pending_reopen_locked(turn_id, revision) - or self._reopen_grace_remaining_locked( - turn_id, - revision, - ) - > 0 - ): - return None - return self._latest_revision.get(turn_id, revision) == revision - - def commit_if_latest_after_pending_reopen(self, turn_id: str | None, revision: int | None) -> bool: - if turn_id is None or revision is None: - return True - with self._condition: - self._wait_for_pending_reopen_locked(turn_id, revision, self._PENDING_REOPEN_WAIT_TIMEOUT_S) - latest = self._latest_revision.get(turn_id, revision) - if revision != latest: - return False - self._committed_revision[turn_id] = revision - logger.debug("Committed speculative turn %s revision %d", turn_id, revision) - self._condition.notify_all() - return True - - def commit_if_latest_after_reopen_grace(self, turn_id: str | None, revision: int | None) -> bool: - if turn_id is None or revision is None: - return True - with self._condition: - self._wait_for_reopen_gate_locked(turn_id, revision) - latest = self._latest_revision.get(turn_id, revision) - if revision != latest: - return False - self._committed_revision[turn_id] = revision - logger.debug("Committed speculative turn %s revision %d", turn_id, revision) - self._condition.notify_all() - return True - - def try_commit_if_latest_after_pending_reopen(self, turn_id: str | None, revision: int | None) -> bool | None: - """Non-blocking variant of ``commit_if_latest_after_pending_reopen``. - - Returns ``None`` when a matching reopen candidate is still pending and - the caller should retry after it resolves. - """ - if turn_id is None or revision is None: - return True - with self._condition: - if self._has_pending_reopen_locked(turn_id, revision): - return None - latest = self._latest_revision.get(turn_id, revision) - if revision != latest: - return False - self._committed_revision[turn_id] = revision - logger.debug("Committed speculative turn %s revision %d", turn_id, revision) - self._condition.notify_all() - return True - - def try_commit_if_latest_after_reopen_grace(self, turn_id: str | None, revision: int | None) -> bool | None: - if turn_id is None or revision is None: - return True - with self._condition: - if ( - self._has_pending_reopen_locked(turn_id, revision) - or self._reopen_grace_remaining_locked( - turn_id, - revision, - ) - > 0 - ): - return None - latest = self._latest_revision.get(turn_id, revision) - if revision != latest: - return False - self._committed_revision[turn_id] = revision - logger.debug("Committed speculative turn %s revision %d", turn_id, revision) - self._condition.notify_all() - return True - - def has_pending_reopen(self, turn_id: str | None, revision: int | None) -> bool: - if turn_id is None or revision is None: - return False - with self._condition: - return self._has_pending_reopen_locked(turn_id, revision) - - def has_pending_reopen_or_grace(self, turn_id: str | None, revision: int | None) -> bool: - if turn_id is None or revision is None: - return False - with self._condition: - return ( - self._has_pending_reopen_locked(turn_id, revision) - or self._reopen_grace_remaining_locked( - turn_id, - revision, - ) - > 0 - ) - - def start_reopen_grace(self, turn_id: str | None, revision: int | None, grace_s: float) -> None: - if turn_id is None or revision is None or grace_s <= 0: - return - with self._condition: - if self._latest_revision.get(turn_id, revision) != revision: - return - if self._committed_revision.get(turn_id, -1) >= revision: - return - deadline = time.monotonic() + grace_s - existing = self._reopen_grace.get(turn_id) - if existing is None or existing.revision != revision or deadline > existing.deadline: - self._reopen_grace[turn_id] = _ReopenGrace(revision=revision, deadline=deadline) - logger.debug( - "Started speculative reopen grace for turn %s revision %d: %.0fms", - turn_id, - revision, - grace_s * 1000, - ) - self._condition.notify_all() - - def is_latest_after_stability_window( - self, - turn_id: str | None, - revision: int | None, - settle_s: float, - ) -> bool: - if turn_id is None or revision is None: - return True - if settle_s <= 0: - return self.is_latest_after_pending_reopen(turn_id, revision) - with self._condition: - deadline = time.monotonic() + settle_s - while self._latest_revision.get(turn_id, revision) == revision: - if self._has_pending_reopen_locked(turn_id, revision): - break - remaining = deadline - time.monotonic() - if remaining <= 0: - break - self._condition.wait(remaining) - self._wait_for_pending_reopen_locked(turn_id, revision, self._PENDING_REOPEN_WAIT_TIMEOUT_S) - return self._latest_revision.get(turn_id, revision) == revision - - def commit(self, turn_id: str | None, revision: int | None) -> None: - if turn_id is None or revision is None: - return - with self._condition: - pending = self._pending_reopen.get(turn_id) - if pending is not None and pending.base_revision == revision: - logger.debug( - "Deferring speculative turn %s revision %d commit while reopen is pending", turn_id, revision - ) - return - latest = self._latest_revision.get(turn_id, revision) - if revision == latest: - self._committed_revision[turn_id] = revision - logger.debug("Committed speculative turn %s revision %d", turn_id, revision) - self._condition.notify_all() - - def is_committed(self, turn_id: str | None, revision: int | None = None) -> bool: - if turn_id is None: - return False - with self._condition: - committed = self._committed_revision.get(turn_id) - if committed is None: - return False - return revision is None or committed >= revision - - def begin_reopen_candidate(self, turn_id: str | None, revision: int | None) -> int | None: - if turn_id is None or revision is None: - return None - with self._condition: - if self._committed_revision.get(turn_id, -1) >= revision: - return None - if self._latest_revision.get(turn_id, revision) != revision: - return None - - pending = self._pending_reopen.get(turn_id) - if pending is not None: - if pending.base_revision == revision: - return pending.candidate_revision - return None - - candidate_revision = revision + 1 - self._pending_reopen[turn_id] = _PendingReopen( - base_revision=revision, - candidate_revision=candidate_revision, - ) - logger.debug( - "Started speculative reopen candidate for turn %s revision %d -> %d", - turn_id, - revision, - candidate_revision, - ) - self._condition.notify_all() - return candidate_revision - - def confirm_reopen_candidate( - self, - turn_id: str | None, - base_revision: int | None, - candidate_revision: int | None, - ) -> bool: - if turn_id is None or base_revision is None or candidate_revision is None: - return False - with self._condition: - pending = self._pending_reopen.get(turn_id) - if ( - pending is None - or pending.base_revision != base_revision - or pending.candidate_revision != candidate_revision - ): - return False - if self._committed_revision.get(turn_id, -1) >= base_revision: - del self._pending_reopen[turn_id] - self._prune_tracked_turns() - self._condition.notify_all() - return False - if self._latest_revision.get(turn_id, base_revision) != base_revision: - del self._pending_reopen[turn_id] - self._prune_tracked_turns() - self._condition.notify_all() - return False - - self._latest_revision[turn_id] = candidate_revision - self._latest_revision.move_to_end(turn_id) - del self._pending_reopen[turn_id] - self._prune_tracked_turns() - logger.debug( - "Confirmed speculative reopen candidate for turn %s revision %d", - turn_id, - candidate_revision, - ) - self._condition.notify_all() - return True - - def cancel_reopen_candidate(self, turn_id: str | None, candidate_revision: int | None = None) -> None: - if turn_id is None: - return - with self._condition: - pending = self._pending_reopen.get(turn_id) - if pending is None: - return - if candidate_revision is not None and pending.candidate_revision != candidate_revision: - return - del self._pending_reopen[turn_id] - self._prune_tracked_turns() - logger.debug("Cancelled speculative reopen candidate for turn %s", turn_id) - self._condition.notify_all() - - def wait_for_pending_reopen( - self, - turn_id: str | None, - revision: int | None, - timeout_s: float = _PENDING_REOPEN_WAIT_TIMEOUT_S, - ) -> None: - if turn_id is None or revision is None: - return - with self._condition: - self._wait_for_pending_reopen_locked(turn_id, revision, timeout_s) - - def _has_pending_reopen_locked(self, turn_id: str, revision: int) -> bool: - pending = self._pending_reopen.get(turn_id) - return pending is not None and pending.base_revision == revision - - def _reopen_grace_remaining_locked(self, turn_id: str, revision: int) -> float: - grace = self._reopen_grace.get(turn_id) - if grace is None or grace.revision != revision: - return 0.0 - if self._latest_revision.get(turn_id, revision) != revision: - del self._reopen_grace[turn_id] - return 0.0 - remaining = grace.deadline - time.monotonic() - if remaining <= 0: - del self._reopen_grace[turn_id] - self._prune_tracked_turns() - return 0.0 - return remaining - - def _wait_for_reopen_gate_locked(self, turn_id: str, revision: int) -> None: - while self._latest_revision.get(turn_id, revision) == revision: - self._wait_for_pending_reopen_locked(turn_id, revision, self._PENDING_REOPEN_WAIT_TIMEOUT_S) - if self._latest_revision.get(turn_id, revision) != revision: - return - remaining = self._reopen_grace_remaining_locked(turn_id, revision) - if remaining <= 0: - return - logger.debug("Waiting for speculative reopen grace turn=%s rev=%s", turn_id, revision) - self._condition.wait(remaining) - - def _wait_for_pending_reopen_locked(self, turn_id: str, revision: int, timeout_s: float) -> None: - deadline = time.monotonic() + timeout_s - pending = self._pending_reopen.get(turn_id) - if pending is None or pending.base_revision != revision: - return - logger.debug("Waiting for pending speculative reopen turn=%s rev=%s", turn_id, revision) - while pending is not None and pending.base_revision == revision: - remaining = deadline - time.monotonic() - if remaining <= 0: - logger.warning("Timed out waiting for pending speculative reopen turn=%s rev=%s", turn_id, revision) - if self._pending_reopen.get(turn_id) == pending: - del self._pending_reopen[turn_id] - self._prune_tracked_turns() - self._condition.notify_all() - return - self._condition.wait(remaining) - pending = self._pending_reopen.get(turn_id) - - def _prune_tracked_turns(self) -> None: - if self._max_tracked_turns <= 0: - return - - self._drop_expired_reopen_graces_locked() - prunable_turn_ids = [ - turn_id - for turn_id in self._latest_revision - if turn_id not in self._pending_reopen and turn_id not in self._reopen_grace - ] - while len(prunable_turn_ids) > self._max_tracked_turns: - turn_id = prunable_turn_ids.pop(0) - self._latest_revision.pop(turn_id, None) - self._committed_revision.pop(turn_id, None) - self._reopen_grace.pop(turn_id, None) - - def _drop_expired_reopen_graces_locked(self) -> None: - now = time.monotonic() - for turn_id, grace in list(self._reopen_grace.items()): - if self._latest_revision.get(turn_id, grace.revision) != grace.revision or grace.deadline <= now: - del self._reopen_grace[turn_id] - - def reset(self) -> None: - with self._condition: - self._latest_revision.clear() - self._committed_revision.clear() - self._pending_reopen.clear() - self._reopen_grace.clear() - self._condition.notify_all() diff --git a/runtime/src/speech_to_speech/s2s_pipeline.py b/runtime/src/speech_to_speech/s2s_pipeline.py deleted file mode 100644 index aa6ee1c40a79ef1ba5f109f18d2d77e284dffd43..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/s2s_pipeline.py +++ /dev/null @@ -1,1099 +0,0 @@ -import argparse -import json -import logging -import os -import signal -import sys -from copy import copy, deepcopy -from dataclasses import dataclass -from pathlib import Path -from queue import Queue -from sys import platform -from threading import Event -from types import FrameType -from typing import Any, Optional - -import nltk -import torch -from openai.types.realtime import RealtimeSessionCreateRequest -from rich.console import Console -from transformers import HfArgumentParser - -from speech_to_speech.api.openai_realtime.pipeline_unit import PipelineUnit -from speech_to_speech.api.openai_realtime.runtime_config import RuntimeConfig -from speech_to_speech.arguments_classes.chat_completions_language_model_arguments import ( - ChatCompletionsLanguageModelHandlerArguments, -) -from speech_to_speech.arguments_classes.chat_tts_arguments import ChatTTSHandlerArguments -from speech_to_speech.arguments_classes.facebookmms_tts_arguments import FacebookMMSTTSHandlerArguments -from speech_to_speech.arguments_classes.faster_whisper_stt_arguments import ( - FasterWhisperSTTHandlerArguments, -) -from speech_to_speech.arguments_classes.kokoro_tts_arguments import KokoroTTSHandlerArguments -from speech_to_speech.arguments_classes.language_model_arguments import LanguageModelHandlerArguments -from speech_to_speech.arguments_classes.mlx_audio_whisper_arguments import ( - MLXAudioWhisperSTTHandlerArguments, -) -from speech_to_speech.arguments_classes.module_arguments import ModuleArguments -from speech_to_speech.arguments_classes.paraformer_stt_arguments import ParaformerSTTHandlerArguments -from speech_to_speech.arguments_classes.parakeet_tdt_arguments import ( - ParakeetTDTSTTHandlerArguments, -) -from speech_to_speech.arguments_classes.pocket_tts_arguments import PocketTTSHandlerArguments -from speech_to_speech.arguments_classes.qwen3_tts_arguments import Qwen3TTSHandlerArguments -from speech_to_speech.arguments_classes.responses_api_language_model_arguments import ( - ResponsesApiLanguageModelHandlerArguments, -) -from speech_to_speech.arguments_classes.socket_receiver_arguments import SocketReceiverArguments -from speech_to_speech.arguments_classes.socket_sender_arguments import SocketSenderArguments -from speech_to_speech.arguments_classes.vad_arguments import VADHandlerArguments -from speech_to_speech.arguments_classes.websocket_streamer_arguments import WebSocketStreamerArguments -from speech_to_speech.arguments_classes.whisper_stt_arguments import WhisperSTTHandlerArguments -from speech_to_speech.baseHandler import BaseHandler -from speech_to_speech.LLM.chat import Chat -from speech_to_speech.pipeline.cancel_scope import CancelScope -from speech_to_speech.pipeline.handler_types import LLMIn, LLMOut, STTIn, STTOut, TTSIn, TTSOut -from speech_to_speech.pipeline.queue_types import ( - AudioInItem, - AudioOutItem, - LMOutItem, - STTOutItem, - TextEventItem, - TextPromptItem, - TTSInItem, - VADOutItem, -) -from speech_to_speech.pipeline.speculative_turns import SpeculativeTurnTracker -from speech_to_speech.STT.transcription_notifier import TranscriptionNotifier -from speech_to_speech.utils.thread_manager import ThreadManager -from speech_to_speech.VAD.vad_handler import VADHandler - -# Ensure that the necessary NLTK resources are available -try: - nltk.data.find("tokenizers/punkt_tab") -except (LookupError, OSError): - nltk.download("punkt_tab") -try: - nltk.data.find("tokenizers/averaged_perceptron_tagger_eng") -except (LookupError, OSError): - nltk.download("averaged_perceptron_tagger_eng") - -# caching allows ~50% compilation time reduction -# see https://docs.google.com/document/d/1y5CRfMLdwEoF1nTk9q8qEu1mgMUuUtvhklPKJ2emLU8/edit#heading=h.o2asbxsrp1ma -CURRENT_DIR = Path(__file__).resolve().parent -os.environ["TORCHINDUCTOR_CACHE_DIR"] = os.path.join(CURRENT_DIR, "tmp") - -console = Console() -logger = logging.getLogger(__name__) -logging.getLogger("numba").setLevel(logging.WARNING) # quiet down numba logs - - -@dataclass -class ParsedArguments: - module_kwargs: ModuleArguments - socket_receiver_kwargs: SocketReceiverArguments - socket_sender_kwargs: SocketSenderArguments - websocket_streamer_kwargs: WebSocketStreamerArguments - vad_handler_kwargs: VADHandlerArguments - whisper_stt_handler_kwargs: WhisperSTTHandlerArguments - paraformer_stt_handler_kwargs: ParaformerSTTHandlerArguments - faster_whisper_stt_handler_kwargs: FasterWhisperSTTHandlerArguments - mlx_audio_whisper_stt_handler_kwargs: MLXAudioWhisperSTTHandlerArguments - parakeet_tdt_stt_handler_kwargs: ParakeetTDTSTTHandlerArguments - language_model_handler_kwargs: LanguageModelHandlerArguments - responses_api_language_model_handler_kwargs: ResponsesApiLanguageModelHandlerArguments - chat_tts_handler_kwargs: ChatTTSHandlerArguments - facebook_mms_tts_handler_kwargs: FacebookMMSTTSHandlerArguments - pocket_tts_handler_kwargs: PocketTTSHandlerArguments - kokoro_tts_handler_kwargs: KokoroTTSHandlerArguments - qwen3_tts_handler_kwargs: Qwen3TTSHandlerArguments - - -def rename_args(args: Any, prefix: str) -> None: - """ - Rename arguments by removing the prefix and prepares the gen_kwargs. - """ - gen_kwargs = {} - for key in copy(args.__dict__): - if key.startswith(prefix): - value = args.__dict__.pop(key) - new_key = key[len(prefix) + 1 :] # Remove prefix and underscore - if new_key.startswith("gen_"): - gen_kwargs[new_key[4:]] = value # Remove 'gen_' and add to dict - else: - args.__dict__[new_key] = value - - args.__dict__["gen_kwargs"] = gen_kwargs - - -def parse_arguments() -> ParsedArguments: - # Pre-parse to determine which LM backend is selected, so only one of the - # mutually exclusive LM argument classes is registered with HfArgumentParser - # (avoids duplicate field names from the shared LanguageModelBaseArguments base). - # chat-completions reuses the responses-api connection fields via subclassing. - _backend_lm_class = { - "responses-api": ResponsesApiLanguageModelHandlerArguments, - "chat-completions": ChatCompletionsLanguageModelHandlerArguments, - } - _is_json = len(sys.argv) == 2 and sys.argv[1].endswith(".json") - if _is_json: - with open(sys.argv[1]) as _f: - _backend = json.load(_f).get("llm_backend") - else: - _pre = argparse.ArgumentParser(add_help=False) - _pre.add_argument("--llm_backend", default="responses-api") - _backend = _pre.parse_known_args()[0].llm_backend - - _lm_class = _backend_lm_class.get(_backend, LanguageModelHandlerArguments) - logger.debug("LLM backend pre-parse: backend=%s, registering %s", _backend, _lm_class.__name__) - - parser = HfArgumentParser( - ( # type: ignore[arg-type] - ModuleArguments, - SocketReceiverArguments, - SocketSenderArguments, - WebSocketStreamerArguments, - VADHandlerArguments, - WhisperSTTHandlerArguments, - ParaformerSTTHandlerArguments, - FasterWhisperSTTHandlerArguments, - MLXAudioWhisperSTTHandlerArguments, - ParakeetTDTSTTHandlerArguments, - _lm_class, - ChatTTSHandlerArguments, - FacebookMMSTTSHandlerArguments, - PocketTTSHandlerArguments, - KokoroTTSHandlerArguments, - Qwen3TTSHandlerArguments, - ) - ) - - if _is_json: - parsed = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]), allow_extra_keys=True) - else: - parsed = parser.parse_args_into_dataclasses() - - # Build a {type: instance} lookup so field assignment is order-independent. - by_type: dict[type, Any] = {type(obj): obj for obj in parsed} - logger.debug("Parsed %d argument classes: %s", len(by_type), [t.__name__ for t in by_type]) - - return ParsedArguments( - module_kwargs=by_type[ModuleArguments], - socket_receiver_kwargs=by_type[SocketReceiverArguments], - socket_sender_kwargs=by_type[SocketSenderArguments], - websocket_streamer_kwargs=by_type[WebSocketStreamerArguments], - vad_handler_kwargs=by_type[VADHandlerArguments], - whisper_stt_handler_kwargs=by_type[WhisperSTTHandlerArguments], - paraformer_stt_handler_kwargs=by_type[ParaformerSTTHandlerArguments], - faster_whisper_stt_handler_kwargs=by_type[FasterWhisperSTTHandlerArguments], - mlx_audio_whisper_stt_handler_kwargs=by_type[MLXAudioWhisperSTTHandlerArguments], - parakeet_tdt_stt_handler_kwargs=by_type[ParakeetTDTSTTHandlerArguments], - language_model_handler_kwargs=by_type.get(LanguageModelHandlerArguments, LanguageModelHandlerArguments()), - # The OpenAI-compatible slot holds whichever class was registered: - # ChatCompletions... (a subclass) for chat-completions, else ResponsesApi.... - responses_api_language_model_handler_kwargs=by_type.get( - ChatCompletionsLanguageModelHandlerArguments, - by_type.get(ResponsesApiLanguageModelHandlerArguments, ResponsesApiLanguageModelHandlerArguments()), - ), - chat_tts_handler_kwargs=by_type[ChatTTSHandlerArguments], - facebook_mms_tts_handler_kwargs=by_type[FacebookMMSTTSHandlerArguments], - pocket_tts_handler_kwargs=by_type[PocketTTSHandlerArguments], - kokoro_tts_handler_kwargs=by_type[KokoroTTSHandlerArguments], - qwen3_tts_handler_kwargs=by_type[Qwen3TTSHandlerArguments], - ) - - -def setup_logger(log_level: str) -> None: - global logger - from speech_to_speech.pipeline.log_context import PipelineLogFilter - - logging.basicConfig( - level=log_level.upper(), - format="%(asctime)s - %(pipeline_prefix)s%(name)s - %(levelname)s - %(message)s", - ) - # Attach the filter to every existing handler so each LogRecord gets a - # `pipeline_prefix` attribute (matching the format string above). - pipeline_filter = PipelineLogFilter() - for h in logging.getLogger().handlers: - h.addFilter(pipeline_filter) - - logger = logging.getLogger(__name__) - - # torch compile logs - if log_level == "debug": - torch._logging.set_logs(graph_breaks=True, recompiles=True, cudagraphs=True) - - -MLX_DEFAULT_LM_MODEL = "mlx-community/Qwen3-4B-Instruct-2507-bf16" -TRANSFORMERS_DEFAULT_LM_MODEL = "Qwen/Qwen3-4B-Instruct-2507" - - -def optimal_mac_settings(mac_optimal_settings: bool, *handler_kwargs: Any) -> None: - if mac_optimal_settings: - for kwargs in handler_kwargs: - if hasattr(kwargs, "device"): - kwargs.device = "mps" - if hasattr(kwargs, "mode"): - kwargs.mode = "local" - if hasattr(kwargs, "stt"): - kwargs.stt = "parakeet-tdt" - if hasattr(kwargs, "llm_backend"): - kwargs.llm_backend = "mlx-lm" - if hasattr(kwargs, "tts"): - kwargs.tts = "qwen3" - if hasattr(kwargs, "model_name"): - if kwargs.model_name == TRANSFORMERS_DEFAULT_LM_MODEL: - kwargs.model_name = MLX_DEFAULT_LM_MODEL - - -def check_mac_settings(module_kwargs: ModuleArguments) -> None: - if platform == "darwin": - if module_kwargs.device == "cuda": - raise ValueError("Cannot use CUDA on macOS. Please set the device to 'cpu' or 'mps'.") - if module_kwargs.llm_backend != "mlx-lm": - logger.warning( - "For macOS users, it is recommended to use mlx-lm. You can activate it by passing --llm_backend mlx-lm." - ) - if module_kwargs.tts not in ("pocket", "kokoro", "qwen3"): - logger.warning( - "For macOS users, it is recommended to use qwen3 for TTS (pocket and kokoro are also valid options)." - ) - - -def overwrite_device_argument(common_device: Optional[str], *handler_kwargs: Any) -> None: - if common_device: - for kwargs in handler_kwargs: - if hasattr(kwargs, "llm_device"): - kwargs.llm_device = common_device - if hasattr(kwargs, "tts_device"): - kwargs.tts_device = common_device - if hasattr(kwargs, "stt_device"): - kwargs.stt_device = common_device - if hasattr(kwargs, "paraformer_stt_device"): - kwargs.paraformer_stt_device = common_device - if hasattr(kwargs, "facebook_mms_device"): - kwargs.facebook_mms_device = common_device - if hasattr(kwargs, "qwen3_tts_device"): - kwargs.qwen3_tts_device = common_device - - -def prepare_module_args(module_kwargs: ModuleArguments, *handler_kwargs: Any) -> None: - optimal_mac_settings(module_kwargs.local_mac_optimal_settings, module_kwargs, *handler_kwargs) - if module_kwargs.tts is None: - module_kwargs.tts = "qwen3" - if module_kwargs.stt == "none" and module_kwargs.llm_backend != "chat-completions": - raise ValueError("--stt none requires --llm_backend chat-completions for audio-input LLM requests.") - if platform == "darwin": - check_mac_settings(module_kwargs) - overwrite_device_argument(module_kwargs.device, *handler_kwargs) - - -def prepare_all_args( - module_kwargs: ModuleArguments, - whisper_stt_handler_kwargs: WhisperSTTHandlerArguments, - paraformer_stt_handler_kwargs: ParaformerSTTHandlerArguments, - faster_whisper_stt_handler_kwargs: FasterWhisperSTTHandlerArguments, - mlx_audio_whisper_stt_handler_kwargs: MLXAudioWhisperSTTHandlerArguments, - parakeet_tdt_stt_handler_kwargs: ParakeetTDTSTTHandlerArguments, - language_model_handler_kwargs: LanguageModelHandlerArguments, - responses_api_language_model_handler_kwargs: ResponsesApiLanguageModelHandlerArguments, - chat_tts_handler_kwargs: ChatTTSHandlerArguments, - facebook_mms_tts_handler_kwargs: FacebookMMSTTSHandlerArguments, - pocket_tts_handler_kwargs: PocketTTSHandlerArguments, - kokoro_tts_handler_kwargs: KokoroTTSHandlerArguments, - qwen3_tts_handler_kwargs: Qwen3TTSHandlerArguments, -) -> None: - prepare_module_args( - module_kwargs, - whisper_stt_handler_kwargs, - faster_whisper_stt_handler_kwargs, - paraformer_stt_handler_kwargs, - mlx_audio_whisper_stt_handler_kwargs, - parakeet_tdt_stt_handler_kwargs, - language_model_handler_kwargs, - responses_api_language_model_handler_kwargs, - chat_tts_handler_kwargs, - facebook_mms_tts_handler_kwargs, - pocket_tts_handler_kwargs, - kokoro_tts_handler_kwargs, - qwen3_tts_handler_kwargs, - ) - - rename_args(whisper_stt_handler_kwargs, "stt") - rename_args(faster_whisper_stt_handler_kwargs, "faster_whisper_stt") - rename_args(paraformer_stt_handler_kwargs, "paraformer_stt") - rename_args(mlx_audio_whisper_stt_handler_kwargs, "mlx_audio_whisper") - rename_args(parakeet_tdt_stt_handler_kwargs, "parakeet_tdt") - rename_args(language_model_handler_kwargs, "llm") - rename_args(responses_api_language_model_handler_kwargs, "responses_api") - rename_args(chat_tts_handler_kwargs, "chat_tts") - rename_args(facebook_mms_tts_handler_kwargs, "facebook_mms") - rename_args(pocket_tts_handler_kwargs, "pocket_tts") - rename_args(kokoro_tts_handler_kwargs, "kokoro") - rename_args(qwen3_tts_handler_kwargs, "qwen3_tts") - - -def initialize_queues_and_events() -> dict[str, Any]: - return { - "stop_event": Event(), - "should_listen": Event(), - "response_playing": Event(), - "cancel_scope": CancelScope(), - "recv_audio_chunks_queue": Queue[AudioInItem](), - "send_audio_chunks_queue": Queue[AudioOutItem](), - "spoken_prompt_queue": Queue[VADOutItem](), - "stt_output_queue": Queue[STTOutItem](), - "text_prompt_queue": Queue[TextPromptItem](), - "lm_response_queue": Queue[LMOutItem](), - "lm_processed_queue": Queue[TTSInItem](), # NEW: LLM -> LM processor -> TTS - "text_output_queue": Queue[TextEventItem](), # NEW: for text messages to WebSocket - } - - -def _build_pipeline_handlers( - *, - stop_event: Event, - should_listen: Event, - recv_audio_chunks_queue: Queue[AudioInItem], - spoken_prompt_queue: Queue[VADOutItem], - stt_output_queue: Queue[STTOutItem], - text_prompt_queue: Queue[TextPromptItem], - lm_response_queue: Queue[LMOutItem], - lm_processed_queue: Queue[TTSInItem], - send_audio_chunks_queue: Queue[AudioOutItem], - text_output_queue: Optional[Queue[TextEventItem]], - transcription_notifier_setup: dict[str, Any], - module_kwargs: ModuleArguments, - vad_handler_kwargs: VADHandlerArguments, - whisper_stt_handler_kwargs: WhisperSTTHandlerArguments, - faster_whisper_stt_handler_kwargs: FasterWhisperSTTHandlerArguments, - paraformer_stt_handler_kwargs: ParaformerSTTHandlerArguments, - mlx_audio_whisper_stt_handler_kwargs: MLXAudioWhisperSTTHandlerArguments, - parakeet_tdt_stt_handler_kwargs: ParakeetTDTSTTHandlerArguments, - language_model_handler_kwargs: LanguageModelHandlerArguments, - responses_api_language_model_handler_kwargs: ResponsesApiLanguageModelHandlerArguments, - chat_tts_handler_kwargs: ChatTTSHandlerArguments, - facebook_mms_tts_handler_kwargs: FacebookMMSTTSHandlerArguments, - pocket_tts_handler_kwargs: PocketTTSHandlerArguments, - kokoro_tts_handler_kwargs: KokoroTTSHandlerArguments, - qwen3_tts_handler_kwargs: Qwen3TTSHandlerArguments, - speculative_turns: SpeculativeTurnTracker | None = None, -) -> list[Any]: - """Build the shared handler chain: VAD → STT/AudioInput → LM → LMOutputProcessor → TTS. - - Callers own the queues, events, and any per-mode mutations to handler kwargs (cancel_scope, - text_output_queue, live transcription flags). `transcription_notifier_setup` differs by mode: - realtime mode passes a service-driven setup; non-realtime modes inject a RuntimeConfig. - """ - from speech_to_speech.LLM.lm_output_processor import LMOutputProcessor - - vad = VADHandler( - stop_event, - queue_in=recv_audio_chunks_queue, - queue_out=spoken_prompt_queue, - setup_args=(should_listen,), - setup_kwargs=vars(vad_handler_kwargs), - ) - - speech_input_handlers: list[Any] - if module_kwargs.stt == "none": - from speech_to_speech.LLM.audio_input_notifier import AudioInputNotifier - - speech_input_handlers = [ - AudioInputNotifier( - stop_event, - queue_in=spoken_prompt_queue, - queue_out=text_prompt_queue, - setup_kwargs={ - "runtime_config": transcription_notifier_setup.get("runtime_config"), - "should_listen": should_listen, - "sample_rate": vad_handler_kwargs.sample_rate, - "speculative_turns": speculative_turns, - "text_output_queue": text_output_queue, - }, - ) - ] - else: - transcription_notifier = TranscriptionNotifier( - stop_event, - queue_in=stt_output_queue, - queue_out=text_prompt_queue, # type: ignore[arg-type] - setup_kwargs=transcription_notifier_setup, - ) - - stt = get_stt_handler( - module_kwargs, - stop_event, - spoken_prompt_queue, - stt_output_queue, - speculative_turns, - whisper_stt_handler_kwargs, - faster_whisper_stt_handler_kwargs, - paraformer_stt_handler_kwargs, - mlx_audio_whisper_stt_handler_kwargs, - parakeet_tdt_stt_handler_kwargs, - ) - speech_input_handlers = [stt, transcription_notifier] - - lm = get_llm_handler( - module_kwargs, - stop_event, - text_prompt_queue, - lm_response_queue, - language_model_handler_kwargs, - responses_api_language_model_handler_kwargs, - ) - - lm_processor = LMOutputProcessor( - stop_event, - queue_in=lm_response_queue, - queue_out=lm_processed_queue, - setup_kwargs={"text_output_queue": text_output_queue, "speculative_turns": speculative_turns}, - ) - - tts = get_tts_handler( - module_kwargs, - stop_event, - lm_processed_queue, - send_audio_chunks_queue, - should_listen, - chat_tts_handler_kwargs, - facebook_mms_tts_handler_kwargs, - pocket_tts_handler_kwargs, - kokoro_tts_handler_kwargs, - qwen3_tts_handler_kwargs, - ) - - return [vad, *speech_input_handlers, lm, lm_processor, tts] - - -def _build_realtime_pipeline_unit( - *, - index: int, - stop_event: Event, - module_kwargs: ModuleArguments, - vad_handler_kwargs: VADHandlerArguments, - whisper_stt_handler_kwargs: WhisperSTTHandlerArguments, - faster_whisper_stt_handler_kwargs: FasterWhisperSTTHandlerArguments, - paraformer_stt_handler_kwargs: ParaformerSTTHandlerArguments, - mlx_audio_whisper_stt_handler_kwargs: MLXAudioWhisperSTTHandlerArguments, - parakeet_tdt_stt_handler_kwargs: ParakeetTDTSTTHandlerArguments, - language_model_handler_kwargs: LanguageModelHandlerArguments, - responses_api_language_model_handler_kwargs: ResponsesApiLanguageModelHandlerArguments, - chat_tts_handler_kwargs: ChatTTSHandlerArguments, - facebook_mms_tts_handler_kwargs: FacebookMMSTTSHandlerArguments, - pocket_tts_handler_kwargs: PocketTTSHandlerArguments, - kokoro_tts_handler_kwargs: KokoroTTSHandlerArguments, - qwen3_tts_handler_kwargs: Qwen3TTSHandlerArguments, -) -> "PipelineUnit": - """Build one isolated realtime pipeline (own queues, events, service, handlers). - - Returns a PipelineUnit that lives inside RealtimeServer's pool. Handler instances - are returned in `unit.handlers`; the caller threads them via ThreadManager. No uvicorn - is created here — the single server is owned by RealtimeServer. - """ - from speech_to_speech.api.openai_realtime.service import RealtimeService - - # Per-unit deep copies so mutations (cancel_scope, text_output_queue) don't bleed across pipelines. - vad_kw = deepcopy(vad_handler_kwargs) - whisper_kw = deepcopy(whisper_stt_handler_kwargs) - faster_whisper_kw = deepcopy(faster_whisper_stt_handler_kwargs) - paraformer_kw = deepcopy(paraformer_stt_handler_kwargs) - mlx_audio_whisper_kw = deepcopy(mlx_audio_whisper_stt_handler_kwargs) - parakeet_kw = deepcopy(parakeet_tdt_stt_handler_kwargs) - lm_kw = deepcopy(language_model_handler_kwargs) - responses_api_kw = deepcopy(responses_api_language_model_handler_kwargs) - chat_tts_kw = deepcopy(chat_tts_handler_kwargs) - facebook_mms_kw = deepcopy(facebook_mms_tts_handler_kwargs) - pocket_tts_kw = deepcopy(pocket_tts_handler_kwargs) - kokoro_tts_kw = deepcopy(kokoro_tts_handler_kwargs) - qwen3_tts_kw = deepcopy(qwen3_tts_handler_kwargs) - - should_listen = Event() - response_playing = Event() - cancel_scope = CancelScope() - speculative_turns = SpeculativeTurnTracker() - recv_audio_chunks_queue: Queue[AudioInItem] = Queue() - send_audio_chunks_queue: Queue[AudioOutItem] = Queue() - spoken_prompt_queue: Queue[VADOutItem] = Queue() - stt_output_queue: Queue[STTOutItem] = Queue() - text_prompt_queue: Queue[TextPromptItem] = Queue() - lm_response_queue: Queue[LMOutItem] = Queue() - lm_processed_queue: Queue[TTSInItem] = Queue() - text_output_queue: Queue[TextEventItem] = Queue() - - vars(vad_kw)["text_output_queue"] = text_output_queue - vars(vad_kw)["speculative_turns"] = speculative_turns - for kw in ( - lm_kw, - responses_api_kw, - kokoro_tts_kw, - qwen3_tts_kw, - pocket_tts_kw, - chat_tts_kw, - facebook_mms_kw, - ): - vars(kw)["cancel_scope"] = cancel_scope - vars(kw)["speculative_turns"] = speculative_turns - - if module_kwargs.llm_backend in ("responses-api", "chat-completions"): - chat_size = vars(responses_api_kw).get("chat_size", 10) - else: - chat_size = vars(lm_kw).get("chat_size", 10) - - service = RealtimeService( - text_prompt_queue=text_prompt_queue, - should_listen=should_listen, - chat_size=chat_size, - speculative_turns=speculative_turns, - ) - - if module_kwargs.enable_live_transcription: - vad_kw.enable_realtime_transcription = True - vad_kw.realtime_processing_pause = module_kwargs.live_transcription_update_interval - - handlers = _build_pipeline_handlers( - stop_event=stop_event, - should_listen=should_listen, - recv_audio_chunks_queue=recv_audio_chunks_queue, - spoken_prompt_queue=spoken_prompt_queue, - stt_output_queue=stt_output_queue, - text_prompt_queue=text_prompt_queue, - lm_response_queue=lm_response_queue, - lm_processed_queue=lm_processed_queue, - send_audio_chunks_queue=send_audio_chunks_queue, - text_output_queue=text_output_queue, - transcription_notifier_setup={ - "text_output_queue": text_output_queue, - "should_listen": should_listen, - }, - module_kwargs=module_kwargs, - vad_handler_kwargs=vad_kw, - whisper_stt_handler_kwargs=whisper_kw, - faster_whisper_stt_handler_kwargs=faster_whisper_kw, - paraformer_stt_handler_kwargs=paraformer_kw, - mlx_audio_whisper_stt_handler_kwargs=mlx_audio_whisper_kw, - parakeet_tdt_stt_handler_kwargs=parakeet_kw, - language_model_handler_kwargs=lm_kw, - responses_api_language_model_handler_kwargs=responses_api_kw, - chat_tts_handler_kwargs=chat_tts_kw, - facebook_mms_tts_handler_kwargs=facebook_mms_kw, - pocket_tts_handler_kwargs=pocket_tts_kw, - kokoro_tts_handler_kwargs=kokoro_tts_kw, - qwen3_tts_handler_kwargs=qwen3_tts_kw, - speculative_turns=speculative_turns, - ) - for h in handlers: - h.pipeline_index = index - - return PipelineUnit( - index=index, - service=service, - cancel_scope=cancel_scope, - should_listen=should_listen, - response_playing=response_playing, - input_queue=recv_audio_chunks_queue, - output_queue=send_audio_chunks_queue, - text_output_queue=text_output_queue, - text_prompt_queue=text_prompt_queue, - handlers=handlers, - ) - - -def build_pipeline( - module_kwargs: ModuleArguments, - socket_receiver_kwargs: SocketReceiverArguments, - socket_sender_kwargs: SocketSenderArguments, - websocket_streamer_kwargs: WebSocketStreamerArguments, - vad_handler_kwargs: VADHandlerArguments, - whisper_stt_handler_kwargs: WhisperSTTHandlerArguments, - faster_whisper_stt_handler_kwargs: FasterWhisperSTTHandlerArguments, - paraformer_stt_handler_kwargs: ParaformerSTTHandlerArguments, - mlx_audio_whisper_stt_handler_kwargs: MLXAudioWhisperSTTHandlerArguments, - parakeet_tdt_stt_handler_kwargs: ParakeetTDTSTTHandlerArguments, - language_model_handler_kwargs: LanguageModelHandlerArguments, - responses_api_language_model_handler_kwargs: ResponsesApiLanguageModelHandlerArguments, - chat_tts_handler_kwargs: ChatTTSHandlerArguments, - facebook_mms_tts_handler_kwargs: FacebookMMSTTSHandlerArguments, - pocket_tts_handler_kwargs: PocketTTSHandlerArguments, - kokoro_tts_handler_kwargs: KokoroTTSHandlerArguments, - qwen3_tts_handler_kwargs: Qwen3TTSHandlerArguments, - queues_and_events: dict[str, Any], -) -> ThreadManager: - stop_event = queues_and_events["stop_event"] - should_listen = queues_and_events["should_listen"] - recv_audio_chunks_queue = queues_and_events["recv_audio_chunks_queue"] - send_audio_chunks_queue = queues_and_events["send_audio_chunks_queue"] - spoken_prompt_queue = queues_and_events["spoken_prompt_queue"] - stt_output_queue = queues_and_events["stt_output_queue"] - text_prompt_queue = queues_and_events["text_prompt_queue"] - lm_response_queue = queues_and_events["lm_response_queue"] - lm_processed_queue = queues_and_events["lm_processed_queue"] - text_output_queue = ( - None # Only set for websocket/realtime modes; kept None otherwise to avoid unbounded queue growth - ) - - comms_handlers: list[Any] = [] - if module_kwargs.mode == "local": - from speech_to_speech.connections.local_audio_streamer import LocalAudioStreamer - - local_audio_streamer = LocalAudioStreamer( - input_queue=recv_audio_chunks_queue, - output_queue=send_audio_chunks_queue, - should_listen=should_listen, - ) - comms_handlers = [local_audio_streamer] - should_listen.set() - elif module_kwargs.mode == "websocket": - from speech_to_speech.connections.websocket_streamer import WebSocketStreamer - - text_output_queue = queues_and_events["text_output_queue"] - websocket_streamer = WebSocketStreamer( - stop_event, - input_queue=recv_audio_chunks_queue, - output_queue=send_audio_chunks_queue, - should_listen=should_listen, - text_output_queue=text_output_queue, - host=websocket_streamer_kwargs.ws_host, - port=websocket_streamer_kwargs.ws_port, - ) - comms_handlers = [websocket_streamer] - elif module_kwargs.mode == "realtime": - from speech_to_speech.api.openai_realtime.server import RealtimeServer - - pool_size = max(1, module_kwargs.num_pipelines) - pool = [ - _build_realtime_pipeline_unit( - index=i, - stop_event=stop_event, - module_kwargs=module_kwargs, - vad_handler_kwargs=vad_handler_kwargs, - whisper_stt_handler_kwargs=whisper_stt_handler_kwargs, - faster_whisper_stt_handler_kwargs=faster_whisper_stt_handler_kwargs, - paraformer_stt_handler_kwargs=paraformer_stt_handler_kwargs, - mlx_audio_whisper_stt_handler_kwargs=mlx_audio_whisper_stt_handler_kwargs, - parakeet_tdt_stt_handler_kwargs=parakeet_tdt_stt_handler_kwargs, - language_model_handler_kwargs=language_model_handler_kwargs, - responses_api_language_model_handler_kwargs=responses_api_language_model_handler_kwargs, - chat_tts_handler_kwargs=chat_tts_handler_kwargs, - facebook_mms_tts_handler_kwargs=facebook_mms_tts_handler_kwargs, - pocket_tts_handler_kwargs=pocket_tts_handler_kwargs, - kokoro_tts_handler_kwargs=kokoro_tts_handler_kwargs, - qwen3_tts_handler_kwargs=qwen3_tts_handler_kwargs, - ) - for i in range(pool_size) - ] - - realtime_server = RealtimeServer( - stop_event=stop_event, - pool=pool, - host=websocket_streamer_kwargs.ws_host, - port=websocket_streamer_kwargs.ws_port, - ) - - all_handlers: list[Any] = [realtime_server] - for unit in pool: - all_handlers.extend(unit.handlers) - return ThreadManager(all_handlers) - else: - from speech_to_speech.connections.socket_receiver import SocketReceiver - from speech_to_speech.connections.socket_sender import SocketSender - - comms_handlers = [ - SocketReceiver( - stop_event, - recv_audio_chunks_queue, - should_listen, - host=socket_receiver_kwargs.recv_host, - port=socket_receiver_kwargs.recv_port, - chunk_size=socket_receiver_kwargs.chunk_size, - ), - SocketSender( - stop_event, - send_audio_chunks_queue, - should_listen, - host=socket_sender_kwargs.send_host, - port=socket_sender_kwargs.send_port, - ), - ] - - # Set VAD realtime transcription parameters from module_kwargs - if module_kwargs.enable_live_transcription: - vad_handler_kwargs.enable_realtime_transcription = True - vad_handler_kwargs.realtime_processing_pause = module_kwargs.live_transcription_update_interval - - if module_kwargs.llm_backend in ("responses-api", "chat-completions"): - _lm_vars = vars(responses_api_language_model_handler_kwargs) - else: - _lm_vars = vars(language_model_handler_kwargs) - transcription_notifier_setup: dict[str, Any] = { - "text_output_queue": text_output_queue, - "should_listen": should_listen, - "runtime_config": RuntimeConfig( - chat=Chat(_lm_vars.get("chat_size", 30)), - session=RealtimeSessionCreateRequest( - type="realtime", - instructions=_lm_vars.get("init_chat_prompt"), - ), - ), - } - - pipeline_handlers = _build_pipeline_handlers( - stop_event=stop_event, - should_listen=should_listen, - recv_audio_chunks_queue=recv_audio_chunks_queue, - spoken_prompt_queue=spoken_prompt_queue, - stt_output_queue=stt_output_queue, - text_prompt_queue=text_prompt_queue, - lm_response_queue=lm_response_queue, - lm_processed_queue=lm_processed_queue, - send_audio_chunks_queue=send_audio_chunks_queue, - text_output_queue=text_output_queue, - transcription_notifier_setup=transcription_notifier_setup, - module_kwargs=module_kwargs, - vad_handler_kwargs=vad_handler_kwargs, - whisper_stt_handler_kwargs=whisper_stt_handler_kwargs, - faster_whisper_stt_handler_kwargs=faster_whisper_stt_handler_kwargs, - paraformer_stt_handler_kwargs=paraformer_stt_handler_kwargs, - mlx_audio_whisper_stt_handler_kwargs=mlx_audio_whisper_stt_handler_kwargs, - parakeet_tdt_stt_handler_kwargs=parakeet_tdt_stt_handler_kwargs, - language_model_handler_kwargs=language_model_handler_kwargs, - responses_api_language_model_handler_kwargs=responses_api_language_model_handler_kwargs, - chat_tts_handler_kwargs=chat_tts_handler_kwargs, - facebook_mms_tts_handler_kwargs=facebook_mms_tts_handler_kwargs, - pocket_tts_handler_kwargs=pocket_tts_handler_kwargs, - kokoro_tts_handler_kwargs=kokoro_tts_handler_kwargs, - qwen3_tts_handler_kwargs=qwen3_tts_handler_kwargs, - ) - - return ThreadManager([*comms_handlers, *pipeline_handlers]) - - -def get_stt_handler( - module_kwargs: ModuleArguments, - stop_event: Event, - spoken_prompt_queue: Queue[VADOutItem], - text_prompt_queue: Queue[STTOutItem], - speculative_turns: SpeculativeTurnTracker | None, - whisper_stt_handler_kwargs: WhisperSTTHandlerArguments, - faster_whisper_stt_handler_kwargs: FasterWhisperSTTHandlerArguments, - paraformer_stt_handler_kwargs: ParaformerSTTHandlerArguments, - mlx_audio_whisper_stt_handler_kwargs: MLXAudioWhisperSTTHandlerArguments, - parakeet_tdt_stt_handler_kwargs: ParakeetTDTSTTHandlerArguments, -) -> BaseHandler[STTIn, STTOut]: - from speech_to_speech.STT.base_stt_handler import BaseSTTHandler - - def with_speculative_turns(handler: BaseSTTHandler) -> BaseSTTHandler: - if speculative_turns is not None: - handler.speculative_turns = speculative_turns - return handler - - if module_kwargs.stt == "whisper": - from speech_to_speech.STT.whisper_stt_handler import WhisperSTTHandler - - return with_speculative_turns( - WhisperSTTHandler( - stop_event, - queue_in=spoken_prompt_queue, - queue_out=text_prompt_queue, - setup_kwargs=vars(whisper_stt_handler_kwargs), - ) - ) - elif module_kwargs.stt == "whisper-mlx": - from speech_to_speech.STT.lightning_whisper_mlx_handler import LightningWhisperSTTHandler - - return with_speculative_turns( - LightningWhisperSTTHandler( - stop_event, - queue_in=spoken_prompt_queue, - queue_out=text_prompt_queue, - setup_kwargs=vars(whisper_stt_handler_kwargs), - ) - ) - elif module_kwargs.stt == "mlx-audio-whisper": - from speech_to_speech.STT.mlx_audio_whisper_handler import MLXAudioWhisperSTTHandler - - # Merge MLX Audio Whisper kwargs with shared language parameter from Whisper kwargs - setup_kwargs = {**vars(mlx_audio_whisper_stt_handler_kwargs), "language": whisper_stt_handler_kwargs.language} - return with_speculative_turns( - MLXAudioWhisperSTTHandler( - stop_event, - queue_in=spoken_prompt_queue, - queue_out=text_prompt_queue, - setup_kwargs=setup_kwargs, - ) - ) - elif module_kwargs.stt == "paraformer": - from speech_to_speech.STT.paraformer_handler import ParaformerSTTHandler - - return with_speculative_turns( - ParaformerSTTHandler( - stop_event, - queue_in=spoken_prompt_queue, - queue_out=text_prompt_queue, - setup_kwargs=vars(paraformer_stt_handler_kwargs), - ) - ) - elif module_kwargs.stt == "faster-whisper": - from speech_to_speech.STT.faster_whisper_handler import FasterWhisperSTTHandler - - return with_speculative_turns( - FasterWhisperSTTHandler( - stop_event, - queue_in=spoken_prompt_queue, - queue_out=text_prompt_queue, - setup_kwargs=vars(faster_whisper_stt_handler_kwargs), - ) - ) - elif module_kwargs.stt == "parakeet-tdt": - from speech_to_speech.STT.parakeet_tdt_handler import ParakeetTDTSTTHandler - - # Add live transcription parameters to setup_kwargs - setup_kwargs = { - **vars(parakeet_tdt_stt_handler_kwargs), - "enable_live_transcription": module_kwargs.enable_live_transcription, - "live_transcription_update_interval": module_kwargs.live_transcription_update_interval, - } - - return with_speculative_turns( - ParakeetTDTSTTHandler( - stop_event, - queue_in=spoken_prompt_queue, - queue_out=text_prompt_queue, - setup_kwargs=setup_kwargs, - ) - ) - else: - raise ValueError( - "The STT should be either none, whisper, whisper-mlx, mlx-audio-whisper, faster-whisper, parakeet-tdt, or paraformer." - ) - - -def get_llm_handler( - module_kwargs: ModuleArguments, - stop_event: Event, - text_prompt_queue: Queue[TextPromptItem], - lm_response_queue: Queue[LMOutItem], - language_model_handler_kwargs: LanguageModelHandlerArguments, - responses_api_language_model_handler_kwargs: ResponsesApiLanguageModelHandlerArguments, -) -> BaseHandler[LLMIn, LLMOut]: - if module_kwargs.llm_backend == "responses-api": - from speech_to_speech.LLM.responses_api_language_model import ResponsesApiModelHandler - - return ResponsesApiModelHandler( - stop_event, - queue_in=text_prompt_queue, - queue_out=lm_response_queue, - setup_kwargs=vars(responses_api_language_model_handler_kwargs), - ) - - if module_kwargs.llm_backend == "chat-completions": - from speech_to_speech.LLM.chat_completions_language_model import ChatCompletionsApiModelHandler - - # Reuses the responses-api argument class (identical fields: model_name, - # base_url, api_key, stream, disable_thinking, ...). - return ChatCompletionsApiModelHandler( - stop_event, - queue_in=text_prompt_queue, - queue_out=lm_response_queue, - setup_kwargs=vars(responses_api_language_model_handler_kwargs), - ) - - if module_kwargs.llm_backend in ("transformers", "mlx-lm"): - lm_kwargs = vars(language_model_handler_kwargs) - is_vlm = lm_kwargs.pop("is_vlm", False) - if module_kwargs.llm_backend == "mlx-lm": - lm_kwargs["backend"] = "mlx" - - if is_vlm: - from speech_to_speech.LLM.language_model import VisionLanguageModelHandler - - return VisionLanguageModelHandler( - stop_event, - queue_in=text_prompt_queue, - queue_out=lm_response_queue, - setup_kwargs=lm_kwargs, - ) - from speech_to_speech.LLM.language_model import LanguageModelHandler - - return LanguageModelHandler( - stop_event, - queue_in=text_prompt_queue, - queue_out=lm_response_queue, - setup_kwargs=lm_kwargs, - ) - - raise ValueError("The LLM should be either transformers, mlx-lm, responses-api or chat-completions") - - -def get_tts_handler( - module_kwargs: ModuleArguments, - stop_event: Event, - lm_response_queue: Queue[TTSInItem], - send_audio_chunks_queue: Queue[AudioOutItem], - should_listen: Event, - chat_tts_handler_kwargs: ChatTTSHandlerArguments, - facebook_mms_tts_handler_kwargs: FacebookMMSTTSHandlerArguments, - pocket_tts_handler_kwargs: PocketTTSHandlerArguments, - kokoro_tts_handler_kwargs: KokoroTTSHandlerArguments, - qwen3_tts_handler_kwargs: Qwen3TTSHandlerArguments, -) -> BaseHandler[TTSIn, TTSOut]: - if module_kwargs.tts == "chatTTS": - try: - from speech_to_speech.TTS.chatTTS_handler import ChatTTSHandler - except (ImportError, RuntimeError) as e: - logger.error('Error importing ChatTTSHandler. Install it with `pip install "speech-to-speech[chattts]"`.') - raise e - return ChatTTSHandler( - stop_event, - queue_in=lm_response_queue, - queue_out=send_audio_chunks_queue, - setup_args=(should_listen,), - setup_kwargs=vars(chat_tts_handler_kwargs), - ) - elif module_kwargs.tts == "facebookMMS": - from speech_to_speech.TTS.facebookmms_handler import FacebookMMSTTSHandler - - return FacebookMMSTTSHandler( - stop_event, - queue_in=lm_response_queue, - queue_out=send_audio_chunks_queue, - setup_args=(should_listen,), - setup_kwargs=vars(facebook_mms_tts_handler_kwargs), - ) - elif module_kwargs.tts == "pocket": - try: - from speech_to_speech.TTS.pocket_tts_handler import PocketTTSHandler - except ImportError as e: - raise ImportError( - 'Pocket TTS is optional. Install it with `pip install "speech-to-speech[pocket]"`.' - ) from e - - return PocketTTSHandler( - stop_event, - queue_in=lm_response_queue, - queue_out=send_audio_chunks_queue, - setup_args=(should_listen,), - setup_kwargs=vars(pocket_tts_handler_kwargs), - ) - elif module_kwargs.tts == "kokoro": - try: - from speech_to_speech.TTS.kokoro_handler import KokoroTTSHandler - except ImportError as e: - raise ImportError('Kokoro is optional. Install it with `pip install "speech-to-speech[kokoro]"`.') from e - - return KokoroTTSHandler( - stop_event, - queue_in=lm_response_queue, - queue_out=send_audio_chunks_queue, - setup_args=(should_listen,), - setup_kwargs=vars(kokoro_tts_handler_kwargs), - ) - elif module_kwargs.tts == "qwen3": - from speech_to_speech.TTS.qwen3_tts_handler import Qwen3TTSHandler - - return Qwen3TTSHandler( - stop_event, - queue_in=lm_response_queue, - queue_out=send_audio_chunks_queue, - setup_args=(should_listen,), - setup_kwargs=vars(qwen3_tts_handler_kwargs), - ) - else: - raise ValueError("The TTS should be either chatTTS, facebookMMS, pocket, kokoro, or qwen3") - - -def main() -> None: - args = parse_arguments() - - setup_logger(args.module_kwargs.log_level) - - if args.module_kwargs.num_pipelines < 1: - raise ValueError(f"--num_pipelines must be >= 1, got {args.module_kwargs.num_pipelines}") - - prepare_all_args( - args.module_kwargs, - args.whisper_stt_handler_kwargs, - args.paraformer_stt_handler_kwargs, - args.faster_whisper_stt_handler_kwargs, - args.mlx_audio_whisper_stt_handler_kwargs, - args.parakeet_tdt_stt_handler_kwargs, - args.language_model_handler_kwargs, - args.responses_api_language_model_handler_kwargs, - args.chat_tts_handler_kwargs, - args.facebook_mms_tts_handler_kwargs, - args.pocket_tts_handler_kwargs, - args.kokoro_tts_handler_kwargs, - args.qwen3_tts_handler_kwargs, - ) - - # Validate after prepare_all_args(): --local_mac_optimal_settings mutates - # module_kwargs.mode to "local", so checking before would let - # --local_mac_optimal_settings --num_pipelines 2 sneak past this guard. - if args.module_kwargs.num_pipelines > 1 and args.module_kwargs.mode != "realtime": - raise ValueError( - f"--num_pipelines > 1 is only supported with --mode realtime " - f"(got mode={args.module_kwargs.mode!r}, num_pipelines={args.module_kwargs.num_pipelines})" - ) - - # On Apple Silicon, all MLX inference serializes through a global lock (utils/mlx_lock.py). - # The progressive STT path uses a short timeout and drops work under contention, producing - # a flood of warnings without affecting final transcripts. With a pool, pre-emptively turn - # it off so logs stay readable; the final STT path is unaffected. Non-darwin platforms - # don't share this lock, so leave their live transcription alone. - if args.module_kwargs.num_pipelines > 1 and platform == "darwin" and args.module_kwargs.enable_live_transcription: - logger.info( - "MLX contention: --num_pipelines=%d > 1 on Apple Silicon → disabling live transcription " - "(progressive STT contends on the global MLX lock)", - args.module_kwargs.num_pipelines, - ) - args.module_kwargs.enable_live_transcription = False - - queues_and_events = initialize_queues_and_events() - - pipeline_manager = build_pipeline( - args.module_kwargs, - args.socket_receiver_kwargs, - args.socket_sender_kwargs, - args.websocket_streamer_kwargs, - args.vad_handler_kwargs, - args.whisper_stt_handler_kwargs, - args.faster_whisper_stt_handler_kwargs, - args.paraformer_stt_handler_kwargs, - args.mlx_audio_whisper_stt_handler_kwargs, - args.parakeet_tdt_stt_handler_kwargs, - args.language_model_handler_kwargs, - args.responses_api_language_model_handler_kwargs, - args.chat_tts_handler_kwargs, - args.facebook_mms_tts_handler_kwargs, - args.pocket_tts_handler_kwargs, - args.kokoro_tts_handler_kwargs, - args.qwen3_tts_handler_kwargs, - queues_and_events, - ) - - # Set up graceful shutdown handler - shutdown_requested = [False] # Use list for nonlocal mutation - - def signal_handler(_sig: int, _frame: Optional[FrameType]) -> None: - if not shutdown_requested[0]: - shutdown_requested[0] = True - console.print("\n[yellow]Shutting down gracefully...[/yellow]") - pipeline_manager.stop() - console.print("[green]✓ Pipeline stopped successfully[/green]") - - signal.signal(signal.SIGINT, signal_handler) - signal.signal(signal.SIGTERM, signal_handler) - - try: - pipeline_manager.start() - pipeline_manager.wait() - except KeyboardInterrupt: - if not shutdown_requested[0]: - console.print("\n[yellow]Shutting down gracefully...[/yellow]") - pipeline_manager.stop() - console.print("[green]✓ Pipeline stopped successfully[/green]") - - -if __name__ == "__main__": - main() diff --git a/runtime/src/speech_to_speech/utils/__init__.py b/runtime/src/speech_to_speech/utils/__init__.py deleted file mode 100644 index e69de29bb2d1d6434b8b29ae775ad8c2e48c5391..0000000000000000000000000000000000000000 diff --git a/runtime/src/speech_to_speech/utils/mlx_lock.py b/runtime/src/speech_to_speech/utils/mlx_lock.py deleted file mode 100644 index 37cb666fd99a7ecb7551e9e594f44cbe5bcd2874..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/utils/mlx_lock.py +++ /dev/null @@ -1,171 +0,0 @@ -""" -Global MLX lock to prevent concurrent Metal/MLX model access. - -MLX models (STT, LLM, TTS) cannot be used concurrently from multiple threads -on Apple Silicon due to Metal command buffer limitations. This module provides -a global lock that all MLX handlers should acquire before using their models. -""" - -import logging -import types -from threading import Lock, RLock, current_thread, get_ident -from time import perf_counter -from typing import Literal - -logger = logging.getLogger(__name__) - -# Global reentrant lock for MLX model access -# Using RLock (reentrant lock) so the same thread can acquire it multiple times -_mlx_lock = RLock() -_mlx_lock_state = Lock() -_lock_owner_ident: int | None = None -_lock_owner_thread: str | None = None -_lock_owner_handler: str | None = None -_lock_acquired_at: float | None = None -_lock_depth = 0 - - -def _owner_snapshot(now: float | None = None) -> str: - if _lock_owner_ident is None: - return "none" - elapsed = "" - if _lock_acquired_at is not None: - elapsed = f" for {(now or perf_counter()) - _lock_acquired_at:.3f}s" - return f"{_lock_owner_handler} on {_lock_owner_thread}{elapsed}" - - -def _record_lock_acquired(handler_name: str) -> int: - global _lock_acquired_at, _lock_depth, _lock_owner_handler, _lock_owner_ident, _lock_owner_thread - ident = get_ident() - thread_name = current_thread().name - now = perf_counter() - with _mlx_lock_state: - if _lock_owner_ident == ident: - _lock_depth += 1 - else: - _lock_owner_ident = ident - _lock_owner_thread = thread_name - _lock_owner_handler = handler_name - _lock_acquired_at = now - _lock_depth = 1 - return _lock_depth - - -def _record_lock_released(handler_name: str) -> tuple[int, float | None]: - global _lock_acquired_at, _lock_depth, _lock_owner_handler, _lock_owner_ident, _lock_owner_thread - ident = get_ident() - now = perf_counter() - with _mlx_lock_state: - if _lock_owner_ident != ident: - logger.warning( - "%s: MLX lock release requested by non-owner thread; owner=%s", - handler_name, - _owner_snapshot(now), - ) - return _lock_depth, None - - hold_s = None if _lock_acquired_at is None else now - _lock_acquired_at - _lock_depth -= 1 - depth = _lock_depth - if _lock_depth <= 0: - _lock_owner_ident = None - _lock_owner_thread = None - _lock_owner_handler = None - _lock_acquired_at = None - _lock_depth = 0 - return depth, hold_s - - -def acquire_mlx_lock(timeout: float | None = None, handler_name: str = "Unknown") -> bool: - """ - Acquire the global MLX lock. - - Args: - timeout: Optional timeout in seconds - handler_name: Name of the handler acquiring the lock (for logging) - - Returns: - True if lock was acquired, False if timeout occurred - """ - with _mlx_lock_state: - owner_before = _owner_snapshot() - logger.debug("%s: Attempting to acquire MLX lock (owner=%s)", handler_name, owner_before) - start = perf_counter() - acquired = _mlx_lock.acquire(timeout=timeout) if timeout else _mlx_lock.acquire(blocking=True) - wait_s = perf_counter() - start - - if acquired: - depth = _record_lock_acquired(handler_name) - if wait_s >= 0.25: - logger.info( - "%s: MLX lock acquired after %.2fs (previous_owner=%s, depth=%d)", - handler_name, - wait_s, - owner_before, - depth, - ) - else: - logger.debug( - "%s: MLX lock acquired after %.3fs (previous_owner=%s, depth=%d)", - handler_name, - wait_s, - owner_before, - depth, - ) - else: - with _mlx_lock_state: - owner_after = _owner_snapshot() - logger.warning( - "%s: Failed to acquire MLX lock after %.3fs (timeout=%s, owner=%s)", - handler_name, - wait_s, - timeout, - owner_after, - ) - - return acquired - - -def release_mlx_lock(handler_name: str = "Unknown") -> None: - """ - Release the global MLX lock. - - Args: - handler_name: Name of the handler releasing the lock (for logging) - """ - try: - depth, hold_s = _record_lock_released(handler_name) - _mlx_lock.release() - if hold_s is not None and depth == 0 and hold_s >= 0.25: - logger.info("%s: MLX lock released after holding %.2fs", handler_name, hold_s) - else: - logger.debug( - "%s: MLX lock released%s", - handler_name, - "" if hold_s is None else f" after holding {hold_s:.3f}s (depth={depth})", - ) - except RuntimeError as e: - logger.error(f"{handler_name}: Failed to release MLX lock: {e}") - - -class MLXLockContext: - """Context manager for MLX lock.""" - - def __init__(self, handler_name: str = "Unknown", timeout: float | None = None) -> None: - self.handler_name = handler_name - self.timeout = timeout - self.acquired = False - - def __enter__(self) -> bool: - self.acquired = acquire_mlx_lock(timeout=self.timeout, handler_name=self.handler_name) - return self.acquired - - def __exit__( - self, - exc_type: type[BaseException] | None, - exc_val: BaseException | None, - exc_tb: types.TracebackType | None, - ) -> Literal[False]: - if self.acquired: - release_mlx_lock(handler_name=self.handler_name) - return False # Do not suppress exceptions diff --git a/runtime/src/speech_to_speech/utils/thread_manager.py b/runtime/src/speech_to_speech/utils/thread_manager.py deleted file mode 100644 index e12a1b710dba6a6a6454fcc54c36670bdea69486..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/utils/thread_manager.py +++ /dev/null @@ -1,39 +0,0 @@ -import logging -import threading -from collections.abc import Sequence -from typing import Any - -logger = logging.getLogger(__name__) - - -class ThreadManager: - """ - Manages multiple threads used to execute given handler tasks. - """ - - def __init__(self, handlers: Sequence[Any]) -> None: - self.handlers = handlers - self.threads: list[threading.Thread] = [] - - def start(self) -> None: - for handler in self.handlers: - thread = threading.Thread(target=handler.run) - thread.daemon = False # Ensure threads are waited for on shutdown - self.threads.append(thread) - thread.start() - - def wait(self) -> None: - for thread in self.threads: - thread.join() - - def stop(self) -> None: - # Signal all handlers to stop - for handler in self.handlers: - handler.stop_event.set() - - # Wait for all threads to finish with timeout - for i, thread in enumerate(self.threads): - if thread.is_alive(): - thread.join(timeout=5.0) - if thread.is_alive(): - logger.warning(f"Thread {i} ({thread.name}) did not terminate within timeout") diff --git a/runtime/src/speech_to_speech/utils/utils.py b/runtime/src/speech_to_speech/utils/utils.py deleted file mode 100644 index 29e112aaef5e471d6718b6a672ad971b6e5f171b..0000000000000000000000000000000000000000 --- a/runtime/src/speech_to_speech/utils/utils.py +++ /dev/null @@ -1,56 +0,0 @@ -from __future__ import annotations - -import uuid -from typing import TYPE_CHECKING - -import numpy as np - -if TYPE_CHECKING: - from openai.types.realtime.realtime_response_create_params import RealtimeResponseCreateParams - - -def response_wants_audio(response: RealtimeResponseCreateParams | None) -> bool: - """Whether a response should produce audio (and audio events) vs. text only. - - Mirrors the OpenAI realtime semantics for ``output_modalities``: an absent - value (``None``) or an empty list, or an explicit ``"audio"`` entry means - audio; a non-empty list without ``"audio"`` (e.g. ``["text"]``) means text - only. - """ - if response is None: - return True - mods = response.output_modalities - return not mods or "audio" in mods - - -def is_out_of_band(response: RealtimeResponseCreateParams | None) -> bool: - """Whether a response is *out-of-band* (``conversation="none"``). - - Out-of-band responses are generated against a temporary context and are - never threaded into the default conversation: their ``input`` (if any) - seeds a throwaway chat, their assistant output is not committed back, and - their ``conversation_id`` is reported as ``null``. Any other ``conversation`` - value (``"auto"``, ``None``, or an arbitrary id) is treated as in-band. - """ - return response is not None and response.conversation == "none" - - -def next_power_of_2(x: int) -> int: - return 1 if x == 0 else 2 ** (x - 1).bit_length() - - -def _generate_id(prefix: str) -> str: - return f"{prefix}_{uuid.uuid4().hex}" - - -def int2float(sound: np.ndarray) -> np.ndarray: - """ - Taken from https://github.com/snakers4/silero-vad - """ - - abs_max = np.abs(sound).max() - sound = sound.astype("float32") - if abs_max > 0: - sound *= 1 / 32768 - sound = sound.squeeze() # depends on the use case - return sound