/livekit-plugins/livekit-plugins-smallestai/livekit/plugins/smallestai/stt.py

# 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 json
import os
import time
import weakref
from collections import Counter
from collections.abc import Callable
from dataclasses import dataclass, replace
from typing import Any, Generic, TypeVar
from urllib.parse import urlencode

import aiohttp

from livekit import rtc
from livekit.agents import (
    DEFAULT_API_CONNECT_OPTIONS,
    APIConnectionError,
    APIConnectOptions,
    APIStatusError,
    APITimeoutError,
    LanguageCode,
    stt,
    utils,
)
from livekit.agents.types import NOT_GIVEN, NotGivenOr, TimedString
from livekit.agents.utils import AudioBuffer, is_given

from .log import logger
from .models import STTEncoding, STTModels
from .version import __version__

NUM_CHANNELS = 1
# Base URL for the Smallest AI API.
# Streaming: wss://api.smallest.ai/waves/v1/stt/live?model={model}
# Batch:     https://api.smallest.ai/waves/v1/stt/?model={model}
SMALLEST_STT_BASE_URL = "https://api.smallest.ai/waves/v1"

# Models that support real-time streaming. All others are batch-only and will be
# wrapped with a StreamAdapter by the agent framework automatically.
_STREAMING_MODELS: frozenset[str] = frozenset({"pulse"})

# ---------------------------------------------------------------------------
# Minimal PeriodicCollector — same logic as livekit-plugins-deepgram/_utils.py
# ---------------------------------------------------------------------------

T = TypeVar("T")


class _PeriodicCollector(Generic[T]):
    def __init__(self, callback: Callable[[T], None], *, duration: float) -> None:
        self._duration = duration
        self._callback = callback
        self._last_flush_time = time.monotonic()
        self._total: T | None = None

    def push(self, value: T) -> None:
        if self._total is None:
            self._total = value
        else:
            self._total += value  # type: ignore[operator]
        if time.monotonic() - self._last_flush_time >= self._duration:
            self.flush()

    def flush(self) -> None:
        if self._total is not None:
            self._callback(self._total)
            self._total = None
        self._last_flush_time = time.monotonic()


# ---------------------------------------------------------------------------


@dataclass
class _STTOptions:
    model: STTModels | str
    api_key: str
    language: str  # BCP-47 code, e.g. "en", "hi"; use "multi" for auto-detection
    sample_rate: int
    encoding: STTEncoding | str
    word_timestamps: bool
    diarize: bool
    eou_timeout_ms: int  # end-of-utterance silence timeout in ms; valid range 100–10000ms
    base_url: str


