wehub-resource-sync
1360 lines
56 KiB
Python
1360 lines
56 KiB
Python
# Copyright 2025 LiveKit, Inc.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import base64
|
|
import json
|
|
import os
|
|
import time
|
|
import uuid
|
|
import weakref
|
|
from collections.abc import Callable
|
|
from dataclasses import dataclass, field, replace
|
|
from enum import Enum
|
|
from typing import Any, Literal, cast, get_args
|
|
from urllib.parse import urljoin
|
|
|
|
import aiohttp
|
|
|
|
from livekit.agents import tokenize, tts, utils
|
|
from livekit.agents._exceptions import (
|
|
APIConnectionError,
|
|
APIError,
|
|
APIStatusError,
|
|
APITimeoutError,
|
|
)
|
|
from livekit.agents.types import (
|
|
DEFAULT_API_CONNECT_OPTIONS,
|
|
NOT_GIVEN,
|
|
APIConnectOptions,
|
|
NotGivenOr,
|
|
)
|
|
from livekit.agents.utils import is_given
|
|
from livekit.agents.voice.io import TimedString
|
|
|
|
from .log import logger
|
|
from .version import __version__
|
|
|
|
USER_AGENT = f"livekit-agents-py/{__version__}"
|
|
|
|
DEFAULT_BIT_RATE = 64000
|
|
DEFAULT_ENCODING: Encoding = "PCM"
|
|
DEFAULT_MODEL = "inworld-tts-1.5-max"
|
|
DEFAULT_SAMPLE_RATE = 24000
|
|
DEFAULT_URL = "https://api.inworld.ai/"
|
|
DEFAULT_WS_URL = "wss://api.inworld.ai/"
|
|
DEFAULT_VOICE = "Ashley"
|
|
DEFAULT_TEMPERATURE = 1.0
|
|
DEFAULT_SPEAKING_RATE = 1.0
|
|
DEFAULT_BUFFER_CHAR_THRESHOLD = 120
|
|
DEFAULT_MAX_BUFFER_DELAY_MS = 3000
|
|
NUM_CHANNELS = 1
|
|
|
|
TTSModels = Literal["inworld-tts-2", "inworld-tts-1.5-max"]
|
|
Encoding = Literal["LINEAR16", "PCM", "MP3", "OGG_OPUS", "FLAC"]
|
|
TimestampType = Literal["TIMESTAMP_TYPE_UNSPECIFIED", "WORD", "CHARACTER"]
|
|
_TextNormalizationStr = Literal["APPLY_TEXT_NORMALIZATION_UNSPECIFIED", "ON", "OFF"]
|
|
TextNormalization = _TextNormalizationStr | bool
|
|
TimestampTransportStrategy = Literal["TIMESTAMP_TRANSPORT_STRATEGY_UNSPECIFIED", "SYNC", "ASYNC"]
|
|
DeliveryMode = Literal["DELIVERY_MODE_UNSPECIFIED", "STABLE", "BALANCED", "CREATIVE"]
|
|
|
|
DEFAULT_TIMESTAMP_TRANSPORT_STRATEGY: TimestampTransportStrategy = "ASYNC"
|
|
|
|
|
|
def _validate_str_param(value: object, name: str, literal_type: Any) -> None:
|
|
valid = get_args(literal_type)
|
|
if not isinstance(value, str) or value not in valid:
|
|
raise ValueError(f"Invalid {name}: {value!r}. Must be one of {sorted(valid)}")
|
|
|
|
|
|
def _resolve_text_normalization(value: Any) -> _TextNormalizationStr:
|
|
if isinstance(value, bool):
|
|
return "ON" if value else "OFF"
|
|
_validate_str_param(value, "text_normalization", _TextNormalizationStr)
|
|
return cast(_TextNormalizationStr, value)
|
|
|
|
|
|
@dataclass
|
|
class _TTSOptions:
|
|
model: TTSModels | str
|
|
encoding: Encoding
|
|
voice: str
|
|
sample_rate: int
|
|
bit_rate: int
|
|
speaking_rate: float
|
|
temperature: float
|
|
language: NotGivenOr[str] = NOT_GIVEN
|
|
timestamp_type: NotGivenOr[TimestampType] = NOT_GIVEN
|
|
text_normalization: NotGivenOr[TextNormalization] = NOT_GIVEN
|
|
delivery_mode: NotGivenOr[DeliveryMode] = NOT_GIVEN
|
|
timestamp_transport_strategy: TimestampTransportStrategy = DEFAULT_TIMESTAMP_TRANSPORT_STRATEGY
|
|
buffer_char_threshold: int = DEFAULT_BUFFER_CHAR_THRESHOLD
|
|
max_buffer_delay_ms: int = DEFAULT_MAX_BUFFER_DELAY_MS
|
|
|
|
@property
|
|
def mime_type(self) -> str:
|
|
if self.encoding == "PCM":
|
|
return "audio/pcm"
|
|
elif self.encoding == "MP3":
|
|
return "audio/mpeg"
|
|
elif self.encoding == "OGG_OPUS":
|
|
return "audio/ogg"
|
|
elif self.encoding == "FLAC":
|
|
return "audio/flac"
|
|
else:
|
|
return "audio/wav"
|
|
|
|
|
|
# Shared WebSocket connection infrastructure
|
|
class _ContextState(Enum):
|
|
ACTIVE = "active"
|
|
CREATING = "creating"
|
|
CLOSING = "closing"
|
|
|
|
|
|
@dataclass
|
|
class _ContextInfo:
|
|
context_id: str
|
|
state: _ContextState
|
|
emitter: tts.AudioEmitter | None = None
|
|
waiter: asyncio.Future[None] | None = None
|
|
segment_started: bool = False
|
|
flush_after_push: bool = False
|
|
created_at: float = field(default_factory=time.time)
|
|
close_started_at: float | None = None
|
|
# Cumulative timestamp tracking for monotonic timestamps across generations.
|
|
# When auto_mode is enabled or flush_context() is called, the server resets
|
|
# timestamps to 0 after each generation. We add cumulative_time to maintain
|
|
# monotonically increasing timestamps within an agent turn.
|
|
cumulative_time: float = 0.0
|
|
generation_end_time: float = 0.0
|
|
|
|
|
|
@dataclass
|
|
class _CreateContextMsg:
|
|
context_id: str
|
|
opts: _TTSOptions
|
|
|
|
|
|
@dataclass
|
|
class _SendTextMsg:
|
|
context_id: str
|
|
text: str
|
|
|
|
|
|
@dataclass
|
|
class _FlushContextMsg:
|
|
context_id: str
|
|
|
|
|
|
@dataclass
|
|
class _CloseContextMsg:
|
|
context_id: str
|
|
|
|
|
|
_OutboundMessage = _CreateContextMsg | _SendTextMsg | _FlushContextMsg | _CloseContextMsg
|
|
|
|
|
|
class _InworldConnection:
|
|
"""Manages a single shared WebSocket connection with up to 5 concurrent contexts."""
|
|
|
|
MAX_CONTEXTS = 5
|
|
|
|
def __init__(
|
|
self,
|
|
session: aiohttp.ClientSession,
|
|
ws_url: str,
|
|
authorization: str,
|
|
on_capacity_available: Callable[[], None] | None = None,
|
|
) -> None:
|
|
self._session = session
|
|
self._ws_url = ws_url
|
|
self._authorization = authorization
|
|
self._on_capacity_available = on_capacity_available
|
|
self._ws: aiohttp.ClientWebSocketResponse | None = None
|
|
|
|
# Context pool
|
|
self._contexts: dict[str, _ContextInfo] = {}
|
|
self._context_available = asyncio.Event()
|
|
self._pending_acquisitions: int = 0 # Reserved slots not yet in _contexts
|
|
|
|
# Message queue for outbound messages
|
|
self._outbound_queue: asyncio.Queue[_OutboundMessage] = asyncio.Queue()
|
|
|
|
# Connection state
|
|
self._closed = False
|
|
self._connect_lock = asyncio.Lock()
|
|
self._acquire_lock = asyncio.Lock() # Ensures atomic capacity check + context creation
|
|
self._last_activity: float = time.time()
|
|
|
|
# Background tasks
|
|
self._send_task: asyncio.Task[None] | None = None
|
|
self._recv_task: asyncio.Task[None] | None = None
|
|
self._cleanup_task: asyncio.Task[None] | None = None
|
|
|
|
@property
|
|
def context_count(self) -> int:
|
|
"""Number of contexts using server resources.
|
|
|
|
Includes all states (CREATING, ACTIVE, CLOSING) until contextClosed is received
|
|
from the server and the context is removed from _contexts.
|
|
"""
|
|
return len(self._contexts)
|
|
|
|
@property
|
|
def has_capacity(self) -> bool:
|
|
"""Whether this connection can accept more contexts.
|
|
|
|
Accounts for both active contexts and pending acquisitions (reserved slots
|
|
that haven't completed context creation yet).
|
|
"""
|
|
return (self.context_count + self._pending_acquisitions) < self.MAX_CONTEXTS
|
|
|
|
def reserve_capacity(self) -> None:
|
|
"""Reserve a slot for a pending context acquisition.
|
|
|
|
Call this before releasing the pool lock to prevent other callers from
|
|
over-subscribing this connection. The reservation is released when
|
|
acquire_context() completes (success or failure).
|
|
"""
|
|
self._pending_acquisitions += 1
|
|
|
|
def release_reservation(self) -> None:
|
|
"""Release a previously reserved slot.
|
|
|
|
Called automatically by acquire_context() on success, or manually on failure.
|
|
"""
|
|
if self._pending_acquisitions > 0:
|
|
self._pending_acquisitions -= 1
|
|
|
|
@property
|
|
def is_idle(self) -> bool:
|
|
"""Whether this connection has no active contexts and no pending acquisitions."""
|
|
return self.context_count == 0 and self._pending_acquisitions == 0
|
|
|
|
@property
|
|
def last_activity(self) -> float:
|
|
"""Timestamp of last context acquisition or release."""
|
|
return self._last_activity
|
|
|
|
async def connect(self) -> None:
|
|
"""Establish WebSocket connection and start background loops."""
|
|
async with self._connect_lock:
|
|
if self._ws is not None or self._closed:
|
|
return
|
|
|
|
url = urljoin(self._ws_url, "/tts/v1/voice:streamBidirectional")
|
|
request_id = str(uuid.uuid4())
|
|
self._ws = await self._session.ws_connect(
|
|
url,
|
|
headers={
|
|
"Authorization": self._authorization,
|
|
"X-User-Agent": USER_AGENT,
|
|
"X-Request-Id": request_id,
|
|
},
|
|
)
|
|
logger.debug(
|
|
"Established Inworld TTS WebSocket connection (shared)",
|
|
extra={"request_id": request_id},
|
|
)
|
|
|
|
self._send_task = asyncio.create_task(self._send_loop())
|
|
self._recv_task = asyncio.create_task(self._recv_loop())
|
|
self._cleanup_task = asyncio.create_task(self._cleanup_stale_contexts())
|
|
|
|
async def acquire_context(
|
|
self,
|
|
emitter: tts.AudioEmitter,
|
|
opts: _TTSOptions,
|
|
timeout: float,
|
|
) -> tuple[str, asyncio.Future[None]]:
|
|
"""Acquire a new context for TTS synthesis.
|
|
|
|
Note: Caller should check has_capacity before calling this method when using
|
|
a connection pool. This method will still wait if at capacity, but the pool
|
|
should route to connections with available capacity first.
|
|
"""
|
|
await self.connect()
|
|
|
|
# Fail fast if connection is closed
|
|
if self._closed or self._ws is None:
|
|
raise APIConnectionError("Connection is closed")
|
|
|
|
start_time = time.time()
|
|
|
|
while True:
|
|
# Check closed state at start of each iteration
|
|
if self._closed or self._ws is None:
|
|
raise APIConnectionError("Connection is closed")
|
|
|
|
# Use lock to ensure atomic capacity check + context creation
|
|
async with self._acquire_lock:
|
|
# Check raw context_count, not has_capacity, because we ARE one of
|
|
# the pending acquisitions - has_capacity would always be false
|
|
if self.context_count < self.MAX_CONTEXTS:
|
|
self._last_activity = time.time()
|
|
ctx_id = utils.shortuuid()
|
|
waiter: asyncio.Future[None] = asyncio.get_event_loop().create_future()
|
|
|
|
ctx_info = _ContextInfo(
|
|
context_id=ctx_id,
|
|
state=_ContextState.CREATING,
|
|
emitter=emitter,
|
|
waiter=waiter,
|
|
flush_after_push=opts.encoding == "PCM",
|
|
)
|
|
self._contexts[ctx_id] = ctx_info
|
|
# Release reservation now that we have a real context
|
|
self.release_reservation()
|
|
|
|
await self._outbound_queue.put(_CreateContextMsg(context_id=ctx_id, opts=opts))
|
|
return ctx_id, waiter
|
|
|
|
# No capacity - wait outside the lock
|
|
elapsed = time.time() - start_time
|
|
remaining = timeout - elapsed
|
|
if remaining <= 0:
|
|
raise APITimeoutError()
|
|
|
|
try:
|
|
await asyncio.wait_for(self._context_available.wait(), timeout=remaining)
|
|
# Clear after wait returns to avoid lost-wakeup race
|
|
self._context_available.clear()
|
|
|
|
# Check closed state after waking
|
|
if self._closed or self._ws is None:
|
|
raise APIConnectionError("Connection is closed")
|
|
except asyncio.TimeoutError:
|
|
raise APITimeoutError() from None
|
|
|
|
def send_text(self, context_id: str, text: str) -> None:
|
|
"""Queue text to be sent to a context."""
|
|
try:
|
|
self._outbound_queue.put_nowait(_SendTextMsg(context_id=context_id, text=text))
|
|
except asyncio.QueueFull:
|
|
logger.warning("Outbound queue full, dropping text")
|
|
|
|
def flush_context(self, context_id: str) -> None:
|
|
"""Queue a flush message for a context."""
|
|
try:
|
|
self._outbound_queue.put_nowait(_FlushContextMsg(context_id=context_id))
|
|
except asyncio.QueueFull:
|
|
logger.warning("Outbound queue full, dropping flush")
|
|
|
|
def _release_context(self, context_id: str) -> None:
|
|
"""Remove a context and signal that capacity is available."""
|
|
self._contexts.pop(context_id, None)
|
|
self._last_activity = time.time()
|
|
self._context_available.set()
|
|
if self._on_capacity_available:
|
|
self._on_capacity_available()
|
|
|
|
def close_context(self, context_id: str) -> None:
|
|
"""Queue a close message for a context (removes from pool)."""
|
|
ctx = self._contexts.get(context_id)
|
|
if ctx:
|
|
ctx.state = _ContextState.CLOSING
|
|
ctx.close_started_at = time.time()
|
|
try:
|
|
self._outbound_queue.put_nowait(_CloseContextMsg(context_id=context_id))
|
|
except asyncio.QueueFull:
|
|
logger.warning("Outbound queue full, dropping close")
|
|
|
|
async def _send_loop(self) -> None:
|
|
"""Process outbound messages to WebSocket."""
|
|
try:
|
|
while not self._closed and self._ws:
|
|
try:
|
|
msg = await asyncio.wait_for(self._outbound_queue.get(), timeout=1.0)
|
|
except asyncio.TimeoutError:
|
|
continue
|
|
|
|
if self._closed or not self._ws:
|
|
break
|
|
|
|
if isinstance(msg, _CreateContextMsg):
|
|
opts = msg.opts
|
|
pkt: dict[str, Any] = {
|
|
"create": {
|
|
"voiceId": opts.voice,
|
|
"modelId": opts.model,
|
|
"audioConfig": {
|
|
"audioEncoding": opts.encoding,
|
|
"sampleRateHertz": opts.sample_rate,
|
|
"bitrate": opts.bit_rate,
|
|
"speakingRate": opts.speaking_rate,
|
|
},
|
|
"temperature": opts.temperature,
|
|
"bufferCharThreshold": opts.buffer_char_threshold,
|
|
"maxBufferDelayMs": opts.max_buffer_delay_ms,
|
|
"timestampTransportStrategy": opts.timestamp_transport_strategy,
|
|
},
|
|
"contextId": msg.context_id,
|
|
}
|
|
if is_given(opts.language):
|
|
pkt["create"]["language"] = opts.language
|
|
if is_given(opts.timestamp_type):
|
|
pkt["create"]["timestampType"] = opts.timestamp_type
|
|
if is_given(opts.text_normalization):
|
|
pkt["create"]["applyTextNormalization"] = opts.text_normalization
|
|
if is_given(opts.delivery_mode):
|
|
pkt["create"]["deliveryMode"] = opts.delivery_mode
|
|
# Always enable auto_mode since we always use SentenceTokenizer
|
|
pkt["create"]["autoMode"] = True
|
|
await self._ws.send_str(json.dumps(pkt))
|
|
|
|
elif isinstance(msg, _SendTextMsg):
|
|
pkt = {
|
|
"send_text": {"text": msg.text},
|
|
"contextId": msg.context_id,
|
|
}
|
|
await self._ws.send_str(json.dumps(pkt))
|
|
|
|
elif isinstance(msg, _FlushContextMsg):
|
|
pkt = {"flush_context": {}, "contextId": msg.context_id}
|
|
await self._ws.send_str(json.dumps(pkt))
|
|
|
|
elif isinstance(msg, _CloseContextMsg):
|
|
pkt = {"close_context": {}, "contextId": msg.context_id}
|
|
await self._ws.send_str(json.dumps(pkt))
|
|
|
|
except Exception as e:
|
|
logger.error("Inworld send loop error", exc_info=e)
|
|
finally:
|
|
if not self._closed:
|
|
await self._handle_connection_error(APIConnectionError())
|
|
|
|
async def _recv_loop(self) -> None:
|
|
"""Process inbound messages and route to correct context."""
|
|
try:
|
|
while not self._closed and self._ws:
|
|
try:
|
|
msg = await self._ws.receive(timeout=60.0)
|
|
except asyncio.TimeoutError:
|
|
continue
|
|
|
|
if msg.type in (
|
|
aiohttp.WSMsgType.CLOSED,
|
|
aiohttp.WSMsgType.CLOSE,
|
|
aiohttp.WSMsgType.CLOSING,
|
|
):
|
|
break
|
|
|
|
if msg.type != aiohttp.WSMsgType.TEXT:
|
|
continue
|
|
|
|
data = json.loads(msg.data)
|
|
|
|
# Check for errors that are not associated with a context
|
|
if err_obj := data.get("error"):
|
|
logger.warning(
|
|
"Received error from Inworld",
|
|
extra={
|
|
"error_code": err_obj.get("code", 0),
|
|
"error_message": err_obj.get("message", "Unknown error"),
|
|
},
|
|
)
|
|
continue
|
|
|
|
result = data.get("result", {})
|
|
context_id = result.get("contextId")
|
|
|
|
ctx = self._contexts.get(context_id) if context_id else None
|
|
|
|
# Check for errors in status
|
|
status = result.get("status", {})
|
|
if status.get("code", 0) != 0:
|
|
error_msg = status.get("message", "Unknown error")
|
|
error = APIError(f"Inworld error: {error_msg}")
|
|
logger.warning(
|
|
"Received error from Inworld",
|
|
extra={
|
|
"context_id": context_id,
|
|
"error_code": status.get("code"),
|
|
"error_message": error_msg,
|
|
"context_state": ctx.state.value if ctx else "unknown",
|
|
"context_known": ctx is not None,
|
|
},
|
|
)
|
|
if ctx:
|
|
if ctx.waiter and not ctx.waiter.done():
|
|
ctx.waiter.set_exception(error)
|
|
self._release_context(context_id)
|
|
continue
|
|
|
|
if not ctx:
|
|
continue
|
|
|
|
if "contextCreated" in result:
|
|
ctx.state = _ContextState.ACTIVE
|
|
logger.info(
|
|
"Context created confirmed by server",
|
|
extra={"context_id": context_id},
|
|
)
|
|
continue
|
|
|
|
if audio_chunk := result.get("audioChunk"):
|
|
if ctx.emitter:
|
|
if not ctx.segment_started:
|
|
ctx.emitter.start_segment(segment_id=context_id)
|
|
ctx.segment_started = True
|
|
|
|
# Adjust timestamps for cumulative offset
|
|
if timestamp_info := audio_chunk.get("timestampInfo"):
|
|
if word_align := timestamp_info.get("wordAlignment"):
|
|
raw_words = word_align.get("words", [])
|
|
raw_starts = word_align.get("wordStartTimeSeconds", [])
|
|
raw_ends = word_align.get("wordEndTimeSeconds", [])
|
|
logger.debug(
|
|
"Raw timestamps from server",
|
|
extra={
|
|
"context_id": context_id,
|
|
"cumulative_offset": ctx.cumulative_time,
|
|
"raw_words": raw_words,
|
|
"raw_starts": raw_starts,
|
|
"raw_ends": raw_ends,
|
|
},
|
|
)
|
|
|
|
timed_strings = _parse_timestamp_info(
|
|
timestamp_info, cumulative_time=ctx.cumulative_time
|
|
)
|
|
# Track generation end time from last word for cumulative offset
|
|
if timed_strings:
|
|
last_ts = timed_strings[-1]
|
|
if utils.is_given(last_ts.end_time):
|
|
ctx.generation_end_time = last_ts.end_time
|
|
|
|
logger.debug(
|
|
"Adjusted timestamps (with cumulative offset)",
|
|
extra={
|
|
"context_id": context_id,
|
|
"words": [str(ts) for ts in timed_strings],
|
|
"adjusted_starts": [ts.start_time for ts in timed_strings],
|
|
"adjusted_ends": [ts.end_time for ts in timed_strings],
|
|
"generation_end_time": ctx.generation_end_time,
|
|
},
|
|
)
|
|
|
|
for ts in timed_strings:
|
|
ctx.emitter.push_timed_transcript(ts)
|
|
|
|
if audio_content := audio_chunk.get("audioContent"):
|
|
ctx.emitter.push(base64.b64decode(audio_content))
|
|
if ctx.flush_after_push:
|
|
ctx.emitter.flush()
|
|
continue
|
|
|
|
if "flushCompleted" in result:
|
|
# Signals the end of a generation, subsequent timestampes from the server
|
|
# will reset offset to 0. We need to update the cumulative time to the
|
|
# generation end time to maintain monotonically increasing timestamps
|
|
# within the agent turn.
|
|
if ctx.generation_end_time > 0:
|
|
logger.debug(
|
|
"flushCompleted - updating cumulative time",
|
|
extra={
|
|
"context_id": context_id,
|
|
"old_cumulative_time": ctx.cumulative_time,
|
|
"new_cumulative_time": ctx.generation_end_time,
|
|
},
|
|
)
|
|
ctx.cumulative_time = ctx.generation_end_time
|
|
continue
|
|
|
|
if "contextClosed" in result:
|
|
if ctx.waiter and not ctx.waiter.done():
|
|
ctx.waiter.set_result(None)
|
|
self._release_context(context_id)
|
|
continue
|
|
|
|
except Exception as e:
|
|
logger.error("Inworld recv loop error", exc_info=e)
|
|
finally:
|
|
if not self._closed:
|
|
await self._handle_connection_error(APIConnectionError())
|
|
|
|
async def _cleanup_stale_contexts(self) -> None:
|
|
"""Periodically clean up orphaned contexts stuck in CLOSING state."""
|
|
while not self._closed:
|
|
await asyncio.sleep(60.0)
|
|
now = time.time()
|
|
for ctx in list(self._contexts.values()):
|
|
# Use close_started_at if available, otherwise fall back to created_at
|
|
close_time = ctx.close_started_at or ctx.created_at
|
|
if ctx.state == _ContextState.CLOSING and now - close_time > 120.0:
|
|
if ctx.waiter and not ctx.waiter.done():
|
|
ctx.waiter.set_result(None)
|
|
self._release_context(ctx.context_id)
|
|
|
|
async def _handle_connection_error(self, error: Exception) -> None:
|
|
"""Handle connection-level error by failing all active contexts."""
|
|
for ctx in list(self._contexts.values()):
|
|
if ctx.waiter and not ctx.waiter.done():
|
|
ctx.waiter.set_exception(error)
|
|
self._contexts.clear()
|
|
self._closed = True
|
|
|
|
# Wake local waiters so they can fail fast
|
|
self._context_available.set()
|
|
|
|
# Notify pool so callers blocked on capacity can create replacement connections
|
|
if self._on_capacity_available:
|
|
self._on_capacity_available()
|
|
|
|
async def aclose(self) -> None:
|
|
"""Close the connection and all contexts."""
|
|
self._closed = True
|
|
|
|
# Cancel background tasks
|
|
for task in [self._send_task, self._recv_task, self._cleanup_task]:
|
|
if task:
|
|
task.cancel()
|
|
try:
|
|
await task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
|
|
# Close all contexts and fail their waiters
|
|
for ctx in list(self._contexts.values()):
|
|
if ctx.waiter and not ctx.waiter.done():
|
|
ctx.waiter.cancel()
|
|
|
|
self._contexts.clear()
|
|
|
|
if self._ws:
|
|
await self._ws.close()
|
|
self._ws = None
|
|
|
|
|
|
DEFAULT_MAX_CONNECTIONS = 20
|
|
DEFAULT_IDLE_CONNECTION_TIMEOUT = 300.0 # 5 minutes
|
|
|
|
|
|
class _ConnectionPool:
|
|
"""Manages a pool of _InworldConnection instances for high-concurrency scenarios.
|
|
|
|
Each connection supports up to 5 concurrent contexts. The pool automatically creates
|
|
new connections when all existing ones are at capacity, up to max_connections (default 20).
|
|
Idle connections are automatically cleaned up after idle_timeout seconds.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
session: aiohttp.ClientSession,
|
|
ws_url: str,
|
|
authorization: str,
|
|
max_connections: int = DEFAULT_MAX_CONNECTIONS,
|
|
idle_timeout: float = DEFAULT_IDLE_CONNECTION_TIMEOUT,
|
|
) -> None:
|
|
self._session = session
|
|
self._ws_url = ws_url
|
|
self._authorization = authorization
|
|
self._max_connections = max_connections
|
|
self._idle_timeout = idle_timeout
|
|
|
|
self._connections: list[_InworldConnection] = []
|
|
self._pool_lock = asyncio.Lock()
|
|
self._closed = False
|
|
|
|
# Event signaled when any connection has capacity available
|
|
self._capacity_available = asyncio.Event()
|
|
self._capacity_available.set() # Initially available since we can create connections
|
|
|
|
# Cleanup task
|
|
self._cleanup_task: asyncio.Task[None] | None = None
|
|
|
|
async def acquire_context(
|
|
self,
|
|
emitter: tts.AudioEmitter,
|
|
opts: _TTSOptions,
|
|
timeout: float,
|
|
) -> tuple[str, asyncio.Future[None], _InworldConnection]:
|
|
"""Acquire a context from the pool, creating new connections as needed.
|
|
|
|
Returns:
|
|
tuple of (context_id, waiter future, connection) - the connection is returned
|
|
so the caller can call send_text/flush_context/close_context on it.
|
|
"""
|
|
if self._closed:
|
|
raise APIConnectionError("Connection pool is closed")
|
|
|
|
start_time = time.time()
|
|
remaining_timeout = timeout
|
|
|
|
while True:
|
|
conn: _InworldConnection | None = None
|
|
created_new = False
|
|
|
|
async with self._pool_lock:
|
|
# Start cleanup task if not already running
|
|
if self._cleanup_task is None:
|
|
self._cleanup_task = asyncio.create_task(self._cleanup_idle_connections())
|
|
|
|
# Prune closed connections first
|
|
self._connections = [c for c in self._connections if not c._closed]
|
|
|
|
# First, try to find a connection with capacity
|
|
for existing in self._connections:
|
|
if not existing._closed and existing.has_capacity:
|
|
conn = existing
|
|
break
|
|
|
|
# No available capacity - can we create a new connection?
|
|
if conn is None and len(self._connections) < self._max_connections:
|
|
conn = _InworldConnection(
|
|
session=self._session,
|
|
ws_url=self._ws_url,
|
|
authorization=self._authorization,
|
|
on_capacity_available=self.notify_capacity_available,
|
|
)
|
|
created_new = True
|
|
# Add to pool IMMEDIATELY so other callers can see it
|
|
self._connections.append(conn)
|
|
logger.debug(
|
|
"Created new Inworld connection",
|
|
extra={"pool_size": len(self._connections)},
|
|
)
|
|
|
|
# Reserve capacity BEFORE releasing lock so other callers see it
|
|
if conn:
|
|
conn.reserve_capacity()
|
|
|
|
# Acquire context OUTSIDE the lock to avoid head-of-line blocking
|
|
if conn:
|
|
try:
|
|
ctx_id, waiter = await conn.acquire_context(emitter, opts, remaining_timeout)
|
|
except BaseException:
|
|
# Release reservation since we didn't get a context
|
|
conn.release_reservation()
|
|
# Remove failed new connection from pool
|
|
if created_new:
|
|
async with self._pool_lock:
|
|
if conn in self._connections:
|
|
self._connections.remove(conn)
|
|
await conn.aclose()
|
|
raise
|
|
|
|
return ctx_id, waiter, conn
|
|
|
|
# At max connections and all at capacity - wait for one to free up
|
|
elapsed = time.time() - start_time
|
|
remaining_timeout = timeout - elapsed
|
|
if remaining_timeout <= 0:
|
|
raise APITimeoutError("Timed out waiting for available connection capacity")
|
|
|
|
try:
|
|
await asyncio.wait_for(
|
|
self._capacity_available.wait(),
|
|
timeout=remaining_timeout,
|
|
)
|
|
# Clear after wait returns to avoid lost-wakeup race
|
|
self._capacity_available.clear()
|
|
except asyncio.TimeoutError:
|
|
raise APITimeoutError(
|
|
"Timed out waiting for available connection capacity"
|
|
) from None
|
|
|
|
def notify_capacity_available(self) -> None:
|
|
"""Called when a context is released and capacity may be available."""
|
|
self._capacity_available.set()
|
|
|
|
async def _cleanup_idle_connections(self) -> None:
|
|
"""Periodically close idle connections that have been unused for a while."""
|
|
while not self._closed:
|
|
await asyncio.sleep(60.0) # Check every minute
|
|
|
|
async with self._pool_lock:
|
|
now = time.time()
|
|
# Keep at least one connection, close others if idle
|
|
connections_to_close: list[_InworldConnection] = []
|
|
|
|
for conn in self._connections:
|
|
if (
|
|
conn.is_idle
|
|
and now - conn.last_activity > self._idle_timeout
|
|
and len(self._connections) - len(connections_to_close) > 1
|
|
):
|
|
connections_to_close.append(conn)
|
|
|
|
for conn in connections_to_close:
|
|
self._connections.remove(conn)
|
|
logger.debug(
|
|
"Closing idle Inworld connection",
|
|
extra={"pool_size": len(self._connections)},
|
|
)
|
|
|
|
# Close connections outside the lock
|
|
for conn in connections_to_close:
|
|
await conn.aclose()
|
|
|
|
async def aclose(self) -> None:
|
|
"""Close all connections in the pool."""
|
|
self._closed = True
|
|
|
|
if self._cleanup_task:
|
|
self._cleanup_task.cancel()
|
|
try:
|
|
await self._cleanup_task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
|
|
async with self._pool_lock:
|
|
connections = list(self._connections)
|
|
self._connections.clear()
|
|
|
|
for conn in connections:
|
|
await conn.aclose()
|
|
|
|
|
|
class TTS(tts.TTS):
|
|
def __init__(
|
|
self,
|
|
*,
|
|
api_key: NotGivenOr[str] = NOT_GIVEN,
|
|
voice: NotGivenOr[str] = NOT_GIVEN,
|
|
model: NotGivenOr[TTSModels | str] = NOT_GIVEN,
|
|
encoding: NotGivenOr[Encoding] = NOT_GIVEN,
|
|
bit_rate: NotGivenOr[int] = NOT_GIVEN,
|
|
sample_rate: NotGivenOr[int] = NOT_GIVEN,
|
|
speaking_rate: NotGivenOr[float] = NOT_GIVEN,
|
|
temperature: NotGivenOr[float] = NOT_GIVEN,
|
|
language: NotGivenOr[str] = NOT_GIVEN,
|
|
timestamp_type: NotGivenOr[TimestampType] = NOT_GIVEN,
|
|
text_normalization: NotGivenOr[TextNormalization] = NOT_GIVEN,
|
|
delivery_mode: NotGivenOr[DeliveryMode] = NOT_GIVEN,
|
|
timestamp_transport_strategy: NotGivenOr[TimestampTransportStrategy] = NOT_GIVEN,
|
|
buffer_char_threshold: NotGivenOr[int] = NOT_GIVEN,
|
|
max_buffer_delay_ms: NotGivenOr[int] = NOT_GIVEN,
|
|
base_url: str = DEFAULT_URL,
|
|
ws_url: str = DEFAULT_WS_URL,
|
|
http_session: aiohttp.ClientSession | None = None,
|
|
tokenizer: NotGivenOr[tokenize.SentenceTokenizer] = NOT_GIVEN,
|
|
retain_format: NotGivenOr[bool] = NOT_GIVEN,
|
|
max_connections: int = DEFAULT_MAX_CONNECTIONS,
|
|
idle_connection_timeout: float = DEFAULT_IDLE_CONNECTION_TIMEOUT,
|
|
) -> None:
|
|
"""
|
|
Create a new instance of Inworld TTS.
|
|
|
|
Args:
|
|
api_key (str, optional): The Inworld API key.
|
|
If not provided, it will be read from the INWORLD_API_KEY environment variable.
|
|
voice (str, optional): The voice to use. Defaults to "Ashley".
|
|
model (str, optional): The Inworld model to use. Defaults to "inworld-tts-1.5-max".
|
|
encoding (str, optional): The encoding to use. Defaults to "PCM".
|
|
bit_rate (int, optional): Bits per second of the audio. Defaults to 64000.
|
|
sample_rate (int, optional): The audio sample rate in Hz. Defaults to 24000.
|
|
speaking_rate (float, optional): The speed of the voice, in the range [0.5, 1.5].
|
|
Defaults to 1.0.
|
|
temperature (float, optional): Determines the degree of randomness when sampling audio
|
|
tokens to generate the response. Range (0, 2]. Defaults to 1.0.
|
|
language (str, optional): BCP-47 language tag (e.g., "en-US", "fr-FR", "ja-JP")
|
|
specifying the language that the given voice should speak the text in.
|
|
If not set, the model default applies.
|
|
timestamp_type (str, optional): Controls timestamp metadata returned with the audio.
|
|
Use "WORD" for word-level timestamps or "CHARACTER" for character-level.
|
|
Useful for karaoke-style captions, word highlighting, and lipsync.
|
|
text_normalization (str, optional): Controls text normalization. When "ON", numbers,
|
|
dates, and abbreviations are expanded (e.g., "Dr." -> "Doctor"). When "OFF",
|
|
text is read exactly as written. Defaults to automatic.
|
|
delivery_mode (str, optional): Controls output variation on ``inworld-tts-2`` only.
|
|
One of "DELIVERY_MODE_UNSPECIFIED", "STABLE", "BALANCED", or "CREATIVE".
|
|
The Inworld API ignores ``temperature`` on ``inworld-tts-2`` — use
|
|
``delivery_mode`` to steer output variation on that model instead.
|
|
Defaults to the server-side default ("BALANCED").
|
|
timestamp_transport_strategy (str, optional): Controls how timestamp info is
|
|
transported relative to audio data. "SYNC" returns timestamps in the same
|
|
message as audio data. "ASYNC" allows timestamps to return in trailing
|
|
messages after the audio data. Defaults to "ASYNC".
|
|
buffer_char_threshold (int, optional): For streaming, the minimum number of characters
|
|
in the buffer that automatically triggers audio generation. Defaults to 1000.
|
|
max_buffer_delay_ms (int, optional): For streaming, the maximum time in ms to buffer
|
|
before starting generation. Defaults to 3000.
|
|
base_url (str, optional): The base URL for the Inworld TTS API.
|
|
Defaults to "https://api.inworld.ai/".
|
|
ws_url (str, optional): The WebSocket URL for streaming TTS.
|
|
Defaults to "wss://api.inworld.ai/".
|
|
http_session (aiohttp.ClientSession, optional): The HTTP session to use.
|
|
tokenizer (tokenize.SentenceTokenizer, optional): The tokenizer to use for streaming.
|
|
Defaults to `livekit.agents.tokenize.blingfire.SentenceTokenizer`.
|
|
retain_format (bool, optional): Whether to retain the format of the text when tokenizing.
|
|
Defaults to True.
|
|
max_connections (int, optional): Maximum number of concurrent WebSocket connections.
|
|
Each connection supports up to 5 concurrent synthesis streams. Defaults to 20.
|
|
idle_connection_timeout (float, optional): Time in seconds after which idle connections
|
|
are closed. Defaults to 300 (5 minutes).
|
|
"""
|
|
if not is_given(sample_rate):
|
|
sample_rate = DEFAULT_SAMPLE_RATE
|
|
super().__init__(
|
|
capabilities=tts.TTSCapabilities(
|
|
streaming=True,
|
|
aligned_transcript=is_given(timestamp_type)
|
|
and timestamp_type != "TIMESTAMP_TYPE_UNSPECIFIED",
|
|
),
|
|
sample_rate=sample_rate,
|
|
num_channels=NUM_CHANNELS,
|
|
)
|
|
|
|
key = api_key if is_given(api_key) else os.getenv("INWORLD_API_KEY")
|
|
if not key:
|
|
raise ValueError(
|
|
"Inworld API key is required, either as argument or set"
|
|
" INWORLD_API_KEY environment variable"
|
|
)
|
|
|
|
self._authorization = f"Basic {key}"
|
|
self._base_url = base_url
|
|
self._ws_url = ws_url
|
|
self._session = http_session
|
|
|
|
if is_given(encoding):
|
|
_validate_str_param(encoding, "encoding", Encoding)
|
|
if is_given(timestamp_type):
|
|
_validate_str_param(timestamp_type, "timestamp_type", TimestampType)
|
|
if is_given(text_normalization):
|
|
text_normalization = _resolve_text_normalization(text_normalization)
|
|
if is_given(delivery_mode):
|
|
_validate_str_param(delivery_mode, "delivery_mode", DeliveryMode)
|
|
if is_given(timestamp_transport_strategy):
|
|
_validate_str_param(
|
|
timestamp_transport_strategy,
|
|
"timestamp_transport_strategy",
|
|
TimestampTransportStrategy,
|
|
)
|
|
|
|
self._opts = _TTSOptions(
|
|
voice=voice if is_given(voice) else DEFAULT_VOICE,
|
|
model=model if is_given(model) else DEFAULT_MODEL,
|
|
encoding=encoding if is_given(encoding) else DEFAULT_ENCODING,
|
|
bit_rate=bit_rate if is_given(bit_rate) else DEFAULT_BIT_RATE,
|
|
sample_rate=sample_rate if is_given(sample_rate) else DEFAULT_SAMPLE_RATE,
|
|
speaking_rate=speaking_rate if is_given(speaking_rate) else DEFAULT_SPEAKING_RATE,
|
|
temperature=temperature if is_given(temperature) else DEFAULT_TEMPERATURE,
|
|
language=language,
|
|
timestamp_type=timestamp_type,
|
|
text_normalization=text_normalization,
|
|
delivery_mode=delivery_mode,
|
|
timestamp_transport_strategy=timestamp_transport_strategy
|
|
if is_given(timestamp_transport_strategy)
|
|
else DEFAULT_TIMESTAMP_TRANSPORT_STRATEGY,
|
|
buffer_char_threshold=buffer_char_threshold
|
|
if is_given(buffer_char_threshold)
|
|
else DEFAULT_BUFFER_CHAR_THRESHOLD,
|
|
max_buffer_delay_ms=max_buffer_delay_ms
|
|
if is_given(max_buffer_delay_ms)
|
|
else DEFAULT_MAX_BUFFER_DELAY_MS,
|
|
)
|
|
|
|
self._max_connections = max_connections
|
|
self._idle_connection_timeout = idle_connection_timeout
|
|
self._pool: _ConnectionPool | None = None
|
|
self._pool_lock = asyncio.Lock()
|
|
self._streams = weakref.WeakSet[SynthesizeStream]()
|
|
self._sentence_tokenizer = (
|
|
tokenizer
|
|
if is_given(tokenizer)
|
|
else tokenize.blingfire.SentenceTokenizer(
|
|
retain_format=retain_format if is_given(retain_format) else True
|
|
)
|
|
)
|
|
|
|
class Markup(tts.TTS.Markup):
|
|
# markup delegation lives in the base class, keyed on _provider_key()
|
|
def _provider_key(self) -> str:
|
|
# only inworld-tts-2 understands the markup tags; older models get no
|
|
# markup so the tags aren't injected, converted, or stripped (matches
|
|
# the inference gateway's behavior)
|
|
return "inworld" if "tts-2" in self._tts.model else ""
|
|
|
|
@property
|
|
def model(self) -> str:
|
|
return self._opts.model
|
|
|
|
@property
|
|
def provider(self) -> str:
|
|
return "Inworld"
|
|
|
|
async def _get_pool(self) -> _ConnectionPool:
|
|
"""Get the connection pool, creating if needed."""
|
|
async with self._pool_lock:
|
|
if self._pool is None or self._pool._closed:
|
|
self._pool = _ConnectionPool(
|
|
session=self._ensure_session(),
|
|
ws_url=self._ws_url,
|
|
authorization=self._authorization,
|
|
max_connections=self._max_connections,
|
|
idle_timeout=self._idle_connection_timeout,
|
|
)
|
|
return self._pool
|
|
|
|
def update_options(
|
|
self,
|
|
*,
|
|
voice: NotGivenOr[str] = NOT_GIVEN,
|
|
model: NotGivenOr[TTSModels | str] = NOT_GIVEN,
|
|
encoding: NotGivenOr[Encoding] = NOT_GIVEN,
|
|
bit_rate: NotGivenOr[int] = NOT_GIVEN,
|
|
sample_rate: NotGivenOr[int] = NOT_GIVEN,
|
|
speaking_rate: NotGivenOr[float] = NOT_GIVEN,
|
|
temperature: NotGivenOr[float] = NOT_GIVEN,
|
|
language: NotGivenOr[str] = NOT_GIVEN,
|
|
timestamp_type: NotGivenOr[TimestampType] = NOT_GIVEN,
|
|
text_normalization: NotGivenOr[TextNormalization] = NOT_GIVEN,
|
|
delivery_mode: NotGivenOr[DeliveryMode] = NOT_GIVEN,
|
|
timestamp_transport_strategy: NotGivenOr[TimestampTransportStrategy] = NOT_GIVEN,
|
|
buffer_char_threshold: NotGivenOr[int] = NOT_GIVEN,
|
|
max_buffer_delay_ms: NotGivenOr[int] = NOT_GIVEN,
|
|
) -> None:
|
|
"""
|
|
Update the TTS configuration options.
|
|
|
|
Args:
|
|
voice (str, optional): The voice to use.
|
|
model (str, optional): The Inworld model to use.
|
|
encoding (str, optional): The encoding to use.
|
|
bit_rate (int, optional): Bits per second of the audio.
|
|
sample_rate (int, optional): The audio sample rate in Hz.
|
|
speaking_rate (float, optional): The speed of the voice.
|
|
temperature (float, optional): Determines the degree of randomness when sampling audio
|
|
tokens to generate the response.
|
|
language (str, optional): BCP-47 language tag (e.g., "en-US", "fr-FR").
|
|
timestamp_type (str, optional): Controls timestamp metadata ("WORD" or "CHARACTER").
|
|
text_normalization (str, optional): Controls text normalization ("ON" or "OFF").
|
|
delivery_mode (str, optional): Controls output variation on ``inworld-tts-2`` only
|
|
("DELIVERY_MODE_UNSPECIFIED", "STABLE", "BALANCED", or "CREATIVE"). The Inworld
|
|
API ignores ``temperature`` on TTS-2; use this instead.
|
|
timestamp_transport_strategy (str, optional): Controls timestamp transport strategy
|
|
("SYNC" or "ASYNC").
|
|
buffer_char_threshold (int, optional): For streaming, min characters before triggering.
|
|
max_buffer_delay_ms (int, optional): For streaming, max time to buffer.
|
|
"""
|
|
if is_given(voice):
|
|
self._opts.voice = voice
|
|
if is_given(model):
|
|
self._opts.model = model
|
|
if is_given(encoding):
|
|
_validate_str_param(encoding, "encoding", Encoding)
|
|
self._opts.encoding = encoding
|
|
if is_given(bit_rate):
|
|
self._opts.bit_rate = bit_rate
|
|
if is_given(sample_rate):
|
|
self._opts.sample_rate = sample_rate
|
|
if is_given(speaking_rate):
|
|
self._opts.speaking_rate = speaking_rate
|
|
if is_given(temperature):
|
|
self._opts.temperature = temperature
|
|
if is_given(language):
|
|
self._opts.language = language
|
|
if is_given(timestamp_type):
|
|
_validate_str_param(timestamp_type, "timestamp_type", TimestampType)
|
|
self._opts.timestamp_type = timestamp_type
|
|
if is_given(text_normalization):
|
|
self._opts.text_normalization = _resolve_text_normalization(text_normalization)
|
|
if is_given(delivery_mode):
|
|
_validate_str_param(delivery_mode, "delivery_mode", DeliveryMode)
|
|
self._opts.delivery_mode = delivery_mode
|
|
if is_given(timestamp_transport_strategy):
|
|
_validate_str_param(
|
|
timestamp_transport_strategy,
|
|
"timestamp_transport_strategy",
|
|
TimestampTransportStrategy,
|
|
)
|
|
self._opts.timestamp_transport_strategy = timestamp_transport_strategy
|
|
if is_given(buffer_char_threshold):
|
|
self._opts.buffer_char_threshold = buffer_char_threshold
|
|
if is_given(max_buffer_delay_ms):
|
|
self._opts.max_buffer_delay_ms = max_buffer_delay_ms
|
|
|
|
def _ensure_session(self) -> aiohttp.ClientSession:
|
|
if not self._session:
|
|
self._session = utils.http_context.http_session()
|
|
|
|
return self._session
|
|
|
|
def prewarm(self) -> None:
|
|
asyncio.create_task(self._prewarm_impl())
|
|
|
|
async def _prewarm_impl(self) -> None:
|
|
# Just ensure the pool is created - first acquire will establish a connection
|
|
await self._get_pool()
|
|
|
|
def synthesize(
|
|
self,
|
|
text: str,
|
|
*,
|
|
conn_options: APIConnectOptions = DEFAULT_API_CONNECT_OPTIONS,
|
|
) -> tts.ChunkedStream:
|
|
return ChunkedStream(tts=self, input_text=text, conn_options=conn_options)
|
|
|
|
def stream(
|
|
self, *, conn_options: APIConnectOptions = DEFAULT_API_CONNECT_OPTIONS
|
|
) -> SynthesizeStream:
|
|
stream = SynthesizeStream(tts=self, conn_options=conn_options)
|
|
self._streams.add(stream)
|
|
return stream
|
|
|
|
async def aclose(self) -> None:
|
|
for stream in list(self._streams):
|
|
await stream.aclose()
|
|
|
|
self._streams.clear()
|
|
if self._pool:
|
|
await self._pool.aclose()
|
|
self._pool = None
|
|
|
|
async def list_voices(self, language: str | None = None) -> list[dict[str, Any]]:
|
|
"""
|
|
List all available voices in the workspace associated with the API key.
|
|
|
|
Args:
|
|
language (str, optional): ISO 639-1 language code to filter voices (e.g., 'en', 'es', 'fr').
|
|
"""
|
|
url = urljoin(self._base_url, "tts/v1/voices")
|
|
params = {}
|
|
if language:
|
|
params["filter"] = f"language={language}"
|
|
|
|
async with self._ensure_session().get(
|
|
url,
|
|
headers={
|
|
"Authorization": self._authorization,
|
|
"X-User-Agent": USER_AGENT,
|
|
"X-Request-Id": str(uuid.uuid4()),
|
|
},
|
|
params=params,
|
|
) as resp:
|
|
if not resp.ok:
|
|
error_body = await resp.json()
|
|
raise APIStatusError(
|
|
message=error_body.get("message"),
|
|
status_code=resp.status,
|
|
request_id=None,
|
|
body=None,
|
|
)
|
|
|
|
data = await resp.json()
|
|
return cast(list[dict[str, Any]], data.get("voices", []))
|
|
|
|
|
|
class ChunkedStream(tts.ChunkedStream):
|
|
def __init__(self, *, tts: TTS, input_text: str, conn_options: APIConnectOptions) -> None:
|
|
super().__init__(tts=tts, input_text=input_text, conn_options=conn_options)
|
|
self._tts: TTS = tts
|
|
self._opts = replace(tts._opts)
|
|
|
|
async def _run(self, output_emitter: tts.AudioEmitter) -> None:
|
|
try:
|
|
audio_config: dict[str, Any] = {
|
|
"audioEncoding": self._opts.encoding,
|
|
"bitrate": self._opts.bit_rate,
|
|
"sampleRateHertz": self._opts.sample_rate,
|
|
"temperature": self._opts.temperature,
|
|
"speakingRate": self._opts.speaking_rate,
|
|
}
|
|
|
|
body_params: dict[str, Any] = {
|
|
"text": self._input_text,
|
|
"voiceId": self._opts.voice,
|
|
"modelId": self._opts.model,
|
|
"audioConfig": audio_config,
|
|
}
|
|
if utils.is_given(self._opts.language):
|
|
body_params["language"] = self._opts.language
|
|
if utils.is_given(self._opts.timestamp_type):
|
|
body_params["timestampType"] = self._opts.timestamp_type
|
|
if utils.is_given(self._opts.text_normalization):
|
|
body_params["applyTextNormalization"] = self._opts.text_normalization
|
|
if utils.is_given(self._opts.delivery_mode):
|
|
body_params["deliveryMode"] = self._opts.delivery_mode
|
|
body_params["timestampTransportStrategy"] = self._opts.timestamp_transport_strategy
|
|
|
|
x_request_id = str(uuid.uuid4())
|
|
async with self._tts._ensure_session().post(
|
|
urljoin(self._tts._base_url, "/tts/v1/voice:stream"),
|
|
headers={
|
|
"Authorization": self._tts._authorization,
|
|
"X-User-Agent": USER_AGENT,
|
|
"X-Request-Id": x_request_id,
|
|
},
|
|
json=body_params,
|
|
timeout=aiohttp.ClientTimeout(sock_connect=self._conn_options.timeout),
|
|
# large read_bufsize to avoid `ValueError: Chunk too big`
|
|
read_bufsize=10 * 1024 * 1024,
|
|
) as resp:
|
|
resp.raise_for_status()
|
|
|
|
request_id = utils.shortuuid()
|
|
output_emitter.initialize(
|
|
request_id=request_id,
|
|
sample_rate=self._opts.sample_rate,
|
|
num_channels=NUM_CHANNELS,
|
|
mime_type=self._opts.mime_type,
|
|
)
|
|
|
|
async for raw_line in resp.content:
|
|
line = raw_line.strip()
|
|
if not line:
|
|
continue
|
|
|
|
try:
|
|
data = json.loads(line)
|
|
except json.JSONDecodeError:
|
|
logger.warning("failed to parse Inworld response line: %s", line)
|
|
continue
|
|
|
|
if result := data.get("result"):
|
|
# Handle timestamp info if present
|
|
if timestamp_info := result.get("timestampInfo"):
|
|
timed_strings = _parse_timestamp_info(timestamp_info)
|
|
if timed_strings:
|
|
output_emitter.push_timed_transcript(timed_strings)
|
|
|
|
if audio_content := result.get("audioContent"):
|
|
output_emitter.push(base64.b64decode(audio_content))
|
|
if self._opts.encoding == "PCM":
|
|
output_emitter.flush()
|
|
elif error := data.get("error"):
|
|
raise APIStatusError(
|
|
message=error.get("message"),
|
|
status_code=error.get("code"),
|
|
request_id=x_request_id,
|
|
body=None,
|
|
)
|
|
except asyncio.TimeoutError:
|
|
raise APITimeoutError() from None
|
|
except aiohttp.ClientResponseError as e:
|
|
raise APIStatusError(
|
|
message=e.message, status_code=e.status, request_id=x_request_id, body=None
|
|
) from None
|
|
except Exception as e:
|
|
raise APIConnectionError() from e
|
|
|
|
|
|
class SynthesizeStream(tts.SynthesizeStream):
|
|
def __init__(self, *, tts: TTS, conn_options: APIConnectOptions):
|
|
super().__init__(tts=tts, conn_options=conn_options)
|
|
self._tts: TTS = tts
|
|
self._opts = replace(tts._opts)
|
|
|
|
async def _run(self, output_emitter: tts.AudioEmitter) -> None:
|
|
request_id = utils.shortuuid()
|
|
sent_tokenizer_stream = self._tts._sentence_tokenizer.stream()
|
|
|
|
output_emitter.initialize(
|
|
request_id=request_id,
|
|
sample_rate=self._opts.sample_rate,
|
|
num_channels=NUM_CHANNELS,
|
|
mime_type=self._opts.mime_type,
|
|
stream=True,
|
|
)
|
|
|
|
pool = await self._tts._get_pool()
|
|
context_id, waiter, connection = await pool.acquire_context(
|
|
emitter=output_emitter,
|
|
opts=self._opts,
|
|
timeout=self._conn_options.timeout,
|
|
)
|
|
|
|
async def _input_task() -> None:
|
|
async for data in self._input_ch:
|
|
if isinstance(data, self._FlushSentinel):
|
|
sent_tokenizer_stream.flush()
|
|
continue
|
|
sent_tokenizer_stream.push_text(data)
|
|
sent_tokenizer_stream.end_input()
|
|
|
|
async def _send_task() -> None:
|
|
async for ev in sent_tokenizer_stream:
|
|
text = ev.token
|
|
# Chunk to stay within Inworld's 1000 char limit
|
|
for i in range(0, len(text), 1000):
|
|
connection.send_text(context_id, text[i : i + 1000])
|
|
self._mark_started()
|
|
connection.flush_context(context_id)
|
|
connection.close_context(context_id)
|
|
|
|
tasks = [
|
|
asyncio.create_task(_input_task()),
|
|
asyncio.create_task(_send_task()),
|
|
]
|
|
|
|
try:
|
|
await asyncio.wait_for(waiter, timeout=self._conn_options.timeout + 60)
|
|
except asyncio.TimeoutError:
|
|
connection.close_context(context_id)
|
|
raise APITimeoutError() from None
|
|
except asyncio.CancelledError:
|
|
connection.close_context(context_id)
|
|
raise
|
|
except APIError:
|
|
raise
|
|
except Exception as e:
|
|
logger.error("Inworld stream error", extra={"context_id": context_id, "error": e})
|
|
connection.close_context(context_id)
|
|
raise APIConnectionError() from e
|
|
finally:
|
|
await utils.aio.gracefully_cancel(*tasks)
|
|
await sent_tokenizer_stream.aclose()
|
|
output_emitter.end_input()
|
|
|
|
|
|
def _parse_timestamp_info(
|
|
timestamp_info: dict[str, Any], cumulative_time: float = 0.0
|
|
) -> list[TimedString]:
|
|
"""Parse timestamp info from API response into TimedString objects.
|
|
|
|
Args:
|
|
timestamp_info: The timestamp info from the API response.
|
|
cumulative_time: Offset to add to all timestamps for monotonic ordering
|
|
across multiple generations within a single context.
|
|
"""
|
|
timed_strings: list[TimedString] = []
|
|
|
|
# Handle word-level alignment
|
|
if word_align := timestamp_info.get("wordAlignment"):
|
|
words = word_align.get("words", [])
|
|
starts = word_align.get("wordStartTimeSeconds", [])
|
|
ends = word_align.get("wordEndTimeSeconds", [])
|
|
|
|
for word, start, end in zip(words, starts, ends, strict=False):
|
|
text = f"{word}"
|
|
timed_strings.append(
|
|
TimedString(
|
|
text,
|
|
start_time=cumulative_time + start,
|
|
end_time=cumulative_time + end,
|
|
)
|
|
)
|
|
|
|
# Handle character-level alignment
|
|
if char_align := timestamp_info.get("characterAlignment"):
|
|
chars = char_align.get("characters", [])
|
|
starts = char_align.get("characterStartTimeSeconds", [])
|
|
ends = char_align.get("characterEndTimeSeconds", [])
|
|
|
|
for char, start, end in zip(chars, starts, ends, strict=False):
|
|
timed_strings.append(
|
|
TimedString(
|
|
char,
|
|
start_time=cumulative_time + start,
|
|
end_time=cumulative_time + end,
|
|
)
|
|
)
|
|
|
|
return timed_strings
|