This commit is contained in:
2026-08-26 17:42:00 -04:00
commit ad95d530ad
134 changed files with 31750 additions and 0 deletions
+139
View File
@@ -0,0 +1,139 @@
# 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 <language>.`
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
```
View File
@@ -0,0 +1,588 @@
from __future__ import annotations
import logging
from abc import ABC, abstractmethod
from collections.abc import Iterator
from typing import Any, Optional
import httpx
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_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
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)
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,
**_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.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
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"
@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]) -> 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."""
...
# ── speculative-turn / cancellation gating ─────────────────────────────────
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 _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:
chat.add_item(pending_item)
state.pending.clear()
chat.add_item(fc_item)
yield self._chunk(turn, tools=[item])
# ── consumption ─────────────────────────────────────────────────────────--
def _consume_streaming(self, events: Iterator[ProviderEvent], state: _GenState, turn: _Turn) -> Iterator[LLMOut]:
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}")
def _consume_nonstreaming(self, events: Iterator[ProviderEvent], state: _GenState, turn: _Turn) -> Iterator[LLMOut]:
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
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}")
# ── orchestration ─────────────────────────────────────────────────────────
def _generate(
self,
active_chat: Chat,
original_chat: Chat,
turn: _Turn,
optional_kwargs: dict[str, Any],
) -> Iterator[LLMOut]:
api_response: Any = None
state = _GenState()
error_message: str | None = None
api_input = 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."
try:
if error_message is None:
api_response = self._request(api_input, optional_kwargs)
if api_response is not None:
events = self._iter_events(api_response)
if self.stream:
yield from self._consume_streaming(events, state, turn)
else:
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}"
finally:
if api_response is not None and hasattr(api_response, "close"):
try:
api_response.close()
except Exception:
pass
if (
error_message is None
and not self._generation_is_stale(turn.gen)
and self._turn_output_allowed(turn.turn_id, turn.turn_revision)
):
# 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:
original_chat.add_item(item)
original_chat.strip_images(consumed_image_ids)
original_chat.trim_if_needed(self.compactor)
if 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
)
def process(self, request: LLMIn) -> Iterator[LLMOut]:
"""Process a language model request and yield LLMResponseChunks."""
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
+760
View File
@@ -0,0 +1,760 @@
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__)
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=<fn>``: 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)
]
if not item.content:
raise ChatItemError(
"Message has no supported content. Supported modalities: input_text, input_image."
)
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=<fn>``: 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 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 = []
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)
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_images = any(p.type == "input_image" for p in item.content)
if has_images:
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 image 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") == "input_image")]
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_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()
@@ -0,0 +1,311 @@
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)
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]) -> ChatCompletionContentPartParam:
"""Convert one transformers content part to Chat-Completions shape.
``to_transformers_chat`` keeps Realtime-style parts (``input_text`` /
``input_image`` with a bare-string ``image_url``). The Chat Completions
HTTP API instead wants ``{type:"text", text}`` and
``{type:"image_url", image_url:{url, detail}}``. Unknown parts pass through.
"""
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)
return cast("ChatCompletionContentPartParam", part)
@classmethod
def _chat_messages(cls, chat: Chat) -> list[dict[str, Any]]:
"""Serialise the chat for the Chat Completions API.
``Chat.to_transformers_chat`` targets HuggingFace ``apply_chat_template``,
so two shapes need fixing up for the OpenAI Chat Completions HTTP API:
tool-call ``arguments`` must be a JSON *string* (not a parsed object), and
multimodal ``content`` parts must use the Chat Completions ``text`` /
``image_url`` shape rather than the Realtime ``input_text`` /
``input_image`` shape.
"""
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"] = [cls._to_chat_content_part(p) for p in content]
if message.get("role") == "tool":
message.pop("name", None)
return messages
# ── base hooks ──────────────────────────────────────────────────────────--
def _serialize(self, active_chat: Chat) -> list[dict[str, Any]]:
return self._chat_messages(active_chat)
def _build_optional_kwargs(self, req_tools: Any, req_tool_choice: Any) -> dict[str, Any]:
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 _request(self, api_input: list[dict[str, Any]], optional_kwargs: dict[str, Any]) -> Any:
create_kwargs: dict[str, Any] = dict(optional_kwargs)
if self.stream:
create_kwargs["stream_options"] = {"include_usage": True}
# Svuota i parametri dei tool per evitare il blocco 400 di NVIDIA NIM
create_kwargs.pop("tools", None)
create_kwargs.pop("tool_choice", None)
return self.client.chat.completions.create(
model=self.model_name,
messages=api_input, # type: ignore[arg-type] # runtime dicts match the Chat Completions message shape
stream=self.stream,
extra_body=self._extra_body,
timeout=self.request_timeout,
**create_kwargs,
)
# def _request(self, api_input: list[dict[str, Any]], optional_kwargs: dict[str, Any]) -> Any:
# create_kwargs: dict[str, Any] = dict(optional_kwargs)
# if self.stream:
# create_kwargs["stream_options"] = {"include_usage": True}
# return self.client.chat.completions.create(
# model=self.model_name,
# messages=api_input, # type: ignore[arg-type] # runtime dicts match the Chat Completions message shape
# stream=self.stream,
# extra_body=self._extra_body,
# timeout=self.request_timeout,
# **create_kwargs,
# )
def _iter_stream_events(self, api_response: Stream[ChatCompletionChunk]) -> Iterator[ProviderEvent]:
# Accumulate streamed tool-call deltas, keyed by their stream index, and the
# raw assistant text, then emit assistant message + tool calls + usage once
# the stream is exhausted.
tool_accum: dict[int, dict[str, str]] = {}
usage: Usage | None = None
raw_text = ""
for chunk in api_response:
# Usage-only trailing chunk (choices == []) when include_usage is set.
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 tc in delta.tool_calls:
entry = tool_accum.setdefault(tc.index, {"name": "", "args": "", "id": ""})
if tc.id:
entry["id"] = tc.id
if tc.function is not None:
if tc.function.name:
entry["name"] = tc.function.name
if tc.function.arguments:
entry["args"] += tc.function.arguments
# A refusal streams as `delta.refusal` with `delta.content` None;
# surface it as assistant text so it is spoken and stored.
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 self._tool_calls_from_accum(tool_accum)
if usage is not None:
yield usage
def _iter_response_events(self, api_response: Any) -> Iterator[ProviderEvent]:
usage = api_response.usage
if usage:
yield Usage(input_tokens=usage.prompt_tokens or 0, output_tokens=usage.completion_tokens or 0)
# A valid-but-empty response (e.g. content filter) returns no choices;
# complete cleanly with no assistant text rather than raising IndexError.
message = api_response.choices[0].message if api_response.choices else None
if message is None:
return
# A refusal arrives as `message.refusal` with `message.content` None; treat
# it as assistant text so it is spoken and stored.
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 tc in message.tool_calls or []:
tool_accum[len(tool_accum)] = {
"name": tc.function.name or "",
"args": tc.function.arguments or "",
"id": tc.id or "",
}
yield from self._tool_calls_from_accum(tool_accum)
@staticmethod
def _tool_calls_from_accum(tool_accum: dict[int, dict[str, str]]) -> Iterator[ToolCall]:
"""Turn accumulated tool-call deltas into ToolCall events.
IDs are regenerated (mirroring the Responses handler) so the rest of the
pipeline pairs each call_id with its function_call_output consistently.
"""
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 on_session_end(self) -> None:
logger.debug("Chat Completions API language model session state reset")
+181
View File
@@ -0,0 +1,181 @@
"""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
File diff suppressed because it is too large Load Diff
+148
View File
@@ -0,0 +1,148 @@
"""
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,
)
@@ -0,0 +1,149 @@
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.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
# ── 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]) -> Any:
return self.client.responses.create(
model=self.model_name,
input=api_input,
stream=self.stream,
extra_body=self._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")
+44
View File
@@ -0,0 +1,44 @@
"""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(),
)
@@ -0,0 +1,321 @@
#!/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. ``<code>func(x=1)</code>``). 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"<code>.*?</code>"``). 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, []
@@ -0,0 +1,34 @@
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, ' ')}"
@@ -0,0 +1,108 @@
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)
@@ -0,0 +1,111 @@
"""
Optional system-prompt builder that instructs a local LLM to output tool calls
inside delimited code blocks (e.g. ``<code>func(arg='val')</code>``).
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 = "<code>"
END_CODE = "</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 <emotion>."; 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("<code>", "</code>")
'<code>.*?</code>'
"""
return f"{re.escape(enter_code)}.*?{re.escape(end_code)}"
+75
View File
@@ -0,0 +1,75 @@
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,<b64>' 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))
+48
View File
@@ -0,0 +1,48 @@
"""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 <emotion>." 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(),
)
+163
View File
@@ -0,0 +1,163 @@
# 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
```
+1
View File
@@ -0,0 +1 @@
+212
View File
@@ -0,0 +1,212 @@
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()
@@ -0,0 +1,68 @@
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
@@ -0,0 +1,103 @@
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,
)
@@ -0,0 +1,155 @@
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
logger = logging.getLogger(__name__)
console = Console()
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
self.last_language = language
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
_ = 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 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
if self.start_language and self.start_language != "auto":
gen_kwargs["language"] = self.start_language
try:
# Generate transcription directly using model.generate
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()
# Try to detect language from result if available
if hasattr(result, "language"):
language_code = result.language
elif self.start_language and self.start_language != "auto":
language_code = self.start_language
else:
# Default to last known language or English
language_code = self.last_language if self.last_language else "en"
# Validate language code
if language_code not in SUPPORTED_LANGUAGES:
logger.warning(f"Detected unsupported language: {language_code}")
if self.last_language in SUPPORTED_LANGUAGES:
language_code = self.last_language
else:
language_code = "en"
else:
self.last_language = language_code
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 "en"
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,
)
@@ -0,0 +1,79 @@
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,
)
@@ -0,0 +1,648 @@
"""
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")
@@ -0,0 +1,343 @@
#!/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()
@@ -0,0 +1,103 @@
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,
)
+147
View File
@@ -0,0 +1,147 @@
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,
)
+155
View File
@@ -0,0 +1,155 @@
# 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.
+1
View File
@@ -0,0 +1 @@
# TTS module
+119
View File
@@ -0,0 +1,119 @@
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])),
)
+212
View File
@@ -0,0 +1,212 @@
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")
+418
View File
@@ -0,0 +1,418 @@
"""
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")
+227
View File
@@ -0,0 +1,227 @@
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
+903
View File
@@ -0,0 +1,903 @@
"""
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}")
Binary file not shown.
+1
View File
@@ -0,0 +1 @@
+855
View File
@@ -0,0 +1,855 @@
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)
else:
# Original mode: yield only when speech ends
yield from self._process_normal(vad_output)
def _process_realtime(self, vad_output: list[torch.Tensor] | 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),
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, 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) -> 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)
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
+170
View File
@@ -0,0 +1,170 @@
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
+3
View File
@@ -0,0 +1,3 @@
"""speech-to-speech: low-latency speech-to-speech pipeline."""
__version__ = "0.2.11"
View File
@@ -0,0 +1,216 @@
## 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`). |
---
## 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 `<code>...</code>` delimiters:
```
<code>
function_name(arg_name_1=value1, arg_name_2='string_value')
</code>
```
After generation, `_extract_tools` uses a regex to find `<code>` 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: "<result>"`
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` -- accepted `response.create`, no audio yet), 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`, 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
```
@@ -0,0 +1,84 @@
from __future__ import annotations
import logging
import os
from dataclasses import dataclass
from typing import Any
import numpy as np
from speech_to_speech.pipeline.messages import TTSInput
logger = logging.getLogger(__name__)
_TRUE_VALUES = {"1", "true", "yes", "on"}
@dataclass(slots=True)
class DelayFillerConfig:
enabled: bool
threshold_s: float
short_phrases: list[str]
full_phrases: list[str]
def _split_env_list(name: str) -> list[str]:
raw = os.environ.get(name, "").strip()
return [part.strip() for part in raw.split("|") if part.strip()]
def _env_flag(name: str, *, default: bool) -> bool:
raw = os.environ.get(name)
if raw is None:
return default
return raw.strip().lower() in _TRUE_VALUES
def load_delay_filler_config() -> DelayFillerConfig:
short_phrases = _split_env_list("TTS_DELAY_FILLERS_SHORT_ES")
full_phrases = _split_env_list("TTS_DELAY_FILLERS_FULL_ES")
# Backward compatibility with the old single-list variable.
if not short_phrases and not full_phrases:
full_phrases = _split_env_list("TTS_DELAY_FILLERS_ES")
enabled = _env_flag("TTS_DELAY_FILLERS_ENABLED", default=bool(short_phrases or full_phrases))
threshold_ms_raw = os.environ.get("TTS_DELAY_FILLERS_THRESHOLD_MS", "1200").strip() or "1200"
try:
threshold_ms = max(0, int(threshold_ms_raw))
except ValueError:
logger.warning("Invalid TTS_DELAY_FILLERS_THRESHOLD_MS=%r; using 1200", threshold_ms_raw)
threshold_ms = 1200
return DelayFillerConfig(
enabled=enabled and bool(short_phrases or full_phrases),
threshold_s=threshold_ms / 1000.0,
short_phrases=short_phrases,
full_phrases=full_phrases,
)
def _chunk_to_bytes(chunk: Any) -> bytes:
audio = getattr(chunk, "audio", chunk)
if isinstance(audio, np.ndarray) or hasattr(audio, "tobytes"):
return audio.tobytes()
if isinstance(audio, bytes):
return audio
raise TypeError(f"Unsupported delay filler audio chunk type: {type(audio)!r}")
def pregenerate_delay_filler_clips(tts_handler: Any, phrases: list[str]) -> list[bytes]:
clips: list[bytes] = []
for phrase in phrases:
try:
clip = b"".join(
_chunk_to_bytes(chunk)
for chunk in tts_handler.process(TTSInput(text=phrase, language_code="es"))
)
except Exception:
logger.exception("Failed to pregenerate TTS delay filler for phrase %r", phrase)
continue
if not clip:
logger.warning("Discarding empty TTS delay filler for phrase %r", phrase)
continue
clips.append(clip)
return clips
@@ -0,0 +1,11 @@
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",
]
@@ -0,0 +1,199 @@
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
pcm_bytes = resample(pcm_bytes, client_in_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 encode_audio_chunk(self, conn_id: str, audio: bytes) -> list[ServerEvent]:
"""Encode a raw PCM audio chunk, emitting ResponseCreated on the first chunk.
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.
"""
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"),
)
)
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
@@ -0,0 +1,57 @@
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)
@@ -0,0 +1,128 @@
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",
),
)
]
@@ -0,0 +1,348 @@
from __future__ import annotations
import logging
from time import perf_counter
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 = []
st.response_audio_wait_started_at_s = None
st.delay_filler_sent = False
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)
if response_wants_audio(event.response):
st.response_audio_wait_started_at_s = perf_counter()
st.delay_filler_sent = False
else:
st.response_audio_wait_started_at_s = None
st.delay_filler_sent = False
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
@@ -0,0 +1,69 @@
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
# Forza il bypass della voce inviata dal browser
if hasattr(s, "voice") and s.voice is not None:
s.voice = None
elif isinstance(s, dict) and "voice" in s:
s["voice"] = None
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,
)
@@ -0,0 +1,70 @@
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.pipeline.cancel_scope import CancelScope
class SessionState(BaseModel):
"""Per-websocket ephemeral state.
Created when the route handler claims a PipelineUnit on `ws.accept()`; dropped
when the websocket disconnects. Holding the websocket 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.
`drained` is set by the send loop when SESSION_END travels through the handler
chain back to the output queue; the route handler 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)
websocket: Any
session_id: str = ""
pending_output_item: Any = None
drained: asyncio.Event = Field(default_factory=asyncio.Event)
# Wall-clock time when the ws 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]
delay_filler_threshold_s: float = 0.0
delay_filler_short_clips: list[bytes] = Field(default_factory=list)
delay_filler_full_clips: list[bytes] = Field(default_factory=list)
http_tts_reserved: bool = False
session: Optional[SessionState] = None
@@ -0,0 +1,81 @@
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)
@@ -0,0 +1,61 @@
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()
@@ -0,0 +1,560 @@
import logging
from collections.abc import Mapping
from queue import Queue
from threading import Event as ThreadingEvent
from time import perf_counter
from typing import Any, Callable, Literal, Optional, TypeVar, Union
from openai.types.realtime import (
ConversationItem,
ConversationItemCreatedEvent,
ConversationItemCreateEvent,
ConversationItemInputAudioTranscriptionCompletedEvent,
ConversationItemInputAudioTranscriptionDeltaEvent,
InputAudioBufferAppendEvent,
InputAudioBufferSpeechStartedEvent,
InputAudioBufferSpeechStoppedEvent,
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,
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,
"session.update": SessionUpdateEvent,
"conversation.item.create": ConversationItemCreateEvent,
"response.create": ResponseCreateEvent,
"response.cancel": ResponseCancelEvent,
}
ClientEvent = Union[
InputAudioBufferAppendEvent,
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)
response_audio_wait_started_at_s: float | None = None
delay_filler_sent: bool = False
delay_filler_cursor: int = 0
delay_filler_last_kind: str | None = None
delay_filler_last_index: int | None = None
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,
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:
# ── Retrocompatibilità session.update ─────────────────────────────────
# Il formato standard OpenAI Realtime GA annida instructions/tools/tool_choice
# DENTRO "session": { ... }. Per tollerare anche client che li mandano
# alla RADICE (es. bug prima della Fase 2 MCP), ricostruiamo qui il payload
# corretto e riproviamo la validazione. Funziona SOLO per session.update.
if event_type == "session.update":
try:
raw_dict = dict(raw) if not isinstance(raw, dict) else raw
session_inner = raw_dict.get("session")
if isinstance(session_inner, dict):
normalized_inner = dict(session_inner)
else:
normalized_inner = {}
# Promuovi campi noti dalla root → inner se mancanti (fallback flat)
SESSION_KNOWN_TOP_LEVEL = (
"instructions",
"tools",
"tool_choice",
"model",
"voice",
"modalities",
"input_audio_format",
"output_audio_format",
"input_audio_transcription",
"turn_detection",
"input_audio_noise_reduction",
)
promoted = 0
for key in SESSION_KNOWN_TOP_LEVEL:
if key in raw_dict and key not in normalized_inner:
normalized_inner[key] = raw_dict[key]
promoted += 1
if promoted > 0 or not isinstance(session_inner, dict):
rebuilt = {"type": "session.update", "session": normalized_inner}
logger.info(
"session.update parser fallback: promoted %d flat root keys → nested session wrapper. Re-validating.",
promoted,
)
return model_cls.model_validate(rebuilt) # type: ignore[return-value]
except Exception as fallback_err:
logger.error(
"session.update fallback validation also failed: %s",
fallback_err,
)
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 handle_audio_commit(self, conn_id: str) -> RealtimeErrorEvent | None:
return self.audio.handle_audio_commit(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, 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
st.response_audio_wait_started_at_s = perf_counter()
st.delay_filler_sent = False
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
# ── 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 response, and ``finish_response`` is itself
a no-op once the slot is closed, so a later EndOfResponse-driven close does
nothing.
"""
logger.info("Response failed: %s", event.message)
if not self._state(conn_id).in_response:
return []
events: list[ServerEvent] = [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"),
)
@@ -0,0 +1,12 @@
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()
@@ -0,0 +1,915 @@
import asyncio
import io
import logging
import random
import shutil
import subprocess
import time
import uuid
import wave
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, Literal, TypeVar
import numpy as np
from fastapi import FastAPI, HTTPException, Response, WebSocket, WebSocketDisconnect
from openai.types.realtime import (
ConversationItemCreateEvent,
InputAudioBufferAppendEvent,
InputAudioBufferCommitEvent,
ResponseCancelEvent,
ResponseCreateEvent,
SessionUpdateEvent,
)
from pydantic import BaseModel, ConfigDict, Field
from starlette.websockets import WebSocketState
from speech_to_speech.api.openai_realtime.pipeline_unit import PipelineUnit, SessionState
from speech_to_speech.api.openai_realtime.service import ServerEvent, build_error_event
from speech_to_speech.pipeline.control import SESSION_END, PipelineControlMessage, is_control_message
from speech_to_speech.pipeline.events import (
AssistantTextEvent,
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, TTSInput
from speech_to_speech.utils.utils import response_wants_audio
logger = logging.getLogger(__name__)
MAX_AUDIO_BATCH_BYTES = 6400
HTTP_TTS_SAMPLE_RATE = 16000
# 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")
class SpeechSynthesisRequest(BaseModel):
model_config = ConfigDict(populate_by_name=True)
input: str = Field(min_length=1, description="Text to synthesize")
model: str | None = Field(default=None, description="Reserved for OpenAI-compatible clients")
voice: str | None = Field(default=None, description="Reserved for future voice override support")
response_format: Literal["wav", "mp3"] | None = Field(
default=None,
description="Desired output audio format",
)
format: Literal["wav", "mp3"] | None = Field(
default=None,
alias="format",
description="Alias for response_format",
)
language: str | None = Field(default=None, description="Optional target language hint")
def resolved_format(self) -> Literal["wav", "mp3"]:
return self.response_format or self.format or "wav"
async def _send_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}")
async def _send_events(ws: WebSocket, events: list[ServerEvent]) -> None:
for event in events:
await _send_event(ws, event)
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, 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(
ws: WebSocket | 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 ws is not None and events:
await _send_events(ws, 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 _mark_audio_started(unit: PipelineUnit, session_id: str | None) -> None:
if session_id is None:
return
st = unit.service._state(session_id)
st.response_pending = False
st.response_audio_wait_started_at_s = None
def _get_qwen_tts_handler(unit: PipelineUnit) -> Any | None:
for handler in unit.handlers:
if handler.__class__.__name__ == "Qwen3TTSHandler":
return handler
return None
def _encode_wav_bytes(pcm_bytes: bytes, sample_rate: int = HTTP_TTS_SAMPLE_RATE) -> bytes:
buffer = io.BytesIO()
with wave.open(buffer, "wb") as wav_file:
wav_file.setnchannels(1)
wav_file.setsampwidth(2)
wav_file.setframerate(sample_rate)
wav_file.writeframes(pcm_bytes)
return buffer.getvalue()
def _encode_mp3_bytes(pcm_bytes: bytes, sample_rate: int = HTTP_TTS_SAMPLE_RATE) -> bytes:
ffmpeg_path = shutil.which("ffmpeg")
if ffmpeg_path is None:
raise RuntimeError("ffmpeg is required to return MP3 audio")
result = subprocess.run(
[
ffmpeg_path,
"-hide_banner",
"-loglevel",
"error",
"-f",
"s16le",
"-ar",
str(sample_rate),
"-ac",
"1",
"-i",
"pipe:0",
"-f",
"mp3",
"-codec:a",
"libmp3lame",
"pipe:1",
],
input=pcm_bytes,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
check=False,
)
if result.returncode != 0:
stderr = result.stderr.decode("utf-8", errors="ignore").strip()
raise RuntimeError(f"ffmpeg mp3 encoding failed: {stderr or 'unknown error'}")
return result.stdout
def _collect_http_tts_pcm(unit: PipelineUnit, request: SpeechSynthesisRequest) -> bytes:
handler = _get_qwen_tts_handler(unit)
if handler is None:
raise RuntimeError("Qwen3-TTS handler not configured on the reserved pipeline unit")
_clean_unit(unit)
tts_input = TTSInput(
text=request.input.strip(),
language_code=request.language,
turn_id=f"http-tts-{uuid.uuid4().hex}",
turn_revision=0,
)
chunks: list[np.ndarray] = []
for chunk in handler.process(tts_input):
payload = chunk.audio if isinstance(chunk, AudioOutput) else chunk
if isinstance(payload, bytes):
if payload in (AUDIO_RESPONSE_DONE, PIPELINE_END):
continue
chunks.append(np.frombuffer(payload, dtype=np.int16).copy())
continue
array = np.asarray(payload, dtype=np.int16).reshape(-1)
if array.size:
chunks.append(array.copy())
_clean_unit(unit)
if not chunks:
raise RuntimeError("Qwen3-TTS produced no audio")
return np.concatenate(chunks).astype(np.int16, copy=False).tobytes()
async def _maybe_send_delay_filler(ws: WebSocket | None, unit: PipelineUnit, session_id: str | None) -> bool:
if (
ws is None
or session_id is None
or unit.delay_filler_threshold_s <= 0
or not (unit.delay_filler_short_clips or unit.delay_filler_full_clips)
):
return False
st = unit.service._state(session_id)
started_at_s = st.response_audio_wait_started_at_s
if started_at_s is None or st.delay_filler_sent or unit.response_playing.is_set():
return False
if not (st.response_pending or st.in_response):
return False
if not response_wants_audio(st.current_response_params):
return False
if time.perf_counter() - started_at_s < unit.delay_filler_threshold_s:
return False
filler_groups: list[tuple[str, list[bytes], float]] = []
if unit.delay_filler_short_clips:
filler_groups.append(("short", unit.delay_filler_short_clips, 0.65))
if unit.delay_filler_full_clips:
filler_groups.append(("full", unit.delay_filler_full_clips, 0.35))
if len(filler_groups) == 1:
filler_kind, filler_clips, _ = filler_groups[0]
else:
filler_kind, filler_clips, _ = random.choices(
filler_groups,
weights=[group[2] for group in filler_groups],
k=1,
)[0]
if len(filler_clips) == 1:
filler_index = 0
else:
candidate_indexes = list(range(len(filler_clips)))
if st.delay_filler_last_kind == filler_kind and st.delay_filler_last_index in candidate_indexes:
candidate_indexes.remove(st.delay_filler_last_index)
filler_index = random.choice(candidate_indexes)
clip = filler_clips[filler_index]
st.delay_filler_cursor += 1
st.delay_filler_last_kind = filler_kind
st.delay_filler_last_index = filler_index
st.delay_filler_sent = True
_mark_audio_started(unit, session_id)
for offset in range(0, len(clip), MAX_AUDIO_BATCH_BYTES):
batch = clip[offset : offset + MAX_AUDIO_BATCH_BYTES]
if batch:
await _send_events(ws, unit.service.encode_audio_chunk(session_id, batch))
unit.response_playing.set()
unit.should_listen.set()
logger.info("Pipeline %d: sent delay filler clip", unit.index)
return True
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)")
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:
try:
await sess.websocket.close()
except Exception:
pass
app = FastAPI(lifespan=lifespan)
# 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 _claim_unit(ws: WebSocket) -> 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().
"""
for unit in pool:
if unit.session is None and not unit.http_tts_reserved:
unit.session = SessionState(websocket=ws)
return unit
return None
def _reserve_http_tts_unit() -> PipelineUnit | None:
for unit in pool:
if unit.session is None and not unit.http_tts_reserved and _get_qwen_tts_handler(unit) is not None:
unit.http_tts_reserved = True
return unit
return None
def _has_qwen_tts_capacity() -> bool:
return any(_get_qwen_tts_handler(unit) is not None for unit in pool)
@app.post("/v1/audio/speech")
async def speech_endpoint(request: SpeechSynthesisRequest) -> Response:
if not _has_qwen_tts_capacity():
raise HTTPException(
status_code=501,
detail="Qwen TTS is not configured on this server instance.",
)
unit = _reserve_http_tts_unit()
if unit is None:
raise HTTPException(
status_code=503,
detail="All Qwen TTS pipeline slots are currently busy.",
)
try:
pcm_bytes = await asyncio.to_thread(_collect_http_tts_pcm, unit, request)
fmt = request.resolved_format()
if fmt == "mp3":
payload = await asyncio.to_thread(_encode_mp3_bytes, pcm_bytes, HTTP_TTS_SAMPLE_RATE)
media_type = "audio/mpeg"
extension = "mp3"
else:
payload = _encode_wav_bytes(pcm_bytes, HTTP_TTS_SAMPLE_RATE)
media_type = "audio/wav"
extension = "wav"
except HTTPException:
raise
except Exception as e:
logger.error("HTTP speech synthesis failed on pipeline %d: %s", unit.index, e, exc_info=True)
raise HTTPException(status_code=500, detail=f"TTS synthesis failed: {e}") from e
finally:
unit.http_tts_reserved = False
return Response(
content=payload,
media_type=media_type,
headers={
"Content-Disposition": f'inline; filename="speech.{extension}"',
"X-Audio-Sample-Rate": str(HTTP_TTS_SAMPLE_RATE),
},
)
@app.websocket("/v1/realtime")
async def realtime_endpoint(ws: WebSocket) -> None:
await ws.accept()
unit = _claim_unit(ws)
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_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_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
event = unit.service.parse_client_event(raw)
if event is None:
await _send_event(
ws,
unit.service.make_error(
f"Unknown or invalid event: {raw.get('type')}", "unknown_or_invalid_event"
),
)
continue
if isinstance(event, InputAudioBufferAppendEvent):
chunks = unit.service.handle_audio_append(session_id, event)
rt_cfg = unit.service._state(session_id).runtime_config
for chunk in chunks:
unit.input_queue.put((chunk, rt_cfg))
elif isinstance(event, InputAudioBufferCommitEvent):
err = unit.service.handle_audio_commit(session_id)
if err:
await _send_event(ws, err)
elif isinstance(event, SessionUpdateEvent):
err = unit.service.handle_session_update(session_id, event)
if err:
await _send_event(ws, err)
elif isinstance(event, ConversationItemCreateEvent):
events = unit.service.handle_conversation_item_create(session_id, event)
if events:
await _send_events(ws, events)
elif isinstance(event, ResponseCreateEvent):
result = unit.service.handle_response_create(session_id, event)
if result:
if result.type != "error":
unit.cancel_scope.new_response()
await _send_event(ws, result)
elif isinstance(event, ResponseCancelEvent):
was_active = unit.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)
events = unit.service.handle_response_cancel(session_id)
if events:
await _send_events(ws, events)
unit.response_playing.clear()
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.
old_session = unit.session
if old_session is not None:
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))
# Spawn the drain-and-release as a separate task so the route handler's
# 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.
# asyncio holds tasks weakly — keep a reference until it completes.
task = asyncio.create_task(_release_unit_after_drain(unit, old_session, session_id))
release_tasks.add(task)
task.add_done_callback(release_tasks.discard)
@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 and u.http_tts_reserved:
return {"index": u.index, "state": "http_tts", "session_id": None}
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],
}
async def _send_loop_for(unit: PipelineUnit) -> None:
"""Per-pipeline send loop. Polls this unit's output queues and forwards
to the websocket currently attached via unit.session.
Per-session scratch (pending_output_item) lives on SessionState, so it
disappears together with the websocket 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 websocket is still valid until ws.close() returns).
session = unit.session
ws = session.websocket 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 ws is not None and isinstance(text_msg, PipelineEvent) and session_id:
events = unit.service.dispatch_pipeline_event(session_id, text_msg)
if events:
await _send_events(ws, events)
if is_speech_start and (was_in_response or was_response_pending):
active_cfg = unit.service._state(session_id).runtime_config if session_id else None
if text_msg.interrupt_response and (
active_cfg is None or active_cfg.interrupt_response_enabled
):
unit.cancel_scope.cancel()
if session_id:
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
if (
session is not None
and session.pending_output_item is None
and unit.output_queue.empty()
and await _maybe_send_delay_filler(ws, unit, session_id)
):
await asyncio.sleep(0.01)
continue
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(ws, unit, session_id)
if ws is not None and session_id:
await _send_events(ws, 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(ws, unit, session_id)
if ws is not None and session_id:
await _send_events(ws, 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()
_mark_audio_started(unit, session_id)
if ws is not None and session_id:
await _send_events(ws, unit.service.encode_audio_chunk(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
@@ -0,0 +1 @@
@@ -0,0 +1,28 @@
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': <value>} 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."
},
)
@@ -0,0 +1,19 @@
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。."
},
)
@@ -0,0 +1,21 @@
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'."},
)
@@ -0,0 +1,54 @@
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."},
)
@@ -0,0 +1,38 @@
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."},
)
@@ -0,0 +1,43 @@
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."
},
)
@@ -0,0 +1,46 @@
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."
},
)
@@ -0,0 +1,15 @@
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."},
)
@@ -0,0 +1,72 @@
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["whisper", "whisper-mlx", "mlx-audio-whisper", "faster-whisper", "parakeet-tdt", "paraformer"]
] = field(
default="parakeet-tdt",
metadata={
"help": "The STT to use. Either '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."
},
)
@@ -0,0 +1,15 @@
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."},
)
@@ -0,0 +1,39 @@
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."
},
)
@@ -0,0 +1,31 @@
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."},
)
@@ -0,0 +1,100 @@
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."
},
)
@@ -0,0 +1,34 @@
from dataclasses import dataclass, field
from typing import 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'."},
)
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."
},
)
@@ -0,0 +1,20 @@
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."},
)
@@ -0,0 +1,16 @@
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."},
)
@@ -0,0 +1,81 @@
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."
},
)
@@ -0,0 +1,16 @@
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."},
)
@@ -0,0 +1,52 @@
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'."""
},
)
+161
View File
@@ -0,0 +1,161 @@
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
+1
View File
@@ -0,0 +1 @@
@@ -0,0 +1,71 @@
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")
@@ -0,0 +1,85 @@
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")
@@ -0,0 +1,69 @@
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")
@@ -0,0 +1,238 @@
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}")
+55
View File
@@ -0,0 +1,55 @@
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,
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",
"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",
]
+64
View File
@@ -0,0 +1,64 @@
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
+26
View File
@@ -0,0 +1,26 @@
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)
+99
View File
@@ -0,0 +1,99 @@
"""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
from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall
from pydantic import BaseModel, 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)
# ── 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
@@ -0,0 +1,44 @@
"""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
+25
View File
@@ -0,0 +1,25 @@
"""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
+172
View File
@@ -0,0 +1,172 @@
"""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
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
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
+46
View File
@@ -0,0 +1,46 @@
"""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
@@ -0,0 +1,418 @@
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()
File diff suppressed because it is too large Load Diff
View File
+171
View File
@@ -0,0 +1,171 @@
"""
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
+39
View File
@@ -0,0 +1,39 @@
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")
+56
View File
@@ -0,0 +1,56 @@
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