class STT(stt.STT):
    def __init__(
        self,
        *,
        model: STTModels | str = "pulse",
        language: str = "en",
        sample_rate: int = 16000,
        encoding: STTEncoding | str = "linear16",
        word_timestamps: bool = True,
        diarize: bool = False,
        eou_timeout_ms: int = 100,
        api_key: str | None = None,
        http_session: aiohttp.ClientSession | None = None,
        base_url: str = SMALLEST_STT_BASE_URL,
    ) -> None:
        """Create a new instance of Smallest AI STT.

        Args:
            model: STT model to use. ``"pulse"`` supports streaming and batch
                transcription across 38 languages. ``"pulse-pro"`` is a
                higher-accuracy English-only model available for batch
                (pre-recorded) transcription only — calling ``stream()`` with
                ``pulse-pro`` raises ``ValueError``.
            language: BCP-47 language code (e.g. "en", "hi", "fr"). Use "multi"
                for automatic language detection across 39 supported languages.
                ``pulse-pro`` only supports ``"en"``.
            sample_rate: Audio sample rate in Hz. Supported: 8000, 16000, 22050,
                24000, 44100, 48000. Defaults to 16000.
            encoding: PCM encoding of the audio stream. Use "linear16" for raw
                16-bit PCM (the default and most compatible choice for streaming).
            word_timestamps: Include per-word start/end timestamps and confidence
                scores in transcripts. Supported by both ``pulse`` and
                ``pulse-pro``. Defaults to True.
            diarize: Enable speaker diarization. When True, each word includes a
                speaker ID (integer during streaming, string label in batch).
                Defaults to False.
            eou_timeout_ms: Milliseconds of silence before the server considers an
                utterance complete and emits a final transcript. Must be between 100 and
                10000ms. Defaults to 100ms (the minimum) so that server-side EOU adds
                minimal latency alongside LiveKit's own end-of-turn detection. If omitted,
                the server applies an 800ms default. Note: the Smallest AI API will soon
                support disabling server-side EOU entirely, which will allow LiveKit's
                end-of-turn detection to be used exclusively.
            api_key: Smallest AI API key. Falls back to the SMALLEST_API_KEY
                environment variable if not provided.
            http_session: An existing aiohttp ClientSession to reuse.
            base_url: Override the default API base URL.
        """
        super().__init__(
            capabilities=stt.STTCapabilities(
                streaming=model in _STREAMING_MODELS,
                interim_results=True,
                diarization=diarize,
                aligned_transcript="word" if word_timestamps else False,
            )
        )

        api_key = api_key or os.environ.get("SMALLEST_API_KEY")
        if not api_key:
            raise ValueError(
                "Smallest AI API key is required, either as argument or set "
                "SMALLEST_API_KEY environment variable"
            )

        self._opts = _STTOptions(
            model=model,
            api_key=api_key,
            language=language,
            sample_rate=sample_rate,
            encoding=encoding,
            word_timestamps=word_timestamps,
            diarize=diarize,
            eou_timeout_ms=eou_timeout_ms,
            base_url=base_url,
        )
        self._session = http_session
        self._streams: weakref.WeakSet[SpeechStream] = weakref.WeakSet()

    @property
    def model(self) -> str:
        return self._opts.model

    @property
    def provider(self) -> str:
        return "SmallestAI"

    def _ensure_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 = DEFAULT_API_CONNECT_OPTIONS,
    ) -> stt.SpeechEvent:
        config = self._sanitize_options(language=language)
        params: dict[str, Any] = {
            "model": config.model,
            "language": config.language,
            "encoding": config.encoding,
            "sample_rate": config.sample_rate,
            "word_timestamps": str(config.word_timestamps).lower(),
            "diarize": str(config.diarize).lower(),
        }

        try:
            async with self._ensure_session().post(
                url=f"{config.base_url}/stt/",
                headers={
                    "Authorization": f"Bearer {config.api_key}",
                    "Content-Type": "application/octet-stream",
                    "X-Source": "livekit",
                    "X-LiveKit-Version": __version__,
                },
                params=params,
                # to_wav_bytes() produces a valid WAV file; the server auto-detects format.
                data=rtc.combine_audio_frames(buffer).to_wav_bytes(),
                timeout=aiohttp.ClientTimeout(
                    total=30,
                    sock_connect=conn_options.timeout,
                ),
            ) as resp:
                resp.raise_for_status()
                data = await resp.json()
                return _batch_transcription_to_speech_event(config.language, data)

        except asyncio.TimeoutError as e:
            raise APITimeoutError() from e
        except aiohttp.ClientResponseError as e:
            raise APIStatusError(
                message=e.message,
                status_code=e.status,
                request_id=None,
                body=None,
            ) from e
        except Exception as e:
            raise APIConnectionError() from e

    def stream(
        self,
        *,
        language: NotGivenOr[str] = NOT_GIVEN,
        conn_options: APIConnectOptions = DEFAULT_API_CONNECT_OPTIONS,
    ) -> SpeechStream:
        if not self.capabilities.streaming:
            raise ValueError(
                f"{self._opts.model} does not support streaming; use recognize() for batch transcription"
            )
        config = self._sanitize_options(language=language)
        stream = SpeechStream(
            stt=self,
            conn_options=conn_options,
            opts=config,
            http_session=self._ensure_session(),
        )
        self._streams.add(stream)
        return stream

    def update_options(
        self,
        *,
        model: NotGivenOr[STTModels | str] = NOT_GIVEN,
        language: NotGivenOr[str] = NOT_GIVEN,
        sample_rate: NotGivenOr[int] = NOT_GIVEN,
        encoding: NotGivenOr[STTEncoding | str] = NOT_GIVEN,
        eou_timeout_ms: NotGivenOr[int] = NOT_GIVEN,
    ) -> None:
        """Update STT options; propagates to all active streams (triggers reconnect)."""
        if is_given(model):
            self._opts.model = model
            self._capabilities.streaming = model in _STREAMING_MODELS
        if is_given(language):
            self._opts.language = language
        if is_given(sample_rate):
            self._opts.sample_rate = sample_rate
        if is_given(encoding):
            self._opts.encoding = encoding
        if is_given(eou_timeout_ms):
            self._opts.eou_timeout_ms = eou_timeout_ms

        for stream in self._streams:
            stream.update_options(
                model=model,
                language=language,
                sample_rate=sample_rate,
                encoding=encoding,
                eou_timeout_ms=eou_timeout_ms,
            )

    def _sanitize_options(self, *, language: NotGivenOr[str] = NOT_GIVEN) -> _STTOptions:
        config = replace(self._opts)
        if is_given(language):
            config.language = language
        return config


