Files
2026-07-13 13:39:38 +08:00

505 lines
19 KiB
Python

# 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.
"""Speech-to-Text implementation for SimpliSmart
This module provides an STT implementation that uses the SimpliSmart API.
"""
import asyncio
import base64
import contextlib
import json
import os
import traceback
import weakref
from typing import Any, Literal
from urllib.parse import urlparse
import aiohttp
from pydantic import BaseModel
from livekit import rtc
from livekit.agents import (
APIConnectionError,
APIConnectOptions,
APIStatusError,
APITimeoutError,
LanguageCode,
stt,
utils,
)
from livekit.agents.types import DEFAULT_API_CONNECT_OPTIONS, NOT_GIVEN, NotGivenOr
from livekit.agents.utils import AudioBuffer
from livekit.agents.utils.misc import is_given
from .log import logger
from .models import STTModels
SIMPLISMART_BASE_URL = "https://api.simplismart.live/predict"
class SimplismartSTTOptions(BaseModel):
language: LanguageCode | None = None
task: Literal["transcribe", "translate"] = "transcribe"
without_timestamps: bool = True
vad_model: Literal["silero", "frame"] = "frame"
vad_filter: bool = True
vad_onset: float | None = 0.5
vad_offset: float | None = None
min_speech_duration_ms: int = 0
max_speech_duration_s: float = 30
min_silence_duration_ms: int = 2000
speech_pad_ms: int = 400
initial_prompt: str | None = None
hotwords: str | None = None
num_speakers: int = 0
compression_ratio_threshold: float | None = 2.4
beam_size: int = 4
temperature: float = 0.0
multilingual: bool = False
max_tokens: float | None = 400
log_prob_threshold: float | None = -1.0
length_penalty: int = 1
repetition_penalty: float = 1.01
strict_hallucination_reduction: bool = False
class STT(stt.STT):
def __init__(
self,
*,
base_url: str = SIMPLISMART_BASE_URL,
api_key: str | None = None,
streaming: bool = False,
model: STTModels | str = "openai/whisper-large-v3-turbo",
language: str = "en",
task: Literal["transcribe", "translate"] = "transcribe",
without_timestamps: bool = True,
vad_model: Literal["silero", "frame"] = "frame",
vad_filter: bool = True,
vad_onset: float | None = 0.5,
vad_offset: float | None = None,
min_speech_duration_ms: int = 0,
max_speech_duration_s: float = 30,
min_silence_duration_ms: int = 2000,
speech_pad_ms: int = 400,
initial_prompt: str | None = None,
hotwords: str | None = None,
num_speakers: int = 0,
compression_ratio_threshold: float | None = 2.4,
beam_size: int = 4,
temperature: float = 0.0,
multilingual: bool = False,
max_tokens: float | None = 400,
log_prob_threshold: float | None = -1.0,
length_penalty: int = 1,
repetition_penalty: float = 1.01,
strict_hallucination_reduction: bool = False,
http_session: aiohttp.ClientSession | None = None,
):
"""
Configuration options for the SimpliSmart STT (Speech-to-Text) engine.
Note:
Streaming transcription is not publicly available at this time.
Args:
language (str): Language code for transcription (default: "en").
task (Literal["transcribe", "translate"]): Operation to perform, either "transcribe" or "translate".
model (STTModels | str): Model identifier for the backend STT model.
without_timestamps (bool): If True, disables timestamp generation in transcripts.
vad_model (Literal["silero", "frame"]): Voice Activity Detection model to use ("silero" or "frame").
vad_filter (bool): Whether to apply VAD to filter input audio.
vad_onset (float | None): Time (in seconds) for VAD onset boundary.
vad_offset (float | None): Time (in seconds) for VAD offset boundary.
min_speech_duration_ms (int): Minimum duration (ms) for a valid speech segment.
max_speech_duration_s (float): Maximum speech segment duration (seconds).
min_silence_duration_ms (int): Minimum silence duration (ms) to split speech.
speech_pad_ms (int): Padding (ms) added to boundaries of detected speech.
initial_prompt (str | None): Optional initial prompt for contextual biasing.
hotwords (str | None): Comma-separated list of hotwords to bias recognition.
num_speakers (int): Number of speakers for diarization.
compression_ratio_threshold (float | None): Threshold for output compression ratio.
beam_size (int): Beam size for the decoder.
temperature (float): Decoding temperature (affects randomness).
multilingual (bool): Whether to permit multilingual recognition.
max_tokens (float | None): Maximum number of output tokens for the model.
log_prob_threshold (float | None): Log probability threshold for word filtering.
length_penalty (int): Penalty for longer transcriptions.
repetition_penalty (float): Penalty for repeated words during decoding.
strict_hallucination_reduction (bool): Whether to apply hallucination reduction.
"""
if streaming:
base_url = f"wss://{urlparse(base_url).netloc}/ws/audio"
super().__init__(
capabilities=stt.STTCapabilities(
streaming=streaming,
interim_results=False,
aligned_transcript="word",
)
)
self._api_key = api_key or os.environ.get("SIMPLISMART_API_KEY")
if not self._api_key:
raise ValueError("SIMPLISMART_API_KEY is not set")
self._model = model
self._opts = SimplismartSTTOptions(
language=LanguageCode(language),
task=task,
without_timestamps=without_timestamps,
vad_model=vad_model,
vad_filter=vad_filter,
vad_onset=vad_onset,
vad_offset=vad_offset,
min_speech_duration_ms=min_speech_duration_ms,
max_speech_duration_s=max_speech_duration_s,
min_silence_duration_ms=min_silence_duration_ms,
speech_pad_ms=speech_pad_ms,
initial_prompt=initial_prompt,
hotwords=hotwords,
num_speakers=num_speakers,
compression_ratio_threshold=compression_ratio_threshold,
beam_size=beam_size,
temperature=temperature,
multilingual=multilingual,
max_tokens=max_tokens,
log_prob_threshold=log_prob_threshold,
length_penalty=length_penalty,
repetition_penalty=repetition_penalty,
strict_hallucination_reduction=strict_hallucination_reduction,
)
self._base_url = base_url
self._session = http_session
self._streams = weakref.WeakSet[SpeechStream]()
@property
def provider(self) -> str:
return "Simplismart"
@property
def model(self) -> str:
return self._model
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:
resolved_language: str | None = language if is_given(language) else self._opts.language
wav_bytes = rtc.combine_audio_frames(buffer).to_wav_bytes()
audio_b64 = base64.b64encode(wav_bytes).decode("utf-8")
payload = self._opts.model_dump()
payload["audio_data"] = audio_b64
payload["language"] = resolved_language
payload["model"] = self._model
try:
async with self._ensure_session().post(
self._base_url,
json=payload,
headers={
"Authorization": f"Bearer {self._api_key}",
"Content-Type": "application/json",
},
timeout=aiohttp.ClientTimeout(
total=conn_options.timeout,
),
) as res:
if res.status != 200:
error_text = await res.text()
logger.error(f"Simplismart API error: {res.status} - {error_text}")
raise APIStatusError(
message=f"Simplismart API Error: {error_text}",
status_code=res.status,
request_id=None,
body=error_text,
)
response_json = await res.json()
timestamps = response_json.get("timestamps", [])
transcription = response_json.get("transcription", [])
info = response_json.get("info", {})
detected_language = LanguageCode(info.get("language", resolved_language or "en"))
start_time = timestamps[0][0] if timestamps else 0.0
end_time = timestamps[-1][1] if timestamps else 0.0
request_id = response_json.get("request_id", "")
text = "".join(transcription)
alternatives = [
stt.SpeechData(
language=detected_language,
text=text,
start_time=start_time,
end_time=end_time,
),
]
return stt.SpeechEvent(
type=stt.SpeechEventType.FINAL_TRANSCRIPT,
request_id=request_id,
alternatives=alternatives,
)
except asyncio.TimeoutError as e:
logger.error(f"Simplismart API timeout: {e}")
raise APITimeoutError("Simplismart API request timed out") from e
except aiohttp.ClientError as e:
logger.error(f"Simplismart API client error: {e}")
raise APIConnectionError(f"Simplismart API connection error: {e}") from e
except APIStatusError:
raise
except Exception as e:
logger.error(f"Error during Simplismart STT processing: {traceback.format_exc()}")
raise APIConnectionError(f"Unexpected error in Simplismart STT: {e}") from e
def stream(
self,
*,
language: NotGivenOr[str] = NOT_GIVEN,
conn_options: APIConnectOptions = DEFAULT_API_CONNECT_OPTIONS,
**kwargs: Any,
) -> "SpeechStream":
"""Create a streaming transcription session."""
opts_language = LanguageCode(language) if is_given(language) else self._opts.language
# Create options for the stream
stream_opts = SimplismartSTTOptions(language=opts_language)
# Create a fresh session for this stream to avoid conflicts
stream_session = aiohttp.ClientSession()
if self._api_key is None:
raise ValueError("API key cannot be None")
stream = SpeechStream(
stt=self,
opts=stream_opts,
conn_options=conn_options,
api_key=self._api_key,
http_session=stream_session,
)
self._streams.add(stream)
return stream
class SpeechStream(stt.SpeechStream):
"""Simplismart streaming speech-to-text implementation."""
_CHUNK_DURATION_MS = 50
_SAMPLE_RATE = 16000
def __init__(
self,
*,
stt: STT,
opts: SimplismartSTTOptions,
conn_options: APIConnectOptions,
api_key: str,
http_session: aiohttp.ClientSession,
) -> None:
self._opts = opts
super().__init__(stt=stt, conn_options=conn_options, sample_rate=self._SAMPLE_RATE)
self._api_key = api_key
self._session = http_session
self._reconnect_event = asyncio.Event()
self._request_id = str(id(self))
self.ws_url = stt._base_url
async def _run(self) -> None:
@utils.log_exceptions(logger=logger)
async def send_task(ws: aiohttp.ClientWebSocketResponse) -> None:
# forward audio to simplismart in chunks of 50ms
samples_50ms = self._SAMPLE_RATE // 20
audio_bstream = utils.audio.AudioByteStream(
sample_rate=self._SAMPLE_RATE,
num_channels=1,
samples_per_channel=samples_50ms,
)
async for data in self._input_ch:
frames: list[rtc.AudioFrame] = []
if isinstance(data, rtc.AudioFrame):
frames.extend(audio_bstream.write(data.data.tobytes()))
elif isinstance(data, self._FlushSentinel):
frames.extend(audio_bstream.flush())
for frame in frames:
await ws.send_bytes(frame.data.tobytes())
@utils.log_exceptions(logger=logger)
async def recv_task(ws: aiohttp.ClientWebSocketResponse) -> None:
while True:
msg = await ws.receive()
if msg.type in (
aiohttp.WSMsgType.CLOSED,
aiohttp.WSMsgType.CLOSE,
aiohttp.WSMsgType.CLOSING,
):
# close is expected, see SpeechStream.aclose
# or when the agent session ends, the http session is closed
if self._session.closed:
return
# this will trigger a reconnection, see the _run loop
self._reconnect_event.set()
return
if msg.type != aiohttp.WSMsgType.BINARY:
logger.warning("unexpected simplismart message type %s", msg.type)
continue
try:
self._handle_transcript_data(msg.data.decode("utf-8"))
except Exception:
logger.exception("failed to process simplismart message")
ws: aiohttp.ClientWebSocketResponse | None = None
while True:
try:
ws = await self._connect_ws()
await self._send_initial_config(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,
)
# propagate exceptions from completed tasks
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()
with contextlib.suppress(asyncio.CancelledError):
tasks_group.exception() # retrieve the exception
finally:
if ws is not None:
await ws.close()
async def _connect_ws(self) -> aiohttp.ClientWebSocketResponse:
try:
ws = await asyncio.wait_for(
self._session.ws_connect(
self.ws_url,
headers={"Authorization": f"Bearer {self._api_key}"},
),
self._conn_options.timeout,
)
except (aiohttp.ClientConnectorError, asyncio.TimeoutError) as e:
raise APIConnectionError("failed to connect to simplismart") from e
return ws
async def aclose(self) -> None:
await super().aclose()
if self._session and not self._session.closed:
await self._session.close()
async def _send_initial_config(self, ws: aiohttp.ClientWebSocketResponse) -> None:
"""Send initial configuration message with language for Simplismart models."""
try:
config_message = {"language": self._opts.language}
await ws.send_json(config_message)
logger.info(
"Sent initial config for Simplismart model",
extra={"request_id": self._request_id, "language": self._opts.language},
)
except Exception as e:
logger.error(
f"Failed to send initial configuration: {e}",
extra={"request_id": self._request_id},
exc_info=True,
)
raise APIConnectionError(f"Failed to send initial config: {e}") from e
def _handle_transcript_data(self, data: str) -> None:
"""Handle transcription result messages."""
transcript_text = data
request_id = self._request_id
try:
# Create usage event with proper metrics extraction
metrics: dict[str, float] = {}
request_data = {
"original_id": request_id,
"processing_latency": metrics.get("processing_latency", 0.0),
}
usage_event = stt.SpeechEvent(
type=stt.SpeechEventType.RECOGNITION_USAGE,
request_id=json.dumps(request_data),
recognition_usage=stt.RecognitionUsage(
audio_duration=metrics.get("audio_duration", 0.0),
),
)
self._event_ch.send_nowait(usage_event)
# Create speech data
speech_data = stt.SpeechData(
language=LanguageCode(self._opts.language or "en"),
text=transcript_text,
)
# Create final transcript event with request_id
speech_event = stt.SpeechEvent(
type=stt.SpeechEventType.FINAL_TRANSCRIPT,
request_id=request_id,
alternatives=[speech_data],
)
self._event_ch.send_nowait(speech_event)
logger.debug(
"Transcript processed successfully",
extra={
"request_id": self._request_id,
"text_length": len(transcript_text),
"language": self._opts.language,
"confidence": speech_data.confidence,
},
)
except Exception as e:
logger.error(
f"Error processing transcript data: {e}",
extra={
"request_id": self._request_id,
"transcript_text": transcript_text,
},
exc_info=True,
)
raise