# 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 concurrent.futures import contextlib import os from dataclasses import dataclass from typing import Any from livekit import rtc from livekit.agents import ( DEFAULT_API_CONNECT_OPTIONS, APIConnectOptions, LanguageCode, stt, utils, ) from livekit.agents.types import NOT_GIVEN, NotGivenOr from livekit.agents.utils import is_given from livekit.agents.voice.io import TimedString from .log import logger from .utils import DEFAULT_REGION try: from aws_sdk_transcribe_streaming.client import TranscribeStreamingClient from aws_sdk_transcribe_streaming.config import Config from aws_sdk_transcribe_streaming.models import ( AudioEvent, AudioStream, AudioStreamAudioEvent, BadRequestException, Result, StartStreamTranscriptionInput, TranscriptEvent, TranscriptResultStream, ) from smithy_aws_core.identity import ( AWSCredentialsIdentity, ContainerCredentialsResolver, EnvironmentCredentialsResolver, IMDSCredentialsResolver, StaticCredentialsResolver, ) from smithy_core.aio.identity import ChainedIdentityResolver from smithy_core.aio.interfaces.eventstream import EventPublisher, EventReceiver from smithy_http.aio.crt import AWSCRTHTTPClient _AWS_SDK_AVAILABLE = True except ImportError: _AWS_SDK_AVAILABLE = False @dataclass class Credentials: access_key_id: str secret_access_key: str session_token: str | None = None @dataclass class STTOptions: sample_rate: int language: LanguageCode | None encoding: str vocabulary_name: NotGivenOr[str] session_id: NotGivenOr[str] vocab_filter_method: NotGivenOr[str] vocab_filter_name: NotGivenOr[str] show_speaker_label: NotGivenOr[bool] enable_channel_identification: NotGivenOr[bool] number_of_channels: NotGivenOr[int] enable_partial_results_stabilization: NotGivenOr[bool] partial_results_stability: NotGivenOr[str] language_model_name: NotGivenOr[str] region: str identify_language: bool identify_multiple_languages: bool language_options: NotGivenOr[str] preferred_language: NotGivenOr[str] vocabulary_names: NotGivenOr[str] vocabulary_filter_names: NotGivenOr[str] class STT(stt.STT): def __init__( self, *, region: NotGivenOr[str] = NOT_GIVEN, sample_rate: int = 24000, language: str | None = "en-US", encoding: str = "pcm", vocabulary_name: NotGivenOr[str] = NOT_GIVEN, session_id: NotGivenOr[str] = NOT_GIVEN, vocab_filter_method: NotGivenOr[str] = NOT_GIVEN, vocab_filter_name: NotGivenOr[str] = NOT_GIVEN, show_speaker_label: NotGivenOr[bool] = NOT_GIVEN, enable_channel_identification: NotGivenOr[bool] = NOT_GIVEN, number_of_channels: NotGivenOr[int] = NOT_GIVEN, enable_partial_results_stabilization: NotGivenOr[bool] = NOT_GIVEN, partial_results_stability: NotGivenOr[str] = NOT_GIVEN, language_model_name: NotGivenOr[str] = NOT_GIVEN, credentials: NotGivenOr[Credentials] = NOT_GIVEN, identify_language: bool = False, identify_multiple_languages: bool = False, language_options: NotGivenOr[str] = NOT_GIVEN, preferred_language: NotGivenOr[str] = NOT_GIVEN, vocabulary_names: NotGivenOr[str] = NOT_GIVEN, vocabulary_filter_names: NotGivenOr[str] = NOT_GIVEN, ): super().__init__( capabilities=stt.STTCapabilities( streaming=True, interim_results=True, aligned_transcript="word", offline_recognize=False, ) ) if not _AWS_SDK_AVAILABLE: raise ImportError( "The 'aws_sdk_transcribe_streaming' package is not installed. " "This implementation requires Python 3.12+ and the 'aws_sdk_transcribe_streaming' dependency." ) if not is_given(region): region = os.getenv("AWS_REGION") or DEFAULT_REGION if identify_language and identify_multiple_languages: raise ValueError( "identify_language and identify_multiple_languages are mutually exclusive. " "Set only one to True." ) # When auto language detection is enabled, language_code must not be set lang: LanguageCode | None = None if not identify_language and not identify_multiple_languages: lang = LanguageCode(language) if language else LanguageCode("en-US") self._config = STTOptions( language=lang, sample_rate=sample_rate, encoding=encoding, vocabulary_name=vocabulary_name, session_id=session_id, vocab_filter_method=vocab_filter_method, vocab_filter_name=vocab_filter_name, show_speaker_label=show_speaker_label, enable_channel_identification=enable_channel_identification, number_of_channels=number_of_channels, enable_partial_results_stabilization=enable_partial_results_stabilization, partial_results_stability=partial_results_stability, language_model_name=language_model_name, region=region, identify_language=identify_language, identify_multiple_languages=identify_multiple_languages, language_options=language_options, preferred_language=preferred_language, vocabulary_names=vocabulary_names, vocabulary_filter_names=vocabulary_filter_names, ) self._credentials = credentials if is_given(credentials) else None @property def model(self) -> str: return ( self._config.language_model_name if is_given(self._config.language_model_name) else "unknown" ) @property def provider(self) -> str: return "Amazon Transcribe" async def aclose(self) -> None: await super().aclose() async def _recognize_impl( self, buffer: utils.AudioBuffer, *, language: NotGivenOr[str] = NOT_GIVEN, conn_options: APIConnectOptions, ) -> stt.SpeechEvent: raise NotImplementedError("Amazon Transcribe does not support single frame recognition") def stream( self, *, language: NotGivenOr[str] = NOT_GIVEN, conn_options: APIConnectOptions = DEFAULT_API_CONNECT_OPTIONS, ) -> SpeechStream: return SpeechStream( stt=self, conn_options=conn_options, opts=self._config, credentials=self._credentials, ) class SpeechStream(stt.SpeechStream): def __init__( self, stt: STT, opts: STTOptions, conn_options: APIConnectOptions = DEFAULT_API_CONNECT_OPTIONS, credentials: Credentials | None = None, ) -> None: super().__init__(stt=stt, conn_options=conn_options, sample_rate=opts.sample_rate) self._opts = opts self._credentials = credentials self._http_client = AWSCRTHTTPClient() async def _run(self) -> None: while True: config_kwargs: dict[str, Any] = {"region": self._opts.region} if self._credentials: # Use a credentials resolver for explicit credentials # for some reason, Config with direct values doesn't work class StaticCredsResolver: def __init__(self, creds: Credentials): self._identity = AWSCredentialsIdentity( access_key_id=creds.access_key_id, secret_access_key=creds.secret_access_key, session_token=creds.session_token, ) async def get_identity(self, **kwargs: Any) -> AWSCredentialsIdentity: return self._identity config_kwargs["aws_credentials_identity_resolver"] = StaticCredsResolver( self._credentials ) else: config_kwargs["aws_credentials_identity_resolver"] = ChainedIdentityResolver( resolvers=( StaticCredentialsResolver(), EnvironmentCredentialsResolver(), ContainerCredentialsResolver(http_client=self._http_client), IMDSCredentialsResolver(http_client=self._http_client), ) ) client: TranscribeStreamingClient = TranscribeStreamingClient( config=Config(**config_kwargs) ) live_config = { "media_sample_rate_hertz": self._opts.sample_rate, "media_encoding": self._opts.encoding, "vocabulary_name": self._opts.vocabulary_name, "session_id": self._opts.session_id, "vocab_filter_method": self._opts.vocab_filter_method, "vocab_filter_name": self._opts.vocab_filter_name, "show_speaker_label": self._opts.show_speaker_label, "enable_channel_identification": self._opts.enable_channel_identification, "number_of_channels": self._opts.number_of_channels, "enable_partial_results_stabilization": self._opts.enable_partial_results_stabilization, "partial_results_stability": self._opts.partial_results_stability, "language_model_name": self._opts.language_model_name, } # Auto language detection is mutually exclusive with language_code if self._opts.identify_language: live_config["identify_language"] = True if is_given(self._opts.language_options): live_config["language_options"] = self._opts.language_options if is_given(self._opts.preferred_language): live_config["preferred_language"] = self._opts.preferred_language if is_given(self._opts.vocabulary_names): live_config["vocabulary_names"] = self._opts.vocabulary_names if is_given(self._opts.vocabulary_filter_names): live_config["vocabulary_filter_names"] = self._opts.vocabulary_filter_names elif self._opts.identify_multiple_languages: live_config["identify_multiple_languages"] = True if is_given(self._opts.language_options): live_config["language_options"] = self._opts.language_options if is_given(self._opts.preferred_language): live_config["preferred_language"] = self._opts.preferred_language if is_given(self._opts.vocabulary_names): live_config["vocabulary_names"] = self._opts.vocabulary_names if is_given(self._opts.vocabulary_filter_names): live_config["vocabulary_filter_names"] = self._opts.vocabulary_filter_names else: if self._opts.language: live_config["language_code"] = self._opts.language filtered_config: dict[str, Any] = {} for k, v in live_config.items(): if isinstance(v, bool): filtered_config[k] = v elif isinstance(v, (int, float)): filtered_config[k] = v elif v is not None and is_given(v): filtered_config[k] = v tasks: list[asyncio.Task[Any]] = [] try: stream = await client.start_stream_transcription( input=StartStreamTranscriptionInput(**filtered_config) ) # Get the output stream _, output_stream = await stream.await_output() async def input_generator( audio_stream: EventPublisher[AudioStream], ) -> None: try: async for frame in self._input_ch: if isinstance(frame, rtc.AudioFrame): await audio_stream.send( AudioStreamAudioEvent( value=AudioEvent(audio_chunk=frame.data.tobytes()) ) ) finally: # Send empty frame to close (required by AWS Transcribe) try: await audio_stream.send( AudioStreamAudioEvent(value=AudioEvent(audio_chunk=b"")) ) except Exception: pass finally: with contextlib.suppress(Exception): await audio_stream.close() async def handle_transcript_events( output_stream: EventReceiver[TranscriptResultStream], ) -> None: try: async for event in output_stream: if isinstance(event.value, TranscriptEvent): self._process_transcript_event(event.value) except BadRequestException as e: if ( e.message and "complete signal was sent without the preceding empty frame" in e.message ): # This can happen during cancellation if the empty frame wasn't sent in time logger.warning( "AWS Transcribe stream closed with empty frame error (this is usually harmless)" ) else: raise except concurrent.futures.InvalidStateError: logger.warning( "AWS Transcribe stream closed unexpectedly (InvalidStateError)" ) pass tasks = [ asyncio.create_task(input_generator(stream.input_stream)), asyncio.create_task(handle_transcript_events(output_stream)), ] gather_future = asyncio.gather(*tasks) await asyncio.shield(gather_future) except BadRequestException as e: if e.message and e.message.startswith("Your request timed out"): # AWS times out after 15s of inactivity, this tends to happen # at the end of the session, when the input is gone, we'll ignore it and # just treat it as a silent retry logger.info("restarting transcribe session") continue else: raise e finally: if tasks: # Close input stream first await utils.aio.gracefully_cancel(tasks[0]) # Wait for output stream to close cleanly try: await asyncio.wait_for(tasks[1], timeout=3.0) except (asyncio.TimeoutError, asyncio.CancelledError): await utils.aio.gracefully_cancel(tasks[1]) # Ensure gather future is retrieved to avoid "exception never retrieved" with contextlib.suppress(Exception): await gather_future def _process_transcript_event(self, transcript_event: TranscriptEvent) -> None: if not transcript_event.transcript or not transcript_event.transcript.results: return stream = transcript_event.transcript.results for resp in stream: if resp.start_time is not None and resp.start_time == 0.0: self._event_ch.send_nowait( stt.SpeechEvent(type=stt.SpeechEventType.START_OF_SPEECH) ) if resp.end_time is not None and resp.end_time > 0.0: if resp.is_partial: self._event_ch.send_nowait( stt.SpeechEvent( type=stt.SpeechEventType.INTERIM_TRANSCRIPT, alternatives=[self._streaming_recognize_response_to_speech_data(resp)], ) ) else: self._event_ch.send_nowait( stt.SpeechEvent( type=stt.SpeechEventType.FINAL_TRANSCRIPT, alternatives=[self._streaming_recognize_response_to_speech_data(resp)], ) ) if not resp.is_partial: self._event_ch.send_nowait(stt.SpeechEvent(type=stt.SpeechEventType.END_OF_SPEECH)) def _streaming_recognize_response_to_speech_data(self, resp: Result) -> stt.SpeechData: confidence = 0.0 if resp.alternatives and (items := resp.alternatives[0].items): confidence = items[0].confidence or 0.0 detected_lang = resp.language_code or self._opts.language or "en-US" # Populate source_languages when language identification is active source_languages: list[LanguageCode] | None = None if ( self._opts.identify_language or self._opts.identify_multiple_languages ) and resp.language_code: source_languages = [LanguageCode(resp.language_code)] return stt.SpeechData( language=LanguageCode(detected_lang), start_time=(resp.start_time or 0.0) + self.start_time_offset, end_time=(resp.end_time or 0.0) + self.start_time_offset, text=resp.alternatives[0].transcript if resp.alternatives else "", confidence=confidence, source_languages=source_languages, words=[ TimedString( text=item.content, start_time=item.start_time + self.start_time_offset, end_time=item.end_time + self.start_time_offset, start_time_offset=self.start_time_offset, confidence=item.confidence or 0.0, ) for item in resp.alternatives[0].items ] if resp.alternatives and resp.alternatives[0].items else None, )