wehub-resource-sync
330 lines
11 KiB
Python
330 lines
11 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import contextlib
|
|
import time
|
|
|
|
import pytest
|
|
|
|
from livekit import rtc
|
|
from livekit.agents import APIConnectionError, APIConnectOptions, APIError, utils
|
|
from livekit.agents.tts import TTS, AvailabilityChangedEvent, FallbackAdapter
|
|
from livekit.agents.tts.tts import SynthesizedAudio, SynthesizeStream
|
|
from livekit.agents.types import USERDATA_TTS_STARTED_TIME
|
|
from livekit.agents.utils.aio.channel import ChanEmpty
|
|
|
|
from .fake_tts import FakeTTS
|
|
|
|
pytestmark = [pytest.mark.unit, pytest.mark.virtual_time, pytest.mark.no_concurrent]
|
|
|
|
|
|
class FallbackAdapterTester(FallbackAdapter):
|
|
def __init__(
|
|
self,
|
|
tts: list[TTS],
|
|
*,
|
|
max_retry_per_tts: int = 1, # only retry once by default
|
|
sample_rate: int | None = None,
|
|
) -> None:
|
|
super().__init__(
|
|
tts,
|
|
max_retry_per_tts=max_retry_per_tts,
|
|
sample_rate=sample_rate,
|
|
)
|
|
|
|
self.on("tts_availability_changed", self._on_tts_availability_changed)
|
|
|
|
self._availability_changed_ch: dict[int, utils.aio.Chan[AvailabilityChangedEvent]] = {
|
|
id(t): utils.aio.Chan[AvailabilityChangedEvent]() for t in tts
|
|
}
|
|
|
|
def _on_tts_availability_changed(self, ev: AvailabilityChangedEvent) -> None:
|
|
self._availability_changed_ch[id(ev.tts)].send_nowait(ev)
|
|
|
|
def availability_changed_ch(
|
|
self,
|
|
tts: TTS,
|
|
) -> utils.aio.ChanReceiver[AvailabilityChangedEvent]:
|
|
return self._availability_changed_ch[id(tts)]
|
|
|
|
|
|
async def test_tts_fallback() -> None:
|
|
fake1 = FakeTTS(fake_exception=APIConnectionError("fake1 failed"))
|
|
fake2 = FakeTTS(fake_audio_duration=5.0, sample_rate=48000)
|
|
|
|
fallback_adapter = FallbackAdapterTester([fake1, fake2])
|
|
|
|
async with fallback_adapter.synthesize("hello test") as stream:
|
|
frames = []
|
|
async for data in stream:
|
|
frames.append(data.frame)
|
|
|
|
assert fake1.synthesize_ch.recv_nowait()
|
|
assert fake2.synthesize_ch.recv_nowait()
|
|
|
|
assert rtc.combine_audio_frames(frames).duration == 5.01
|
|
|
|
assert not fallback_adapter.availability_changed_ch(fake1).recv_nowait().available
|
|
|
|
fake2.update_options(fake_audio_duration=0.0)
|
|
|
|
with pytest.raises(APIConnectionError):
|
|
async with fallback_adapter.synthesize("hello test") as stream:
|
|
async for _ in stream:
|
|
pass
|
|
|
|
assert not fallback_adapter.availability_changed_ch(fake2).recv_nowait().available
|
|
|
|
await fallback_adapter.aclose()
|
|
|
|
|
|
async def test_no_audio() -> None:
|
|
fake1 = FakeTTS(fake_audio_duration=0.0)
|
|
|
|
fallback_adapter = FallbackAdapterTester([fake1])
|
|
|
|
with pytest.raises(APIConnectionError):
|
|
async with fallback_adapter.synthesize("hello test chunked") as stream:
|
|
async for _ in stream:
|
|
pass
|
|
|
|
# stream
|
|
fake1.update_options(fake_audio_duration=5.0)
|
|
|
|
async def _input_task(stream: SynthesizeStream, text: str):
|
|
with contextlib.suppress(RuntimeError):
|
|
stream.push_text(text)
|
|
stream.flush()
|
|
stream.end_input()
|
|
|
|
async def _stream_task(text: str, audio_duration: float) -> None:
|
|
fake1.update_options(fake_timeout=0.5, fake_audio_duration=audio_duration)
|
|
|
|
async with fallback_adapter.stream() as stream:
|
|
input_task = asyncio.create_task(_input_task(stream, text))
|
|
|
|
segments = set()
|
|
try:
|
|
async for ev in stream:
|
|
segments.add(ev.segment_id)
|
|
finally:
|
|
await input_task
|
|
|
|
if audio_duration > 0.0:
|
|
assert len(segments) == 1
|
|
else:
|
|
assert len(segments) == 0
|
|
|
|
await _stream_task("hello test normal", 1.0)
|
|
|
|
with pytest.raises(APIError):
|
|
await _stream_task("hello test no audio", 0.0)
|
|
|
|
await fallback_adapter.aclose()
|
|
|
|
|
|
async def test_tts_stream_fallback() -> None:
|
|
fake1 = FakeTTS(fake_exception=APIConnectionError("fake1 failed"))
|
|
fake2 = FakeTTS(fake_audio_duration=5.0)
|
|
|
|
fallback_adapter = FallbackAdapterTester([fake1, fake2])
|
|
|
|
async with fallback_adapter.stream() as stream:
|
|
stream.push_text("hello test")
|
|
stream.end_input()
|
|
|
|
async for _ in stream:
|
|
pass
|
|
|
|
assert fake1.stream_ch.recv_nowait()
|
|
assert fake2.stream_ch.recv_nowait()
|
|
|
|
assert not fallback_adapter.availability_changed_ch(fake1).recv_nowait().available
|
|
|
|
await fallback_adapter.aclose()
|
|
|
|
|
|
async def test_tts_stream_ttfb_includes_failover_time() -> None:
|
|
# the fallback adapter is measured as a single TTS node: ttfb anchors on the first
|
|
# time a sentence was handed to any underlying TTS and is kept when falling back,
|
|
# so time spent failing over counts towards ttfb
|
|
fake1 = FakeTTS(fake_timeout=0.5, fake_exception=APIConnectionError("fake1 failed"))
|
|
fake2 = FakeTTS(fake_audio_duration=5.0)
|
|
|
|
fallback_adapter = FallbackAdapterTester([fake1, fake2])
|
|
|
|
async with fallback_adapter.stream(
|
|
conn_options=APIConnectOptions(timeout=10.0, max_retry=0, retry_interval=0.1)
|
|
) as stream:
|
|
start = time.perf_counter()
|
|
stream.push_text("hello test")
|
|
stream.end_input()
|
|
|
|
frames: list[SynthesizedAudio] = []
|
|
async for data in stream:
|
|
frames.append(data)
|
|
|
|
assert frames
|
|
started_times = {f.frame.userdata.get(USERDATA_TTS_STARTED_TIME) for f in frames}
|
|
assert len(started_times) == 1, "all frames should carry the same started time"
|
|
started_time = started_times.pop()
|
|
assert started_time is not None
|
|
# fake1 receives the sentence 0.5s in (after its fake connection delay) and fails
|
|
# without audio after two attempts (~1.1s); the anchor must stay on fake1's first
|
|
# attempt rather than moving to fake2's
|
|
assert started_time == pytest.approx(start + 0.5, abs=0.1)
|
|
|
|
await fallback_adapter.aclose()
|
|
|
|
|
|
async def test_tts_chunked_ttfb_includes_failover_time() -> None:
|
|
# same as above for the chunked path: the ChunkedStream base stamps its creation time
|
|
# (when the full text was submitted), which must not be overridden on failover
|
|
fake1 = FakeTTS(fake_timeout=10.0) # always times out
|
|
fake2 = FakeTTS(fake_audio_duration=5.0)
|
|
|
|
fallback_adapter = FallbackAdapterTester([fake1, fake2])
|
|
|
|
start = time.perf_counter()
|
|
async with fallback_adapter.synthesize(
|
|
"hello test",
|
|
conn_options=APIConnectOptions(timeout=0.5, max_retry=0, retry_interval=0.1),
|
|
) as stream:
|
|
frames: list[SynthesizedAudio] = []
|
|
async for data in stream:
|
|
frames.append(data)
|
|
|
|
assert frames
|
|
started_times = {f.frame.userdata.get(USERDATA_TTS_STARTED_TIME) for f in frames}
|
|
assert len(started_times) == 1, "all frames should carry the same started time"
|
|
started_time = started_times.pop()
|
|
assert started_time is not None
|
|
assert started_time == pytest.approx(start, abs=0.1)
|
|
|
|
await fallback_adapter.aclose()
|
|
|
|
|
|
async def test_tts_recover() -> None:
|
|
fake1 = FakeTTS(fake_exception=APIConnectionError("fake1 failed"))
|
|
fake2 = FakeTTS(fake_exception=APIConnectionError("fake2 failed"), fake_timeout=0.5)
|
|
|
|
fallback_adapter = FallbackAdapterTester([fake1, fake2])
|
|
|
|
with pytest.raises(APIConnectionError):
|
|
async with fallback_adapter.synthesize("hello test") as stream:
|
|
async for _ in stream:
|
|
pass
|
|
|
|
assert fake1.synthesize_ch.recv_nowait()
|
|
assert fake2.synthesize_ch.recv_nowait()
|
|
|
|
fake2.update_options(fake_exception=None, fake_audio_duration=5.0)
|
|
|
|
assert not fallback_adapter.availability_changed_ch(fake1).recv_nowait().available
|
|
assert not fallback_adapter.availability_changed_ch(fake2).recv_nowait().available
|
|
|
|
assert (
|
|
await asyncio.wait_for(fallback_adapter.availability_changed_ch(fake2).recv(), 1.0)
|
|
).available, "fake2 should have recovered"
|
|
|
|
async with fallback_adapter.synthesize("hello test") as stream:
|
|
async for _ in stream:
|
|
pass
|
|
|
|
assert fake1.synthesize_ch.recv_nowait()
|
|
assert fake2.synthesize_ch.recv_nowait()
|
|
|
|
with pytest.raises(ChanEmpty):
|
|
fallback_adapter.availability_changed_ch(fake1).recv_nowait()
|
|
|
|
with pytest.raises(ChanEmpty):
|
|
fallback_adapter.availability_changed_ch(fake2).recv_nowait()
|
|
|
|
await fallback_adapter.aclose()
|
|
|
|
|
|
async def test_audio_resampled() -> None:
|
|
fake1 = FakeTTS(sample_rate=48000, fake_exception=APIConnectionError("fake1 failed"))
|
|
fake2 = FakeTTS(fake_audio_duration=5.0, sample_rate=16000)
|
|
|
|
fallback_adapter = FallbackAdapterTester([fake1, fake2])
|
|
|
|
async with fallback_adapter.synthesize("hello test") as stream:
|
|
frames: list[SynthesizedAudio] = []
|
|
async for data in stream:
|
|
frames.append(data)
|
|
|
|
assert fake1.synthesize_ch.recv_nowait()
|
|
assert fake2.synthesize_ch.recv_nowait()
|
|
|
|
assert not fallback_adapter.availability_changed_ch(fake1).recv_nowait().available
|
|
|
|
assert frames[-1].is_final is True, "last frame should be final"
|
|
|
|
combined_frame = rtc.combine_audio_frames([f.frame for f in frames])
|
|
assert combined_frame.duration == 5.01
|
|
assert combined_frame.sample_rate == 48000
|
|
|
|
assert await asyncio.wait_for(fake1.synthesize_ch.recv(), 1.0)
|
|
|
|
async with fallback_adapter.stream() as stream:
|
|
stream.push_text("hello test")
|
|
stream.end_input()
|
|
|
|
frames: list[SynthesizedAudio] = []
|
|
async for data in stream:
|
|
frames.append(data)
|
|
|
|
assert fake2.stream_ch.recv_nowait()
|
|
|
|
assert frames[-1].is_final is True, "last frame should be final"
|
|
|
|
combined_frame = rtc.combine_audio_frames([f.frame for f in frames])
|
|
assert combined_frame.duration == 5.01 # 5.0 + 0.01 final flag
|
|
assert combined_frame.sample_rate == 48000
|
|
|
|
await fallback_adapter.aclose()
|
|
|
|
|
|
async def test_timeout():
|
|
fake1 = FakeTTS(fake_timeout=0.5, sample_rate=48000)
|
|
fake2 = FakeTTS(fake_timeout=0.5, sample_rate=48000)
|
|
|
|
fallback_adapter = FallbackAdapterTester([fake1, fake2])
|
|
|
|
with pytest.raises(APIConnectionError):
|
|
async with fallback_adapter.synthesize(
|
|
"hello test",
|
|
conn_options=APIConnectOptions(timeout=0.1, max_retry=0),
|
|
) as stream:
|
|
async for _ in stream:
|
|
pass
|
|
|
|
assert fake1.synthesize_ch.recv_nowait()
|
|
assert fake2.synthesize_ch.recv_nowait()
|
|
|
|
assert not fallback_adapter.availability_changed_ch(fake1).recv_nowait().available
|
|
assert not fallback_adapter.availability_changed_ch(fake2).recv_nowait().available
|
|
|
|
assert await asyncio.wait_for(fake1.synthesize_ch.recv(), 1.0)
|
|
assert await asyncio.wait_for(fake2.synthesize_ch.recv(), 1.0)
|
|
|
|
# stream
|
|
with pytest.raises(APIConnectionError):
|
|
async with fallback_adapter.stream(
|
|
conn_options=APIConnectOptions(timeout=0.1, max_retry=0)
|
|
) as stream:
|
|
stream.push_text("hello test")
|
|
stream.end_input()
|
|
|
|
async for _ in stream:
|
|
pass
|
|
|
|
assert fake1.stream_ch.recv_nowait()
|
|
assert fake2.stream_ch.recv_nowait()
|
|
|
|
assert await asyncio.wait_for(fake1.stream_ch.recv(), 1.0)
|
|
assert await asyncio.wait_for(fake2.stream_ch.recv(), 1.0)
|
|
|
|
await fallback_adapter.aclose()
|