livekit--agents
wehub-resource-sync
510 行
19 KiB
Python
510 行
19 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import os
|
|
import weakref
|
|
from dataclasses import dataclass, replace
|
|
from typing import Any
|
|
|
|
import aiohttp
|
|
import msgpack
|
|
|
|
from livekit.agents import (
|
|
APIConnectionError,
|
|
APIConnectOptions,
|
|
APIStatusError,
|
|
APITimeoutError,
|
|
tokenize,
|
|
tts,
|
|
utils,
|
|
)
|
|
from livekit.agents.types import DEFAULT_API_CONNECT_OPTIONS, NOT_GIVEN, NotGivenOr
|
|
from livekit.agents.utils import is_given
|
|
|
|
from .log import logger
|
|
from .models import LatencyMode, OutputFormat, TTSModels
|
|
from .version import __version__
|
|
|
|
DEFAULT_MODEL: TTSModels = "s2.1-pro"
|
|
DEFAULT_VOICE_ID = "933563129e564b19a115bedd57b7406a"
|
|
DEFAULT_BASE_URL = "https://api.fish.audio"
|
|
NUM_CHANNELS = 1
|
|
USER_AGENT = f"livekit-plugins-fishaudio/{__version__}"
|
|
|
|
# Hold the first N audio chunks before releasing any audio, so playout starts with
|
|
# enough buffered to ride out the gap after Fish's small first chunk (else jitter
|
|
# underruns it into a crackle). Internal stopgap for Fish's cold-start pacing; drop
|
|
# to 1 once inference streams the opening chunks smoothly. Values <= 1 disable it.
|
|
_PREBUFFER_CHUNKS = 2
|
|
|
|
# Fish Audio's default sample rate per output format. Opus only supports 48 kHz;
|
|
# the other formats default to 24 kHz, which matches the previous plugin default.
|
|
_DEFAULT_SAMPLE_RATE: dict[OutputFormat, int] = {
|
|
"opus": 48000,
|
|
"pcm": 24000,
|
|
"wav": 24000,
|
|
"mp3": 32000,
|
|
}
|
|
|
|
|
|
@dataclass
|
|
class _TTSOptions:
|
|
model: TTSModels | str
|
|
output_format: OutputFormat
|
|
sample_rate: int
|
|
voice_id: NotGivenOr[str]
|
|
base_url: str
|
|
api_key: str
|
|
latency_mode: LatencyMode
|
|
chunk_length: int
|
|
speed: NotGivenOr[float]
|
|
volume: NotGivenOr[float]
|
|
|
|
def get_http_url(self, path: str) -> str:
|
|
return f"{self.base_url}{path}"
|
|
|
|
def get_ws_url(self, path: str) -> str:
|
|
return f"{self.base_url.replace('http', 'ws', 1)}{path}"
|
|
|
|
|
|
class TTS(tts.TTS):
|
|
def __init__(
|
|
self,
|
|
*,
|
|
api_key: NotGivenOr[str] = NOT_GIVEN,
|
|
model: TTSModels | str = DEFAULT_MODEL,
|
|
voice_id: NotGivenOr[str] = DEFAULT_VOICE_ID,
|
|
output_format: OutputFormat = "wav",
|
|
sample_rate: NotGivenOr[int] = NOT_GIVEN,
|
|
base_url: NotGivenOr[str] = NOT_GIVEN,
|
|
latency_mode: LatencyMode = "balanced",
|
|
chunk_length: int = 100,
|
|
speed: NotGivenOr[float] = NOT_GIVEN,
|
|
volume: NotGivenOr[float] = NOT_GIVEN,
|
|
tokenizer: NotGivenOr[tokenize.SentenceTokenizer] = NOT_GIVEN,
|
|
http_session: aiohttp.ClientSession | None = None,
|
|
) -> None:
|
|
"""
|
|
Create a new instance of Fish Audio TTS.
|
|
|
|
See https://docs.fish.audio/api-reference/endpoint/websocket/tts-live for more details
|
|
on the Fish Audio Live TTS WebSocket API.
|
|
|
|
Args:
|
|
api_key (NotGivenOr[str]): Fish Audio API key. Reads ``FISH_API_KEY`` if unset.
|
|
model (TTSModels | str): TTS model to use. Defaults to ``"s2.1-pro"``.
|
|
voice_id (NotGivenOr[str]): Voice model ID. Fish Audio's API refers to this
|
|
as ``reference_id``; it's the same value either way.
|
|
output_format (OutputFormat): Audio output format. Defaults to ``"wav"``.
|
|
sample_rate (int): Audio sample rate in Hz.
|
|
base_url (NotGivenOr[str]): Custom base URL. Defaults to ``https://api.fish.audio``.
|
|
latency_mode (LatencyMode): Streaming latency mode. ``"normal"``, ``"balanced"``,
|
|
or ``"low"``. Defaults to ``"balanced"``.
|
|
chunk_length (int): Upper bound on text Fish buffers before auto-synthesizing
|
|
(100–300). With sentence-level flushing this is only hit by sentences longer
|
|
than ``chunk_length``; otherwise audio is produced when each sentence is
|
|
flushed. Defaults to 100.
|
|
speed (NotGivenOr[float]): Speaking rate multiplier (Fish ``prosody.speed``).
|
|
``1.0`` is normal; below 1.0 is slower, above is faster. Unset uses the
|
|
voice's natural pace.
|
|
volume (NotGivenOr[float]): Loudness adjustment in decibels (Fish
|
|
``prosody.volume``). ``0`` is the voice's natural level. Unset leaves it
|
|
unchanged.
|
|
tokenizer (tokenize.SentenceTokenizer): Sentence tokenizer used to detect
|
|
sentence boundaries. Defaults to ``tokenize.blingfire.SentenceTokenizer()``.
|
|
http_session (aiohttp.ClientSession | None): Optional aiohttp session.
|
|
"""
|
|
if is_given(sample_rate):
|
|
if output_format == "opus" and sample_rate != 48000:
|
|
raise ValueError(
|
|
"Fish Audio only supports 48000 Hz for opus output; "
|
|
f"got sample_rate={sample_rate}"
|
|
)
|
|
resolved_sample_rate = sample_rate
|
|
else:
|
|
resolved_sample_rate = _DEFAULT_SAMPLE_RATE[output_format]
|
|
|
|
super().__init__(
|
|
capabilities=tts.TTSCapabilities(streaming=True),
|
|
sample_rate=resolved_sample_rate,
|
|
num_channels=NUM_CHANNELS,
|
|
)
|
|
|
|
fish_api_key = api_key if is_given(api_key) else os.getenv("FISH_API_KEY")
|
|
if not fish_api_key:
|
|
raise ValueError(
|
|
"Fish Audio API key is required, either as argument or set "
|
|
"FISH_API_KEY environment variable"
|
|
)
|
|
|
|
if not 100 <= chunk_length <= 300:
|
|
raise ValueError("chunk_length must be between 100 and 300")
|
|
|
|
self._opts = _TTSOptions(
|
|
model=model,
|
|
output_format=output_format,
|
|
sample_rate=resolved_sample_rate,
|
|
voice_id=voice_id,
|
|
base_url=base_url if is_given(base_url) else DEFAULT_BASE_URL,
|
|
api_key=fish_api_key,
|
|
latency_mode=latency_mode,
|
|
chunk_length=chunk_length,
|
|
speed=speed,
|
|
volume=volume,
|
|
)
|
|
|
|
self._session = http_session
|
|
self._pool = utils.ConnectionPool[aiohttp.ClientWebSocketResponse](
|
|
connect_cb=self._connect_ws,
|
|
close_cb=self._close_ws,
|
|
max_session_duration=300,
|
|
mark_refreshed_on_get=True,
|
|
)
|
|
# min_sentence_len=1 emits each sentence as soon as the next one starts,
|
|
# rather than batching short sentences together — minimizes TTFB on the
|
|
# first sentence and keeps Fish synthesizing continuously.
|
|
self._sentence_tokenizer = (
|
|
tokenizer
|
|
if is_given(tokenizer)
|
|
else tokenize.blingfire.SentenceTokenizer(min_sentence_len=1)
|
|
)
|
|
self._streams = weakref.WeakSet[SynthesizeStream]()
|
|
|
|
@property
|
|
def model(self) -> TTSModels | str:
|
|
return self._opts.model
|
|
|
|
@property
|
|
def provider(self) -> str:
|
|
return "FishAudio"
|
|
|
|
@property
|
|
def output_format(self) -> OutputFormat:
|
|
return self._opts.output_format
|
|
|
|
@property
|
|
def voice_id(self) -> NotGivenOr[str]:
|
|
return self._opts.voice_id
|
|
|
|
@property
|
|
def latency_mode(self) -> LatencyMode:
|
|
return self._opts.latency_mode
|
|
|
|
def _ensure_session(self) -> aiohttp.ClientSession:
|
|
if not self._session:
|
|
self._session = utils.http_context.http_session()
|
|
return self._session
|
|
|
|
async def _connect_ws(self, timeout: float) -> aiohttp.ClientWebSocketResponse:
|
|
session = self._ensure_session()
|
|
return await asyncio.wait_for(
|
|
session.ws_connect(
|
|
self._opts.get_ws_url("/v1/tts/live"),
|
|
headers={
|
|
"Authorization": f"Bearer {self._opts.api_key}",
|
|
"User-Agent": USER_AGENT,
|
|
"model": self._opts.model,
|
|
},
|
|
heartbeat=30.0,
|
|
),
|
|
timeout,
|
|
)
|
|
|
|
async def _close_ws(self, ws: aiohttp.ClientWebSocketResponse) -> None:
|
|
await ws.close()
|
|
|
|
def prewarm(self) -> None:
|
|
self._pool.prewarm()
|
|
|
|
def update_options(
|
|
self,
|
|
*,
|
|
model: NotGivenOr[TTSModels | str] = NOT_GIVEN,
|
|
voice_id: NotGivenOr[str] = NOT_GIVEN,
|
|
latency_mode: NotGivenOr[LatencyMode] = NOT_GIVEN,
|
|
chunk_length: NotGivenOr[int] = NOT_GIVEN,
|
|
speed: NotGivenOr[float] = NOT_GIVEN,
|
|
volume: NotGivenOr[float] = NOT_GIVEN,
|
|
) -> None:
|
|
if is_given(model) and model != self._opts.model:
|
|
self._opts.model = model
|
|
# The model is sent as a connection header at ws-handshake time, not in the
|
|
# per-request body, so a pooled socket keeps the old model. Drop pooled
|
|
# connections so the next stream reconnects with the new model. Other
|
|
# options ride in the per-request body and need no reconnect.
|
|
self._pool.invalidate()
|
|
if is_given(voice_id):
|
|
self._opts.voice_id = voice_id
|
|
if is_given(latency_mode):
|
|
self._opts.latency_mode = latency_mode
|
|
if is_given(chunk_length):
|
|
if not 100 <= chunk_length <= 300:
|
|
raise ValueError("chunk_length must be between 100 and 300")
|
|
self._opts.chunk_length = chunk_length
|
|
if is_given(speed):
|
|
self._opts.speed = speed
|
|
if is_given(volume):
|
|
self._opts.volume = volume
|
|
|
|
def synthesize(
|
|
self,
|
|
text: str,
|
|
*,
|
|
conn_options: APIConnectOptions = DEFAULT_API_CONNECT_OPTIONS,
|
|
) -> 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()
|
|
await self._pool.aclose()
|
|
|
|
|
|
def _build_tts_request(opts: _TTSOptions, *, text: str = "") -> dict[str, Any]:
|
|
# Send the same field set the upstream Fish Audio Python SDK sends so the
|
|
# server doesn't fall back to its own (larger) defaults — in particular the
|
|
# docs default of `chunk_length=300` produces large bursts that leave audible
|
|
# gaps between Fish's chunk boundaries.
|
|
# `prosody` stays None unless the caller set speed/volume, so the default
|
|
# request is byte-for-byte unchanged.
|
|
prosody: dict[str, float] | None = None
|
|
if is_given(opts.speed) or is_given(opts.volume):
|
|
prosody = {}
|
|
if is_given(opts.speed):
|
|
prosody["speed"] = opts.speed
|
|
if is_given(opts.volume):
|
|
prosody["volume"] = opts.volume
|
|
return {
|
|
"text": text,
|
|
"chunk_length": opts.chunk_length,
|
|
"format": opts.output_format,
|
|
"sample_rate": opts.sample_rate,
|
|
"mp3_bitrate": 64,
|
|
"opus_bitrate": 64000,
|
|
"references": [],
|
|
# Fish Audio's wire field is `reference_id`; we expose it as `voice_id` on
|
|
# the plugin for consistency with other TTS plugins.
|
|
"reference_id": opts.voice_id if is_given(opts.voice_id) else None,
|
|
"normalize": True,
|
|
"latency": opts.latency_mode,
|
|
"prosody": prosody,
|
|
"top_p": 0.7,
|
|
"temperature": 0.7,
|
|
}
|
|
|
|
|
|
class ChunkedStream(tts.ChunkedStream):
|
|
"""Synthesize via the Fish Audio HTTP /v1/tts endpoint."""
|
|
|
|
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:
|
|
payload = _build_tts_request(self._opts, text=self._input_text)
|
|
|
|
try:
|
|
async with self._tts._ensure_session().post(
|
|
self._opts.get_http_url("/v1/tts"),
|
|
headers={
|
|
"Authorization": f"Bearer {self._opts.api_key}",
|
|
"Content-Type": "application/msgpack",
|
|
"model": self._opts.model,
|
|
},
|
|
data=msgpack.packb(payload, use_bin_type=True),
|
|
timeout=aiohttp.ClientTimeout(total=30, sock_connect=self._conn_options.timeout),
|
|
) as resp:
|
|
resp.raise_for_status()
|
|
|
|
output_emitter.initialize(
|
|
request_id=utils.shortuuid(),
|
|
sample_rate=self._opts.sample_rate,
|
|
num_channels=NUM_CHANNELS,
|
|
mime_type=f"audio/{self._opts.output_format}",
|
|
)
|
|
|
|
async for data, _ in resp.content.iter_chunks():
|
|
output_emitter.push(data)
|
|
|
|
output_emitter.flush()
|
|
|
|
except asyncio.TimeoutError:
|
|
raise APITimeoutError() from None
|
|
except aiohttp.ClientResponseError as e:
|
|
raise APIStatusError(
|
|
message=e.message, status_code=e.status, request_id=None, body=None
|
|
) from None
|
|
except Exception as e:
|
|
raise APIConnectionError() from e
|
|
|
|
|
|
class SynthesizeStream(tts.SynthesizeStream):
|
|
"""Streaming TTS via the Fish Audio /v1/tts/live WebSocket endpoint."""
|
|
|
|
def __init__(self, *, tts: TTS, conn_options: APIConnectOptions) -> None:
|
|
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()
|
|
output_emitter.initialize(
|
|
request_id=request_id,
|
|
sample_rate=self._opts.sample_rate,
|
|
num_channels=NUM_CHANNELS,
|
|
mime_type=f"audio/{self._opts.output_format}",
|
|
stream=True,
|
|
)
|
|
output_emitter.start_segment(segment_id=request_id)
|
|
|
|
try:
|
|
async with self._tts._pool.connection(timeout=self._conn_options.timeout) as ws:
|
|
await self._run_ws(ws, output_emitter)
|
|
|
|
except asyncio.TimeoutError:
|
|
raise APITimeoutError() from None
|
|
except aiohttp.ClientResponseError as e:
|
|
raise APIStatusError(
|
|
message=e.message, status_code=e.status, request_id=request_id, body=None
|
|
) from None
|
|
except APIStatusError:
|
|
raise
|
|
except Exception as e:
|
|
raise APIConnectionError() from e
|
|
finally:
|
|
output_emitter.end_segment()
|
|
|
|
async def _run_ws(
|
|
self,
|
|
ws: aiohttp.ClientWebSocketResponse,
|
|
output_emitter: tts.AudioEmitter,
|
|
) -> None:
|
|
# Tokenize incoming text by sentence and flush after each sentence so
|
|
# Fish synthesizes immediately at sentence boundaries instead of waiting
|
|
# for `chunk_length` characters to accumulate. The result is much smoother
|
|
# audio: gaps line up with sentence breaks (where pauses are natural)
|
|
# rather than mid-clause.
|
|
sent_stream = self._tts._sentence_tokenizer.stream()
|
|
|
|
async def input_task() -> None:
|
|
try:
|
|
first_token = True
|
|
async for data in self._input_ch:
|
|
if isinstance(data, self._FlushSentinel):
|
|
sent_stream.flush()
|
|
continue
|
|
if not data:
|
|
continue
|
|
if first_token:
|
|
self._mark_started()
|
|
first_token = False
|
|
sent_stream.push_text(data)
|
|
finally:
|
|
sent_stream.end_input()
|
|
|
|
async def send_task() -> None:
|
|
start_msg = {"event": "start", "request": _build_tts_request(self._opts)}
|
|
await ws.send_bytes(msgpack.packb(start_msg, use_bin_type=True))
|
|
|
|
async for ev in sent_stream:
|
|
sentence = ev.token
|
|
if not sentence:
|
|
continue
|
|
await ws.send_bytes(
|
|
msgpack.packb({"event": "text", "text": sentence + " "}, use_bin_type=True)
|
|
)
|
|
await ws.send_bytes(msgpack.packb({"event": "flush"}, use_bin_type=True))
|
|
|
|
await ws.send_bytes(msgpack.packb({"event": "stop"}, use_bin_type=True))
|
|
|
|
prebuffer_chunks = _PREBUFFER_CHUNKS
|
|
prebuffer = bytearray()
|
|
chunks_seen = 0
|
|
prebuffering = prebuffer_chunks > 1
|
|
|
|
def push_audio(audio: bytes) -> None:
|
|
nonlocal prebuffering, chunks_seen
|
|
if prebuffering:
|
|
prebuffer.extend(audio)
|
|
chunks_seen += 1
|
|
if chunks_seen < prebuffer_chunks:
|
|
return
|
|
output_emitter.push(bytes(prebuffer))
|
|
prebuffer.clear()
|
|
prebuffering = False
|
|
else:
|
|
output_emitter.push(audio)
|
|
|
|
def flush_prebuffer() -> None:
|
|
# Stream ended before we reached `prebuffer_chunks` (short utterance) —
|
|
# release whatever is held so nothing is dropped or left hanging.
|
|
nonlocal prebuffering
|
|
if prebuffering and prebuffer:
|
|
output_emitter.push(bytes(prebuffer))
|
|
prebuffer.clear()
|
|
prebuffering = False
|
|
|
|
async def recv_task() -> None:
|
|
# No per-receive timeout: Fish has natural inter-sentence gaps that
|
|
# can exceed `_conn_options.timeout` when the LLM is slow. Dead
|
|
# connections are detected by aiohttp's ws heartbeat (see ws_connect).
|
|
while True:
|
|
msg = await ws.receive()
|
|
if msg.type in (
|
|
aiohttp.WSMsgType.CLOSED,
|
|
aiohttp.WSMsgType.CLOSE,
|
|
aiohttp.WSMsgType.CLOSING,
|
|
):
|
|
raise APIStatusError(
|
|
"Fish Audio websocket connection closed unexpectedly",
|
|
status_code=ws.close_code or -1,
|
|
request_id=None,
|
|
body=f"{msg.data=} {msg.extra=}",
|
|
)
|
|
|
|
if msg.type != aiohttp.WSMsgType.BINARY:
|
|
logger.debug("unexpected Fish Audio message type %s", msg.type)
|
|
continue
|
|
|
|
data = msgpack.unpackb(msg.data, raw=False)
|
|
event = data.get("event")
|
|
if event == "audio":
|
|
audio = data.get("audio")
|
|
if audio:
|
|
push_audio(audio)
|
|
elif event == "finish":
|
|
reason = data.get("reason")
|
|
if reason == "error":
|
|
raise APIStatusError(
|
|
"Fish Audio TTS reported an error",
|
|
status_code=-1,
|
|
request_id=None,
|
|
body=str(data),
|
|
)
|
|
flush_prebuffer()
|
|
break
|
|
else:
|
|
logger.debug("unknown Fish Audio event: %s", data)
|
|
|
|
tasks = [
|
|
asyncio.create_task(input_task()),
|
|
asyncio.create_task(send_task()),
|
|
asyncio.create_task(recv_task()),
|
|
]
|
|
try:
|
|
await asyncio.gather(*tasks)
|
|
finally:
|
|
await sent_stream.aclose()
|
|
await utils.aio.gracefully_cancel(*tasks)
|