# 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 dataclasses import json import os import weakref from collections.abc import Callable from dataclasses import dataclass from urllib.parse import urlencode import aiohttp from livekit.agents import ( DEFAULT_API_CONNECT_OPTIONS, APIConnectOptions, APIStatusError, LanguageCode, stt, utils, ) from livekit.agents.types import NOT_GIVEN, NotGivenOr from livekit.agents.utils import AudioBuffer, is_given from .log import logger _STREAMING_PATH = "/audio/transcriptions/streaming" class _PeriodicCollector: def __init__(self, duration: float, callback: Callable[[float], None]): self._duration = duration self._callback = callback self._collected_value = 0.0 self._task: asyncio.Task | None = None self._lock = asyncio.Lock() async def push(self, value: float) -> None: async with self._lock: self._collected_value += value if not self._task: self._task = asyncio.create_task(self._run()) async def flush(self) -> None: async with self._lock: if self._task: self._task.cancel() try: await self._task except asyncio.CancelledError: pass self._task = None if self._collected_value > 0: self._callback(self._collected_value) self._collected_value = 0.0 async def _run(self) -> None: await asyncio.sleep(self._duration) async with self._lock: self._callback(self._collected_value) self._collected_value = 0.0 self._task = None @dataclass class STTOptions: model: NotGivenOr[str] sample_rate: int language: NotGivenOr[LanguageCode] = NOT_GIVEN prompt: NotGivenOr[str] = NOT_GIVEN temperature: NotGivenOr[float] = NOT_GIVEN skip_vad: NotGivenOr[bool] = NOT_GIVEN vad_kwargs: NotGivenOr[dict] = NOT_GIVEN text_timeout_seconds: float = 1.0 response_format: str = "verbose_json" timestamp_granularities: NotGivenOr[list[str]] = NOT_GIVEN base_url: NotGivenOr[str] = NOT_GIVEN class STT(stt.STT): def __init__( self, *, model: NotGivenOr[str] = NOT_GIVEN, api_key: NotGivenOr[str] = NOT_GIVEN, sample_rate: int = 16000, language: NotGivenOr[str] = NOT_GIVEN, prompt: NotGivenOr[str] = NOT_GIVEN, temperature: NotGivenOr[float] = NOT_GIVEN, skip_vad: NotGivenOr[bool] = NOT_GIVEN, vad_kwargs: NotGivenOr[dict] = NOT_GIVEN, text_timeout_seconds: float = 1.0, timestamp_granularities: NotGivenOr[list[str]] = NOT_GIVEN, response_format: str = "verbose_json", http_session: aiohttp.ClientSession | None = None, base_url: str = "wss://audio-streaming.us-virginia-1.direct.fireworks.ai/v1", ): """ Create a new instance of Fireworks AI STT. Args: model: The Fireworks AI STT model to use. Defaults to NOT_GIVEN (server uses default model). language: The target language for transcription. Defaults to NOT_GIVEN (server detects language automatically). Full list: https://fireworks.ai/docs/api-reference/audio-streaming-transcriptions#supported-languages prompt: The input prompt that the model will use when generating the transcription. Defaults to NOT_GIVEN. temperature: Sampling temperature to use when decoding text tokens during transcription. Defaults to NOT_GIVEN. skip_vad: Whether to skip server-side VAD. Defaults to NOT_GIVEN. vad_kwargs: The optional kwargs to pass to the VAD model. Defaults to NOT_GIVEN. Example: Set to {"threshold": 0.15} to adjust the VAD threshold. text_timeout_seconds: Duration of silence before marking transcript as final. Defaults to 1.0. timestamp_granularities: The timestamp granularities to populate for this streaming transcription. Defaults to NOT_GIVEN. Set to "word,segment" to enable timestamp granularities. response_format: The format in which to return the response. Default to "verbose_json". base_url: The base URL for the Fireworks AI STT. Defaults to "wss://audio-streaming.us-virginia-1.direct.fireworks.ai/v1". api_key: The Fireworks AI API key. If not provided, it will be read from the FIREWORKS_API_KEY environment variable. http_session: Optional aiohttp ClientSession to use for requests. Raises: ValueError: If no API key is provided, found in environment variables, or if a parameter is invalid. """ super().__init__( capabilities=stt.STTCapabilities( streaming=True, interim_results=True, aligned_transcript=False, offline_recognize=False, ), ) if sample_rate != 16000: raise ValueError("FireworksAI STT only supports a sample rate of 16000") if not 1.0 <= text_timeout_seconds <= 29.0: raise ValueError("text_timeout_seconds must be between 1.0 and 29.0") fireworks_api_key = api_key if is_given(api_key) else os.environ.get("FIREWORKS_API_KEY") if fireworks_api_key is None: raise ValueError( "Fireworks API key is required. " "Pass one in via the `api_key` parameter, " "or set it as the `FIREWORKS_API_KEY` environment variable" ) self._api_key = fireworks_api_key self._opts = STTOptions( model=model, sample_rate=sample_rate, language=LanguageCode(language) if isinstance(language, str) else language, prompt=prompt, temperature=temperature, skip_vad=skip_vad, vad_kwargs=vad_kwargs, text_timeout_seconds=text_timeout_seconds, response_format=response_format, timestamp_granularities=timestamp_granularities, base_url=base_url, ) self._session = http_session self._streams = weakref.WeakSet[SpeechStream]() @property def model(self) -> str: return self._opts.model if is_given(self._opts.model) else "unknown" @property def provider(self) -> str: return "FireworksAI" @property def session(self) -> aiohttp.ClientSession: if not self._session: self._session = utils.http_context.http_session() return self._session async def _recognize_impl( self, buffer: AudioBuffer, *, language: NotGivenOr[str] = NOT_GIVEN, conn_options: APIConnectOptions, ) -> stt.SpeechEvent: raise NotImplementedError( "FireworksAI STT does not support batch recognition, use stream() instead" ) def stream( self, *, language: NotGivenOr[str] = NOT_GIVEN, conn_options: APIConnectOptions = DEFAULT_API_CONNECT_OPTIONS, ) -> SpeechStream: config = dataclasses.replace(self._opts) stream = SpeechStream( stt=self, opts=config, conn_options=conn_options, api_key=self._api_key, http_session=self.session, ) self._streams.add(stream) return stream def update_options( self, *, model: NotGivenOr[str] = NOT_GIVEN, language: NotGivenOr[str] = NOT_GIVEN, prompt: NotGivenOr[str] = NOT_GIVEN, temperature: NotGivenOr[float] = NOT_GIVEN, skip_vad: NotGivenOr[bool] = NOT_GIVEN, vad_kwargs: NotGivenOr[dict] = NOT_GIVEN, text_timeout_seconds: NotGivenOr[float] = NOT_GIVEN, timestamp_granularities: NotGivenOr[list[str]] = NOT_GIVEN, ) -> None: if is_given(model): self._opts.model = model if is_given(language): self._opts.language = LanguageCode(language) if is_given(prompt): self._opts.prompt = prompt if is_given(temperature): self._opts.temperature = temperature if is_given(skip_vad): self._opts.skip_vad = skip_vad if is_given(vad_kwargs): self._opts.vad_kwargs = vad_kwargs if is_given(text_timeout_seconds): if not 1.0 <= text_timeout_seconds <= 29.0: raise ValueError("text_timeout_seconds must be between 1.0 and 29.0") self._opts.text_timeout_seconds = text_timeout_seconds if is_given(timestamp_granularities): self._opts.timestamp_granularities = timestamp_granularities for stream in self._streams: stream.update_options( model=model, language=language, prompt=prompt, temperature=temperature, skip_vad=skip_vad, vad_kwargs=vad_kwargs, text_timeout_seconds=text_timeout_seconds, timestamp_granularities=timestamp_granularities, ) class SpeechStream(stt.SpeechStream): _CLOSE_MSG: str = json.dumps({"checkpoint_id": "final"}) def __init__( self, *, stt: STT, opts: STTOptions, conn_options: APIConnectOptions, api_key: str, http_session: aiohttp.ClientSession, ) -> None: super().__init__(stt=stt, conn_options=conn_options, sample_rate=opts.sample_rate) self._opts = opts self._api_key = api_key self._session = http_session self._transcript_state: dict[str, str] = {} self._reconnect_event = asyncio.Event() self._speaking = False self._final_segments_length: dict[int, int] = {} self._last_final_segment_id = -1 self._audio_duration_collector = _PeriodicCollector( callback=self._on_audio_duration_report, duration=10.0, ) def update_options( self, *, model: NotGivenOr[str] = NOT_GIVEN, language: NotGivenOr[str] = NOT_GIVEN, prompt: NotGivenOr[str] = NOT_GIVEN, temperature: NotGivenOr[float] = NOT_GIVEN, skip_vad: NotGivenOr[bool] = NOT_GIVEN, vad_kwargs: NotGivenOr[dict] = NOT_GIVEN, text_timeout_seconds: NotGivenOr[float] = NOT_GIVEN, timestamp_granularities: NotGivenOr[list[str]] = NOT_GIVEN, ) -> None: if is_given(model): self._opts.model = model if is_given(language): self._opts.language = LanguageCode(language) if is_given(prompt): self._opts.prompt = prompt if is_given(temperature): self._opts.temperature = temperature if is_given(skip_vad): self._opts.skip_vad = skip_vad if is_given(vad_kwargs): self._opts.vad_kwargs = vad_kwargs if is_given(text_timeout_seconds): self._opts.text_timeout_seconds = text_timeout_seconds if is_given(timestamp_granularities): self._opts.timestamp_granularities = timestamp_granularities self._reconnect_event.set() async def _run(self) -> None: """ Run a single websocket connection to Fireworks and make sure to reconnect when something went wrong. """ closing_ws = False async def send_task(ws: aiohttp.ClientWebSocketResponse) -> None: nonlocal closing_ws samples_per_buffer = self._opts.sample_rate // 20 # 50ms chunk audio_bstream = utils.audio.AudioByteStream( sample_rate=self._opts.sample_rate, num_channels=1, samples_per_channel=samples_per_buffer, ) async for data in self._input_ch: if isinstance(data, self._FlushSentinel): frames = audio_bstream.flush() else: frames = audio_bstream.write(data.data.tobytes()) for frame in frames: await self._audio_duration_collector.push(frame.duration) await ws.send_bytes(frame.data.tobytes()) closing_ws = True await ws.send_str(self._CLOSE_MSG) async def recv_task(ws: aiohttp.ClientWebSocketResponse) -> None: nonlocal closing_ws while True: try: msg = await asyncio.wait_for(ws.receive(), timeout=5) except asyncio.TimeoutError: if closing_ws: break continue if msg.type in ( aiohttp.WSMsgType.CLOSED, aiohttp.WSMsgType.CLOSE, aiohttp.WSMsgType.CLOSING, ): if closing_ws: return raise APIStatusError( "Fireworks connection closed unexpectedly", status_code=ws.close_code or -1, body=f"{msg.data=} {msg.extra=}", ) if msg.type != aiohttp.WSMsgType.TEXT: logger.error("unexpected FireworksAI message type %s", msg.type) continue try: self._process_stream_event(json.loads(msg.data)) except Exception: logger.exception("failed to process FireworksAI message") ws: aiohttp.ClientWebSocketResponse | None = None while True: try: ws = await self._connect_ws() tasks = [ asyncio.create_task(send_task(ws)), asyncio.create_task(recv_task(ws)), ] wait_reconnect_task = asyncio.create_task(self._reconnect_event.wait()) try: done, _ = await asyncio.wait( (asyncio.gather(*tasks), wait_reconnect_task), return_when=asyncio.FIRST_COMPLETED, ) for task in done: if task != wait_reconnect_task: task.result() if wait_reconnect_task not in done: break self._reconnect_event.clear() finally: await utils.aio.gracefully_cancel(*tasks, wait_reconnect_task) finally: if self._speaking: self._speaking = False end_event = stt.SpeechEvent(type=stt.SpeechEventType.END_OF_SPEECH) self._event_ch.send_nowait(end_event) if ws is not None: await ws.close() await self._audio_duration_collector.flush() async def _connect_ws(self) -> aiohttp.ClientWebSocketResponse: live_config = { "model": self._opts.model if is_given(self._opts.model) else None, "language": self._opts.language if is_given(self._opts.language) else None, "prompt": self._opts.prompt if is_given(self._opts.prompt) else None, "temperature": self._opts.temperature if is_given(self._opts.temperature) else None, "skip_vad": self._opts.skip_vad if is_given(self._opts.skip_vad) else None, "vad_kwargs": self._opts.vad_kwargs if is_given(self._opts.vad_kwargs) else None, "text_timeout_seconds": self._opts.text_timeout_seconds, "response_format": self._opts.response_format, "timestamp_granularities": ( self._opts.timestamp_granularities if is_given(self._opts.timestamp_granularities) else None ), } headers = { "User-Agent": "LiveKit Agents", "Authorization": self._api_key, } ws_url = str(self._opts.base_url).rstrip("/") + _STREAMING_PATH filtered_config = {k: v for k, v in live_config.items() if v is not None} url = f"{ws_url}?{urlencode(filtered_config, doseq=True)}" ws = await self._session.ws_connect(url, headers=headers) logger.info("connected to Fireworks AI STT", extra={"url": url}) return ws def _process_stream_event(self, data: dict) -> None: if "segments" in data and data["segments"]: latest_segment = max(data["segments"], key=lambda s: s["id"]) max_segment_id = latest_segment["id"] for segment in data["segments"]: segment_id = segment["id"] if segment_id < self._last_final_segment_id: continue if segment_id == self._last_final_segment_id: finalized_word_count = self._final_segments_length.get(segment_id, 0) words = segment.get("words", []) if isinstance(words, list) and finalized_word_count < len(words): new_words = words[finalized_word_count:] new_text = " ".join(w["word"] for w in new_words if "word" in w).strip() self._transcript_state[segment_id] = new_text elif segment_id in self._transcript_state: del self._transcript_state[segment_id] else: self._transcript_state[segment["id"]] = segment["text"] for local_segment_id in list(self._transcript_state.keys()): if local_segment_id > max_segment_id: del self._transcript_state[local_segment_id] # The state dictionary may not be sorted, so we must sort it by the segment ID # before joining the text. sorted_segments = sorted(self._transcript_state.items(), key=lambda item: int(item[0])) full_transcript = " ".join([text for _, text in sorted_segments]) if not full_transcript: return if not self._speaking: self._speaking = True start_event = stt.SpeechEvent(type=stt.SpeechEventType.START_OF_SPEECH) self._event_ch.send_nowait(start_event) is_final = False words = latest_segment.get("words") if words and isinstance(words, list) and words: last_word = words[-1] if isinstance(last_word, dict) and last_word.get("is_final") is True: is_final = True if is_final: final_event = stt.SpeechEvent( type=stt.SpeechEventType.FINAL_TRANSCRIPT, alternatives=[ stt.SpeechData( language=LanguageCode(self._opts.language or ""), text=full_transcript ) ], ) self._event_ch.send_nowait(final_event) self._transcript_state.clear() self._last_final_segment_id = max_segment_id words = latest_segment.get("words") if isinstance(words, list): self._final_segments_length[max_segment_id] = len(words) else: interim_event = stt.SpeechEvent( type=stt.SpeechEventType.INTERIM_TRANSCRIPT, alternatives=[ stt.SpeechData( language=LanguageCode(self._opts.language or ""), text=full_transcript ) ], ) self._event_ch.send_nowait(interim_event) def _on_audio_duration_report(self, duration: float) -> None: usage_event = stt.SpeechEvent( type=stt.SpeechEventType.RECOGNITION_USAGE, recognition_usage=stt.RecognitionUsage(audio_duration=duration), ) self._event_ch.send_nowait(usage_event)