first
This commit is contained in:
@@ -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
|
||||
```
|
||||
@@ -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
|
||||
@@ -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")
|
||||
@@ -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
@@ -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")
|
||||
@@ -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)}"
|
||||
@@ -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))
|
||||
@@ -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(),
|
||||
)
|
||||
@@ -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
|
||||
```
|
||||
@@ -0,0 +1 @@
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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.
|
||||
@@ -0,0 +1 @@
|
||||
# TTS module
|
||||
@@ -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])),
|
||||
)
|
||||
@@ -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")
|
||||
@@ -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")
|
||||
@@ -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
|
||||
@@ -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.
@@ -0,0 +1 @@
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -0,0 +1,3 @@
|
||||
"""speech-to-speech: low-latency speech-to-speech pipeline."""
|
||||
|
||||
__version__ = "0.2.11"
|
||||
@@ -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'."""
|
||||
},
|
||||
)
|
||||
@@ -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
|
||||
@@ -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}")
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
@@ -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
|
||||
@@ -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")
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user