wehub-resource-sync
608 lines
19 KiB
Python
608 lines
19 KiB
Python
import asyncio
|
|
import contextlib
|
|
import gc
|
|
|
|
import pytest
|
|
from google.cloud.speech_v1.types import cloud_speech as cloud_speech_v1
|
|
from google.cloud.speech_v2.types import cloud_speech as cloud_speech_v2
|
|
from google.protobuf.duration_pb2 import Duration
|
|
|
|
from livekit.agents import APIConnectOptions, LanguageCode
|
|
from livekit.agents.stt import SpeechData, SpeechEvent, SpeechEventType
|
|
from livekit.agents.types import TimedString
|
|
from livekit.agents.utils.aio import ChanClosed
|
|
from livekit.plugins.google.stt import (
|
|
SpeechStream,
|
|
STTOptions,
|
|
_recognize_response_to_speech_event, # pyright: ignore[reportPrivateUsage]
|
|
_streaming_recognize_response_to_speech_data, # pyright: ignore[reportPrivateUsage]
|
|
)
|
|
|
|
pytestmark = pytest.mark.plugin("google")
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_google_adc(monkeypatch):
|
|
from livekit.plugins.google import stt as google_stt
|
|
|
|
monkeypatch.setattr(google_stt, "gauth_default", lambda: (None, "test-project"))
|
|
|
|
|
|
class _FakeSTT:
|
|
_label = "google.STT"
|
|
model = "default"
|
|
provider = "Google Cloud Platform"
|
|
|
|
def emit(self, *_args, **_kwargs):
|
|
pass
|
|
|
|
|
|
class _FakeStreamingCall:
|
|
def __init__(self, requests):
|
|
self._requests = requests
|
|
self._consumer_task = asyncio.create_task(self._consume_requests())
|
|
self.awaiting_audio = asyncio.Event()
|
|
self.consumer_done = asyncio.Event()
|
|
|
|
async def _consume_requests(self):
|
|
try:
|
|
await self._requests.__anext__()
|
|
self.awaiting_audio.set()
|
|
await self._requests.__anext__()
|
|
finally:
|
|
self.consumer_done.set()
|
|
|
|
def cancel(self):
|
|
self._consumer_task.cancel()
|
|
|
|
def __aiter__(self):
|
|
return self
|
|
|
|
async def __anext__(self):
|
|
await self.awaiting_audio.wait()
|
|
raise StopAsyncIteration
|
|
|
|
|
|
class _FakeSpeechClient:
|
|
def __init__(self):
|
|
self.call: _FakeStreamingCall | None = None
|
|
self.call_ready = asyncio.Event()
|
|
|
|
async def streaming_recognize(self, *, requests):
|
|
self.call = _FakeStreamingCall(requests)
|
|
self.call_ready.set()
|
|
return self.call
|
|
|
|
|
|
class _FakeConnectionPool:
|
|
last_acquire_time = 0.0
|
|
last_connection_reused = False
|
|
|
|
def __init__(self, client):
|
|
self._client = client
|
|
|
|
@contextlib.asynccontextmanager
|
|
async def connection(self, *, timeout):
|
|
yield self._client
|
|
|
|
def remove(self, _client):
|
|
pass
|
|
|
|
|
|
def _default_stt_options() -> STTOptions:
|
|
return STTOptions(
|
|
languages=[LanguageCode("en-US")],
|
|
detect_language=True,
|
|
interim_results=True,
|
|
punctuate=True,
|
|
spoken_punctuation=False,
|
|
enable_word_time_offsets=True,
|
|
enable_word_confidence=False,
|
|
enable_voice_activity_events=False,
|
|
model="default",
|
|
sample_rate=16000,
|
|
min_confidence_threshold=0.65,
|
|
profanity_filter=False,
|
|
)
|
|
|
|
|
|
async def test_google_stt_stream_cancel_does_not_leak_chanclosed_task():
|
|
client = _FakeSpeechClient()
|
|
stream = SpeechStream(
|
|
stt=_FakeSTT(),
|
|
conn_options=APIConnectOptions(max_retry=0, timeout=0.1),
|
|
pool=_FakeConnectionPool(client),
|
|
recognizer_cb=lambda _client: "",
|
|
config=_default_stt_options(),
|
|
)
|
|
|
|
loop = asyncio.get_running_loop()
|
|
original_exception_handler = loop.get_exception_handler()
|
|
contexts = []
|
|
loop.set_exception_handler(lambda _loop, context: contexts.append(context))
|
|
|
|
try:
|
|
await asyncio.wait_for(client.call_ready.wait(), timeout=1)
|
|
assert client.call is not None
|
|
await asyncio.wait_for(client.call.awaiting_audio.wait(), timeout=1)
|
|
|
|
await asyncio.wait_for(stream._task, timeout=1)
|
|
await asyncio.wait_for(client.call.consumer_done.wait(), timeout=1)
|
|
|
|
await stream.aclose()
|
|
for _ in range(3):
|
|
gc.collect()
|
|
await asyncio.sleep(0)
|
|
finally:
|
|
loop.set_exception_handler(original_exception_handler)
|
|
|
|
unhandled_chanclosed = [
|
|
context
|
|
for context in contexts
|
|
if context.get("message") == "Task exception was never retrieved"
|
|
and isinstance(context.get("exception"), ChanClosed)
|
|
]
|
|
assert unhandled_chanclosed == []
|
|
|
|
|
|
async def test_google_stt_stream_cancel_waits_for_cleanup_if_cancelled_twice(monkeypatch):
|
|
from livekit.plugins.google import stt as google_stt
|
|
|
|
client = _FakeSpeechClient()
|
|
stream = SpeechStream(
|
|
stt=_FakeSTT(),
|
|
conn_options=APIConnectOptions(max_retry=0, timeout=0.1),
|
|
pool=_FakeConnectionPool(client),
|
|
recognizer_cb=lambda _client: "",
|
|
config=_default_stt_options(),
|
|
)
|
|
cleanup_started = asyncio.Event()
|
|
cleanup_release = asyncio.Event()
|
|
cleanup_done = asyncio.Event()
|
|
original_gracefully_cancel = google_stt.utils.aio.gracefully_cancel
|
|
|
|
async def delayed_input_cleanup(*tasks):
|
|
if any(
|
|
isinstance(task, asyncio.Task) and task.get_coro().__qualname__ == "Chan.recv"
|
|
for task in tasks
|
|
):
|
|
cleanup_started.set()
|
|
try:
|
|
await cleanup_release.wait()
|
|
await original_gracefully_cancel(*tasks)
|
|
finally:
|
|
cleanup_done.set()
|
|
return
|
|
|
|
await original_gracefully_cancel(*tasks)
|
|
|
|
monkeypatch.setattr(google_stt.utils.aio, "gracefully_cancel", delayed_input_cleanup)
|
|
|
|
try:
|
|
await asyncio.wait_for(client.call_ready.wait(), timeout=1)
|
|
assert client.call is not None
|
|
await asyncio.wait_for(client.call.awaiting_audio.wait(), timeout=1)
|
|
await asyncio.wait_for(cleanup_started.wait(), timeout=1)
|
|
|
|
client.call.cancel()
|
|
await asyncio.sleep(0)
|
|
|
|
cleanup_release.set()
|
|
await asyncio.wait_for(cleanup_done.wait(), timeout=1)
|
|
await asyncio.wait_for(client.call.consumer_done.wait(), timeout=1)
|
|
finally:
|
|
cleanup_release.set()
|
|
await stream.aclose()
|
|
|
|
|
|
async def test_streaming_recognize_response_to_speech_data_01():
|
|
srr = cloud_speech_v2.StreamingRecognizeResponse(
|
|
results=[cloud_speech_v2.StreamingRecognitionResult()]
|
|
)
|
|
assert (
|
|
_streaming_recognize_response_to_speech_data(
|
|
srr, min_confidence_threshold=1.0, start_time_offset=0.0
|
|
)
|
|
is None
|
|
)
|
|
|
|
srr = cloud_speech_v1.StreamingRecognizeResponse(
|
|
results=[cloud_speech_v1.StreamingRecognitionResult()]
|
|
)
|
|
assert (
|
|
_streaming_recognize_response_to_speech_data(
|
|
srr, min_confidence_threshold=1.0, start_time_offset=0.0
|
|
)
|
|
is None
|
|
)
|
|
|
|
|
|
async def test_streaming_recognize_response_to_speech_data_02():
|
|
# final result should be returned regardless of confidence
|
|
srr = cloud_speech_v2.StreamingRecognizeResponse(
|
|
results=[
|
|
cloud_speech_v2.StreamingRecognitionResult(
|
|
alternatives=[
|
|
cloud_speech_v2.SpeechRecognitionAlternative(confidence=0.0, transcript="test")
|
|
],
|
|
is_final=True,
|
|
language_code="te-ST",
|
|
)
|
|
]
|
|
)
|
|
result = _streaming_recognize_response_to_speech_data(
|
|
srr, min_confidence_threshold=0.5, start_time_offset=0.0
|
|
)
|
|
assert type(result) is SpeechData
|
|
assert result.text == "test"
|
|
assert result.language == "te-ST"
|
|
assert result.confidence == 0.0
|
|
|
|
srr = cloud_speech_v1.StreamingRecognizeResponse(
|
|
results=[
|
|
cloud_speech_v1.StreamingRecognitionResult(
|
|
alternatives=[
|
|
cloud_speech_v1.SpeechRecognitionAlternative(confidence=0.0, transcript="test")
|
|
],
|
|
is_final=True,
|
|
language_code="te-ST",
|
|
)
|
|
]
|
|
)
|
|
result = _streaming_recognize_response_to_speech_data(
|
|
srr, min_confidence_threshold=0.5, start_time_offset=0.0
|
|
)
|
|
assert type(result) is SpeechData
|
|
assert result.text == "test"
|
|
assert result.language == "te-ST"
|
|
assert result.confidence == 0.0
|
|
|
|
|
|
async def test_streaming_recognize_response_to_speech_data_03():
|
|
srr = cloud_speech_v2.StreamingRecognizeResponse(
|
|
results=[
|
|
cloud_speech_v2.StreamingRecognitionResult(
|
|
alternatives=[
|
|
cloud_speech_v2.SpeechRecognitionAlternative(confidence=0.0, transcript="test")
|
|
],
|
|
is_final=False,
|
|
)
|
|
]
|
|
)
|
|
result = _streaming_recognize_response_to_speech_data(
|
|
srr, min_confidence_threshold=0.5, start_time_offset=0.0
|
|
)
|
|
assert result is None
|
|
|
|
srr = cloud_speech_v1.StreamingRecognizeResponse(
|
|
results=[
|
|
cloud_speech_v1.StreamingRecognitionResult(
|
|
alternatives=[
|
|
cloud_speech_v1.SpeechRecognitionAlternative(confidence=0.0, transcript="test")
|
|
],
|
|
is_final=False,
|
|
)
|
|
]
|
|
)
|
|
result = _streaming_recognize_response_to_speech_data(
|
|
srr, min_confidence_threshold=0.5, start_time_offset=0.0
|
|
)
|
|
assert result is None
|
|
|
|
|
|
async def test_streaming_recognize_response_to_speech_data_04():
|
|
srr = cloud_speech_v2.StreamingRecognizeResponse(
|
|
results=[
|
|
cloud_speech_v2.StreamingRecognitionResult(
|
|
alternatives=[
|
|
cloud_speech_v2.SpeechRecognitionAlternative(
|
|
confidence=1.0, transcript="test01"
|
|
)
|
|
],
|
|
is_final=False,
|
|
language_code="te-ST",
|
|
),
|
|
cloud_speech_v2.StreamingRecognitionResult(
|
|
alternatives=[
|
|
cloud_speech_v2.SpeechRecognitionAlternative(
|
|
confidence=1.0, transcript="test02"
|
|
)
|
|
],
|
|
is_final=False,
|
|
language_code="te-ST",
|
|
),
|
|
]
|
|
)
|
|
result = _streaming_recognize_response_to_speech_data(
|
|
srr, min_confidence_threshold=0.5, start_time_offset=0.0
|
|
)
|
|
assert type(result) is SpeechData
|
|
assert result.text == "test01test02"
|
|
assert result.language == "te-ST"
|
|
assert result.confidence == 1.0
|
|
|
|
|
|
async def test_streaming_recognize_response_to_speech_data_05():
|
|
srr = cloud_speech_v2.StreamingRecognizeResponse(
|
|
results=[
|
|
cloud_speech_v2.StreamingRecognitionResult(
|
|
alternatives=[
|
|
cloud_speech_v2.SpeechRecognitionAlternative(
|
|
confidence=1.0, transcript="test01"
|
|
)
|
|
],
|
|
is_final=False,
|
|
language_code="te-ST",
|
|
),
|
|
cloud_speech_v2.StreamingRecognitionResult(
|
|
alternatives=[
|
|
cloud_speech_v2.SpeechRecognitionAlternative(
|
|
confidence=1.0, transcript="test02"
|
|
)
|
|
],
|
|
is_final=False,
|
|
language_code="te-ST",
|
|
),
|
|
cloud_speech_v2.StreamingRecognitionResult(
|
|
alternatives=[
|
|
cloud_speech_v2.SpeechRecognitionAlternative(confidence=1.0, transcript="best")
|
|
],
|
|
is_final=True,
|
|
language_code="te-ST",
|
|
),
|
|
]
|
|
)
|
|
result = _streaming_recognize_response_to_speech_data(
|
|
srr, min_confidence_threshold=0.5, start_time_offset=0.0
|
|
)
|
|
assert type(result) is SpeechData
|
|
assert result.text == "best"
|
|
assert result.language == "te-ST"
|
|
assert result.confidence == 1.0
|
|
|
|
|
|
async def test_streaming_recognize_response_to_speech_data_words():
|
|
srr = cloud_speech_v2.StreamingRecognizeResponse(
|
|
results=[
|
|
cloud_speech_v2.StreamingRecognitionResult(
|
|
alternatives=[
|
|
cloud_speech_v2.SpeechRecognitionAlternative(
|
|
confidence=1.0,
|
|
transcript="test",
|
|
words=[
|
|
cloud_speech_v2.WordInfo(
|
|
word="test",
|
|
start_offset=Duration(seconds=0),
|
|
end_offset=Duration(seconds=1),
|
|
confidence=1.0,
|
|
)
|
|
],
|
|
)
|
|
],
|
|
is_final=True,
|
|
)
|
|
]
|
|
)
|
|
result = _streaming_recognize_response_to_speech_data(
|
|
srr, min_confidence_threshold=0.5, start_time_offset=0.0
|
|
)
|
|
assert type(result) is SpeechData
|
|
assert result.text == "test"
|
|
assert result.confidence == 1.0
|
|
assert result.words == [
|
|
TimedString(
|
|
text="test", start_time=0.0, end_time=1.0, start_time_offset=0.0, confidence=1.0
|
|
)
|
|
]
|
|
|
|
srr = cloud_speech_v1.StreamingRecognizeResponse(
|
|
results=[
|
|
cloud_speech_v1.StreamingRecognitionResult(
|
|
alternatives=[
|
|
cloud_speech_v1.SpeechRecognitionAlternative(
|
|
confidence=1.0,
|
|
transcript="test",
|
|
words=[
|
|
cloud_speech_v1.WordInfo(
|
|
word="test",
|
|
start_time=Duration(seconds=0),
|
|
end_time=Duration(seconds=1),
|
|
confidence=1.0,
|
|
)
|
|
],
|
|
)
|
|
],
|
|
is_final=True,
|
|
)
|
|
]
|
|
)
|
|
result = _streaming_recognize_response_to_speech_data(
|
|
srr, min_confidence_threshold=0.5, start_time_offset=0.1
|
|
)
|
|
assert type(result) is SpeechData
|
|
assert result.text == "test"
|
|
assert result.confidence == 1.0
|
|
assert result.words == [
|
|
TimedString(
|
|
text="test", start_time=0.1, end_time=1.1, start_time_offset=0.1, confidence=1.0
|
|
)
|
|
]
|
|
|
|
|
|
async def test_recognize_response_to_speech_event_words():
|
|
resp = cloud_speech_v2.RecognizeResponse(
|
|
results=[
|
|
cloud_speech_v2.SpeechRecognitionResult(
|
|
alternatives=[
|
|
cloud_speech_v2.SpeechRecognitionAlternative(
|
|
confidence=1.0,
|
|
transcript="test",
|
|
words=[
|
|
cloud_speech_v2.WordInfo(
|
|
word="test",
|
|
start_offset=Duration(seconds=0),
|
|
end_offset=Duration(seconds=1),
|
|
confidence=1.0,
|
|
)
|
|
],
|
|
)
|
|
],
|
|
language_code="te-ST",
|
|
)
|
|
]
|
|
)
|
|
result = _recognize_response_to_speech_event(resp)
|
|
assert type(result) is SpeechEvent
|
|
assert result.type == SpeechEventType.FINAL_TRANSCRIPT
|
|
assert result.alternatives == [
|
|
SpeechData(
|
|
language=LanguageCode("te-ST"),
|
|
start_time=0.0,
|
|
end_time=1.0,
|
|
text="test",
|
|
confidence=1.0,
|
|
words=[
|
|
TimedString(
|
|
text="test", start_time=0.0, end_time=1.0, start_time_offset=0.0, confidence=1.0
|
|
)
|
|
],
|
|
)
|
|
]
|
|
|
|
resp = cloud_speech_v1.RecognizeResponse(
|
|
results=[
|
|
cloud_speech_v1.SpeechRecognitionResult(
|
|
alternatives=[
|
|
cloud_speech_v1.SpeechRecognitionAlternative(
|
|
confidence=1.0,
|
|
transcript="test",
|
|
words=[
|
|
cloud_speech_v1.WordInfo(
|
|
word="test",
|
|
start_time=Duration(seconds=0),
|
|
end_time=Duration(seconds=1),
|
|
confidence=1.0,
|
|
)
|
|
],
|
|
)
|
|
],
|
|
language_code="te-ST",
|
|
)
|
|
]
|
|
)
|
|
result = _recognize_response_to_speech_event(resp)
|
|
assert type(result) is SpeechEvent
|
|
assert result.type == SpeechEventType.FINAL_TRANSCRIPT
|
|
assert result.alternatives == [
|
|
SpeechData(
|
|
language=LanguageCode("te-ST"),
|
|
start_time=0.0,
|
|
end_time=1.0,
|
|
text="test",
|
|
confidence=1.0,
|
|
words=[
|
|
TimedString(
|
|
text="test", start_time=0.0, end_time=1.0, start_time_offset=0.0, confidence=1.0
|
|
)
|
|
],
|
|
)
|
|
]
|
|
|
|
|
|
async def test_voice_activity_timeout_defaults(mock_google_adc):
|
|
"""Test voice activity timeouts are not set by default."""
|
|
from livekit.agents.types import NOT_GIVEN
|
|
from livekit.plugins.google import STT
|
|
|
|
stt = STT()
|
|
assert stt._config.speech_start_timeout is NOT_GIVEN
|
|
assert stt._config.speech_end_timeout is NOT_GIVEN
|
|
|
|
|
|
async def test_voice_activity_timeout_set(mock_google_adc):
|
|
"""Test voice activity timeouts can be set."""
|
|
from livekit.plugins.google import STT
|
|
|
|
stt = STT(
|
|
speech_start_timeout=10.0,
|
|
speech_end_timeout=2.5,
|
|
)
|
|
assert stt._config.speech_start_timeout == 10.0
|
|
assert stt._config.speech_end_timeout == 2.5
|
|
|
|
|
|
async def test_voice_activity_timeout_fractional_seconds(mock_google_adc):
|
|
"""Test voice activity timeouts handle fractional seconds."""
|
|
from livekit.plugins.google import STT
|
|
|
|
stt = STT(
|
|
speech_start_timeout=5.5,
|
|
speech_end_timeout=1.25,
|
|
)
|
|
assert stt._config.speech_start_timeout == 5.5
|
|
assert stt._config.speech_end_timeout == 1.25
|
|
|
|
|
|
async def test_voice_activity_timeout_speech_start_only(mock_google_adc):
|
|
"""Test setting only speech_start_timeout."""
|
|
from livekit.agents.types import NOT_GIVEN
|
|
from livekit.plugins.google import STT
|
|
|
|
stt = STT(speech_start_timeout=15.0)
|
|
assert stt._config.speech_start_timeout == 15.0
|
|
assert stt._config.speech_end_timeout is NOT_GIVEN
|
|
|
|
|
|
async def test_voice_activity_timeout_speech_end_only(mock_google_adc):
|
|
"""Test setting only speech_end_timeout."""
|
|
from livekit.agents.types import NOT_GIVEN
|
|
from livekit.plugins.google import STT
|
|
|
|
stt = STT(speech_end_timeout=3.0)
|
|
assert stt._config.speech_end_timeout == 3.0
|
|
assert stt._config.speech_start_timeout is NOT_GIVEN
|
|
|
|
|
|
async def test_voice_activity_timeout_v2_model(mock_google_adc):
|
|
"""Test that V2 model detection works correctly."""
|
|
from livekit.plugins.google import STT
|
|
|
|
stt_v2 = STT(model="chirp_3")
|
|
assert stt_v2._config.version == 2
|
|
|
|
stt_v1 = STT(model="default")
|
|
assert stt_v1._config.version == 1
|
|
|
|
|
|
async def test_voice_activity_timeout_update(mock_google_adc):
|
|
"""Test that timeout options can be updated dynamically."""
|
|
from livekit.plugins.google import STT
|
|
|
|
stt = STT(
|
|
speech_start_timeout=10.0,
|
|
speech_end_timeout=2.0,
|
|
)
|
|
stt.update_options(
|
|
speech_start_timeout=15.0,
|
|
speech_end_timeout=3.0,
|
|
)
|
|
assert stt._config.speech_start_timeout == 15.0
|
|
assert stt._config.speech_end_timeout == 3.0
|
|
|
|
|
|
async def test_voice_activity_timeout_partial_update(mock_google_adc):
|
|
"""Test updating only one timeout at a time."""
|
|
from livekit.plugins.google import STT
|
|
|
|
stt = STT(
|
|
speech_start_timeout=10.0,
|
|
speech_end_timeout=2.0,
|
|
)
|
|
stt.update_options(speech_start_timeout=20.0)
|
|
assert stt._config.speech_start_timeout == 20.0
|
|
assert stt._config.speech_end_timeout == 2.0
|
|
|
|
stt.update_options(speech_end_timeout=5.0)
|
|
assert stt._config.speech_start_timeout == 20.0
|
|
assert stt._config.speech_end_timeout == 5.0
|