class SpeechStream(stt.SpeechStream):
    # Signals end of stream: server flushes remaining audio, emits final transcripts,
    # and responds with is_last=True before closing the session.
    # Use {"type": "finalize"} mid-session to force is_final without closing.
    _CLOSE_STREAM_MSG: str = json.dumps({"type": "close_stream"})

    def __init__(
        self,
        *,
        stt: STT,
        opts: _STTOptions,
        conn_options: APIConnectOptions,
        http_session: aiohttp.ClientSession,
    ) -> None:
        super().__init__(stt=stt, conn_options=conn_options, sample_rate=opts.sample_rate)
        self._opts = opts
        self._session = http_session
        self._speaking = False
        self._session_id = ""
        self._reconnect_event = asyncio.Event()
        self._audio_duration_collector = _PeriodicCollector(
            callback=self._on_audio_duration_report,
            duration=5.0,
        )

    def update_options(
        self,
        *,
        model: NotGivenOr[STTModels | str] = NOT_GIVEN,
        language: NotGivenOr[str] = NOT_GIVEN,
        sample_rate: NotGivenOr[int] = NOT_GIVEN,
        encoding: NotGivenOr[STTEncoding | str] = NOT_GIVEN,
        eou_timeout_ms: NotGivenOr[int] = NOT_GIVEN,
    ) -> None:
        if is_given(model):
            self._opts.model = model
        if is_given(language):
            self._opts.language = language
        if is_given(sample_rate):
            self._opts.sample_rate = sample_rate
        if is_given(encoding):
            self._opts.encoding = encoding
        if is_given(eou_timeout_ms):
            self._opts.eou_timeout_ms = eou_timeout_ms
        self._reconnect_event.set()

    async def _run(self) -> None:
        closing_ws = False

        @utils.log_exceptions(logger=logger)
        async def send_task(ws: aiohttp.ClientWebSocketResponse) -> None:
            nonlocal closing_ws

            # Send audio in 50ms chunks; matches the 50–100ms guidance from Smallest AI docs.
            samples_per_chunk = self._opts.sample_rate // 20
            audio_bstream = utils.audio.AudioByteStream(
                sample_rate=self._opts.sample_rate,
                num_channels=NUM_CHANNELS,
                samples_per_channel=samples_per_chunk,
            )

            async for data in self._input_ch:
                if isinstance(data, rtc.AudioFrame):
                    for frame in audio_bstream.write(data.data.tobytes()):
                        self._audio_duration_collector.push(frame.duration)
                        await ws.send_bytes(frame.data.tobytes())
                elif isinstance(data, self._FlushSentinel):
                    # User paused: drain the accumulator so the server gets all buffered
                    # audio. The server's eou_timeout_ms will then detect the silence and
                    # emit a final transcript — no explicit flush message is needed.
                    for frame in audio_bstream.flush():
                        self._audio_duration_collector.push(frame.duration)
                        await ws.send_bytes(frame.data.tobytes())
                    self._audio_duration_collector.flush()

            # Input channel closed: close the stream so the server flushes remaining
            # audio, emits final transcripts, and sends is_last=True.
            closing_ws = True
            await ws.send_str(SpeechStream._CLOSE_STREAM_MSG)

        @utils.log_exceptions(logger=logger)
        async def recv_task(ws: aiohttp.ClientWebSocketResponse) -> None:
            nonlocal closing_ws
            while True:
                msg = await ws.receive()
                if msg.type in (
                    aiohttp.WSMsgType.CLOSED,
                    aiohttp.WSMsgType.CLOSE,
                    aiohttp.WSMsgType.CLOSING,
                ):
                    if closing_ws or self._session.closed:
                        return
                    raise APIStatusError(
                        message="Smallest AI STT connection closed unexpectedly",
                        status_code=ws.close_code or -1,
                        body=f"{msg.data=} {msg.extra=}",
                    )

                if msg.type != aiohttp.WSMsgType.TEXT:
                    logger.warning("unexpected Smallest AI STT message type: %s", msg.type)
                    continue

                try:
                    data = json.loads(msg.data)
                except json.JSONDecodeError:
                    logger.warning("failed to parse Smallest AI STT message: %s", msg.data)
                    continue

                self._process_stream_event(data)

                # Server confirms the session is fully flushed; recv loop can exit.
                if data.get("is_last"):
                    return

        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)),
                ]
                tasks_group = asyncio.gather(*tasks)
                wait_reconnect_task = asyncio.create_task(self._reconnect_event.wait())
                try:
                    done, _ = await asyncio.wait(
                        (tasks_group, 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)
                    tasks_group.cancel()
                    tasks_group.exception()
            finally:
                if ws is not None:
                    await ws.close()

    async def _connect_ws(self) -> aiohttp.ClientWebSocketResponse:
        params: dict[str, Any] = {
            "model": self._opts.model,
            "language": self._opts.language,
            "encoding": self._opts.encoding,
            "sample_rate": self._opts.sample_rate,
            "word_timestamps": str(self._opts.word_timestamps).lower(),
            "diarize": str(self._opts.diarize).lower(),
        }
        params["eou_timeout_ms"] = self._opts.eou_timeout_ms
        ws_url = (
            self._opts.base_url.replace("https://", "wss://", 1).replace("http://", "ws://", 1)
            + "/stt/live"
            + f"?{urlencode(params)}"
        )

        t0 = time.perf_counter()
        try:
            # heartbeat sends standard WebSocket ping frames every 5s, which is sufficient
            # to keep the Smallest AI connection alive without a custom JSON message.
            ws = await asyncio.wait_for(
                self._session.ws_connect(
                    ws_url,
                    headers={
                        "Authorization": f"Bearer {self._opts.api_key}",
                        "X-Source": "livekit",
                        "X-LiveKit-Version": __version__,
                    },
                    heartbeat=5.0,
                ),
                self._conn_options.timeout,
            )
            self._report_connection_acquired(time.perf_counter() - t0, False)
            logger.debug("established Smallest AI STT WebSocket connection")
        except (aiohttp.ClientConnectorError, asyncio.TimeoutError) as e:
            raise APIConnectionError("failed to connect to Smallest AI STT") from e
        return ws

    def _on_audio_duration_report(self, duration: float) -> None:
        self._event_ch.send_nowait(
            stt.SpeechEvent(
                type=stt.SpeechEventType.RECOGNITION_USAGE,
                request_id=self._session_id,
                alternatives=[],
                recognition_usage=stt.RecognitionUsage(audio_duration=duration),
            )
        )

    def _process_stream_event(self, data: dict[str, Any]) -> None:
        # Streaming WebSocket response schema (Smallest AI Pulse API):
        # {
        #   "session_id":   str,
        #   "transcript":   str,        # partial or final text for this utterance
        #   "is_final":     bool,       # True when the utterance is complete
        #   "is_last":      bool,       # True when the session itself is done (after close_stream)
        #   "language":     str,        # present when is_final=True (detected or echoed)
        #   "words":        [           # present when word_timestamps=True
        #     {"word": str, "start": float, "end": float,
        #      "confidence": float, "speaker": int}  # speaker only when diarize=True
        #   ]
        # }
        session_id = data.get("session_id", "")
        if session_id:
            self._session_id = session_id

        transcript = data.get("transcript", "")
        is_final = data.get("is_final", False)

        if not transcript:
            return

        # Infer START_OF_SPEECH — the Pulse API does not emit a dedicated speech-start event.
        if not self._speaking:
            self._speaking = True
            self._event_ch.send_nowait(stt.SpeechEvent(type=stt.SpeechEventType.START_OF_SPEECH))

        alts = _transcript_to_speech_data(
            language=self._opts.language,
            data=data,
            start_time_offset=self.start_time_offset,
            diarize=self._opts.diarize,
        )

        if is_final:
            self._event_ch.send_nowait(
                stt.SpeechEvent(
                    type=stt.SpeechEventType.FINAL_TRANSCRIPT,
                    request_id=self._session_id,
                    alternatives=alts,
                )
            )
            if self._speaking:
                self._speaking = False
                self._event_ch.send_nowait(stt.SpeechEvent(type=stt.SpeechEventType.END_OF_SPEECH))
        else:
            self._event_ch.send_nowait(
                stt.SpeechEvent(
                    type=stt.SpeechEventType.INTERIM_TRANSCRIPT,
                    request_id=self._session_id,
                    alternatives=alts,
                )
            )


def _transcript_to_speech_data(
    language: str,
    data: dict[str, Any],
    *,
    start_time_offset: float,
    diarize: bool,
) -> list[stt.SpeechData]:
    transcript = data.get("transcript", "")
    raw_words: list[dict[str, Any]] = data.get("words") or []

    words: list[TimedString] | None = (
        [
            TimedString(
                text=w.get("word", ""),
                start_time=w.get("start", 0.0) + start_time_offset,
                end_time=w.get("end", 0.0) + start_time_offset,
            )
            for w in raw_words
        ]
        if raw_words
        else None
    )

    start_time = raw_words[0].get("start", 0.0) + start_time_offset if raw_words else 0.0
    end_time = raw_words[-1].get("end", 0.0) + start_time_offset if raw_words else 0.0

    # Streaming diarization: per-word speaker IDs are integers (0, 1, …).
    # Pick the most frequent speaker across the utterance for top-level speaker_id.
    speaker_id: str | None = None
    if diarize and raw_words:
        speaker_counts = Counter(w["speaker"] for w in raw_words if "speaker" in w)
        if speaker_counts:
            speaker_id = f"S{speaker_counts.most_common(1)[0][0]}"

    # When language="multi", the server echoes the detected language in is_final responses.
    detected_language = data.get("language", language) or language

    return [
        stt.SpeechData(
            language=LanguageCode(detected_language),
            text=transcript,
            start_time=start_time,
            end_time=end_time,
            confidence=raw_words[0].get("confidence", 0.0) if raw_words else 0.0,
            words=words,
            speaker_id=speaker_id,
        )
    ]


def _batch_transcription_to_speech_event(
    language: str,
    data: dict[str, Any],
) -> stt.SpeechEvent:
    # Batch HTTP response schema (Smallest AI Pulse API):
    # {
    #   "status":       str,
    #   "transcription": str,
    #   "audio_length": str,   # duration in seconds as a string
    #   "words":        [{"word": str, "start": float, "end": float,
    #                     "confidence": float, "speaker": str}],
    #   "language":     str,
    #   "metadata":     {"filename": str, "duration": float, "fileSize": int}
    # }
    transcript = data.get("transcription", "")
    raw_words: list[dict[str, Any]] = data.get("words") or []
    detected_language = data.get("language", language) or language

    words: list[TimedString] | None = (
        [
            TimedString(
                text=w.get("word", ""),
                start_time=w.get("start", 0.0),
                end_time=w.get("end", 0.0),
            )
            for w in raw_words
        ]
        if raw_words
        else None
    )

    start_time = raw_words[0].get("start", 0.0) if raw_words else 0.0
    end_time = raw_words[-1].get("end", 0.0) if raw_words else 0.0

    return stt.SpeechEvent(
        type=stt.SpeechEventType.FINAL_TRANSCRIPT,
        request_id=utils.shortuuid(),
        alternatives=[
            stt.SpeechData(
                language=LanguageCode(detected_language),
                text=transcript,
                start_time=start_time,
                end_time=end_time,
                confidence=raw_words[0].get("confidence", 0.0) if raw_words else 0.0,
                words=words,
            )
        ],
    )