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

478 lines
16 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 Gnani Vachana
This module provides an STT implementation that uses the Gnani Vachana API,
supporting both REST recognition and real-time streaming (WebSocket).
"""
from __future__ import annotations
import asyncio
import json
import os
from dataclasses import dataclass, replace
from typing import Any, Literal
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
from livekit.agents.utils import AudioBuffer
from livekit.agents.utils.misc import is_given
from .log import logger
GnaniSTTFormat = Literal["verbatim", "transcribe"]
GNANI_STT_BASE_URL = "https://api.vachana.ai"
GnaniSTTLanguages = Literal[
"bn-IN",
"en-IN",
"gu-IN",
"hi-IN",
"kn-IN",
"ml-IN",
"mr-IN",
"pa-IN",
"ta-IN",
"te-IN",
"en-IN,hi-IN",
]
SUPPORTED_LANGUAGES: set[str] = {
"bn-IN",
"en-IN",
"gu-IN",
"hi-IN",
"kn-IN",
"ml-IN",
"mr-IN",
"pa-IN",
"ta-IN",
"te-IN",
"en-IN,hi-IN",
}
STREAM_SUPPORTED_LANGUAGES: set[str] = {
"bn-IN",
"en-IN",
"gu-IN",
"hi-IN",
"kn-IN",
"ml-IN",
"mr-IN",
"pa-IN",
"ta-IN",
"te-IN",
}
SAMPLE_RATE_16K = 16000
SAMPLE_RATE_8K = 8000
STREAM_CHUNK_BYTES = 1024
@dataclass
class GnaniSTTOptions:
api_key: str
language: str
sample_rate: int = SAMPLE_RATE_16K
base_url: str = GNANI_STT_BASE_URL
preferred_language: str | None = None
format: str = "verbatim"
itn_native_numerals: bool = False
_DEPRECATED_STT_KWARGS = frozenset(("organization_id", "user_id", "http_session"))
def _check_deprecated_args(kwargs: dict[str, Any], *, caller: str = "STT.__init__") -> None:
"""Warn about deprecated kwargs and raise on truly unknown ones."""
for name in _DEPRECATED_STT_KWARGS:
if name in kwargs:
logger.warning(f"`{name}` is deprecated and no longer used")
unknown = set(kwargs) - _DEPRECATED_STT_KWARGS
if unknown:
raise TypeError(
f"{caller}() got unexpected keyword argument(s): {', '.join(sorted(unknown))}"
)
class STT(stt.STT):
"""Gnani Vachana Speech-to-Text implementation.
Provides speech-to-text functionality using Gnani's Vachana platform.
Supports REST recognition and real-time streaming via WebSocket.
Args:
language: BCP-47 language code (e.g. "hi-IN", "en-IN").
api_key: Gnani API key (falls back to GNANI_API_KEY env var).
sample_rate: Audio sample rate for streaming (8000 or 16000).
base_url: Vachana API base URL.
preferred_language: Force single-language model for this code.
format: "verbatim" (default) or "transcribe" (enables ITN).
itn_native_numerals: Render digits in native script when format="transcribe".
"""
def __init__(
self,
*,
language: str = "en-IN",
api_key: str | None = None,
sample_rate: int = SAMPLE_RATE_16K,
base_url: str = GNANI_STT_BASE_URL,
preferred_language: str | None = None,
format: GnaniSTTFormat = "verbatim",
itn_native_numerals: bool = False,
**kwargs: Any,
) -> None:
super().__init__(
capabilities=stt.STTCapabilities(
streaming=True,
interim_results=False,
aligned_transcript=False,
)
)
_check_deprecated_args(kwargs)
self._api_key = api_key or os.environ.get("GNANI_API_KEY")
if not self._api_key:
raise ValueError(
"Gnani API key is required. "
"Provide it directly or set GNANI_API_KEY environment variable."
)
if sample_rate not in (SAMPLE_RATE_8K, SAMPLE_RATE_16K):
raise ValueError("sample_rate must be 8000 or 16000")
self._opts = GnaniSTTOptions(
api_key=self._api_key,
language=language,
sample_rate=sample_rate,
base_url=base_url,
preferred_language=preferred_language,
format=format,
itn_native_numerals=itn_native_numerals,
)
self._session: aiohttp.ClientSession | None = None
@property
def model(self) -> str:
return "vachana-stt-v3"
@property
def provider(self) -> str:
return "Gnani"
def _ensure_session(self) -> aiohttp.ClientSession:
if not self._session:
self._session = utils.http_context.http_session()
return self._session
@staticmethod
def _single_attempt(conn_options: APIConnectOptions) -> APIConnectOptions:
return APIConnectOptions(
max_retry=0,
retry_interval=conn_options.retry_interval,
timeout=conn_options.timeout,
)
async def recognize(
self,
buffer: AudioBuffer,
*,
language: NotGivenOr[str] = NOT_GIVEN,
conn_options: APIConnectOptions = DEFAULT_API_CONNECT_OPTIONS,
) -> stt.SpeechEvent:
return await super().recognize(
buffer,
language=language,
conn_options=self._single_attempt(conn_options),
)
async def _recognize_impl(
self,
buffer: AudioBuffer,
*,
language: NotGivenOr[str] = NOT_GIVEN,
conn_options: APIConnectOptions = DEFAULT_API_CONNECT_OPTIONS,
) -> stt.SpeechEvent:
lang = language if is_given(language) else self._opts.language
wav_bytes = rtc.combine_audio_frames(buffer).to_wav_bytes()
form_data = aiohttp.FormData()
form_data.add_field("audio_file", wav_bytes, filename="audio.wav", content_type="audio/wav")
form_data.add_field("language_code", lang)
form_data.add_field("format", self._opts.format)
if self._opts.preferred_language is not None:
form_data.add_field("preferred_language", self._opts.preferred_language)
if self._opts.itn_native_numerals:
form_data.add_field("itn_native_numerals", "true")
headers: dict[str, str] = {
"X-API-Key-ID": self._opts.api_key,
}
try:
async with self._ensure_session().post(
url=f"{self._opts.base_url}/stt/v3",
data=form_data,
headers=headers,
timeout=aiohttp.ClientTimeout(
total=conn_options.timeout,
sock_connect=conn_options.timeout,
),
) as res:
if res.status != 200:
error_text = await res.text()
logger.error(f"Gnani STT API error: {res.status} - {error_text}")
raise APIStatusError(
message=f"Gnani STT API Error ({res.status}): {error_text}",
status_code=res.status,
body=error_text,
)
response_json = await res.json()
transcript = response_json.get("transcript", "")
request_id = response_json.get("request_id", "")
return stt.SpeechEvent(
type=stt.SpeechEventType.FINAL_TRANSCRIPT,
request_id=request_id,
alternatives=[
stt.SpeechData(
language=LanguageCode(lang),
text=transcript,
confidence=1.0,
)
],
)
except asyncio.TimeoutError as e:
raise APITimeoutError("Gnani STT API request timed out") from e
except (APIStatusError, APIConnectionError, APITimeoutError):
raise
except Exception as e:
raise APIConnectionError(f"Gnani STT error: {e}") from e
def stream(
self,
*,
language: NotGivenOr[str] = NOT_GIVEN,
conn_options: APIConnectOptions = DEFAULT_API_CONNECT_OPTIONS,
) -> SpeechStream:
opts = replace(self._opts)
if is_given(language):
opts.language = language
return SpeechStream(
stt=self,
opts=opts,
conn_options=self._single_attempt(conn_options),
)
async def aclose(self) -> None:
pass
class SpeechStream(stt.RecognizeStream):
"""WebSocket-based streaming STT for Gnani Vachana.
Connects to wss://api.vachana.ai/stt/v3/stream and sends raw PCM audio
in 1024-byte chunks (512 samples, 16-bit mono).
"""
def __init__(
self,
*,
stt: STT,
opts: GnaniSTTOptions,
conn_options: APIConnectOptions,
) -> None:
super().__init__(
stt=stt,
conn_options=conn_options,
sample_rate=opts.sample_rate,
)
self._opts = opts
def _build_ws_url(self) -> str:
base = self._opts.base_url
if base.startswith("https://"):
ws_base = "wss://" + base[len("https://") :]
elif base.startswith("http://"):
ws_base = "ws://" + base[len("http://") :]
else:
ws_base = "wss://" + base
return f"{ws_base}/stt/v3/stream"
async def _run(self) -> None:
import websockets
ws_url = self._build_ws_url()
headers: dict[str, str] = {
"x-api-key-id": self._opts.api_key,
"lang_code": self._opts.language,
"x-sample-rate": str(self._opts.sample_rate),
}
if self._opts.format != "verbatim":
headers["x-format"] = self._opts.format
if self._opts.preferred_language is not None:
headers["preferred_language"] = self._opts.preferred_language
if self._opts.itn_native_numerals:
headers["itn_native_numerals"] = "true"
try:
async with websockets.connect(
ws_url,
additional_headers=headers,
ping_interval=20,
ping_timeout=20,
close_timeout=10,
) as ws:
connected_msg = await asyncio.wait_for(ws.recv(), timeout=10)
connected_data = json.loads(connected_msg)
if connected_data.get("type") != "connected":
logger.warning(f"Unexpected first message from Gnani STT: {connected_data}")
send_task = asyncio.create_task(self._send_audio(ws), name="gnani-stt-send")
recv_task = asyncio.create_task(self._recv_messages(ws), name="gnani-stt-recv")
try:
# Wait for send to finish; if recv errors first, propagate it.
done, _ = await asyncio.wait(
[send_task, recv_task],
return_when=asyncio.FIRST_COMPLETED,
)
for task in done:
task.result()
if send_task.done() and not recv_task.done():
# All audio sent. The Gnani API has no application-level
# end-of-stream message, so give the server a short
# window to flush final transcripts before closing.
try:
await asyncio.wait_for(asyncio.shield(recv_task), timeout=1.0)
except asyncio.TimeoutError:
pass
finally:
await utils.aio.gracefully_cancel(send_task, recv_task)
except websockets.exceptions.ConnectionClosed as e:
raise APIConnectionError(f"Gnani STT WebSocket closed unexpectedly: {e}") from e
except asyncio.TimeoutError as e:
raise APITimeoutError("Gnani STT WebSocket connection timed out") from e
except (APIConnectionError, APIStatusError, APITimeoutError):
raise
except Exception as e:
raise APIConnectionError(f"Gnani STT WebSocket error: {e}") from e
async def _send_audio(self, ws: Any) -> None:
audio_buffer = bytearray()
async for data in self._input_ch:
if isinstance(data, self._FlushSentinel):
if audio_buffer:
await ws.send(bytes(audio_buffer))
audio_buffer.clear()
continue
frame: rtc.AudioFrame = data
raw_pcm = frame.data.tobytes()
audio_buffer.extend(raw_pcm)
while len(audio_buffer) >= STREAM_CHUNK_BYTES:
chunk = bytes(audio_buffer[:STREAM_CHUNK_BYTES])
audio_buffer = audio_buffer[STREAM_CHUNK_BYTES:]
await ws.send(chunk)
if audio_buffer:
await ws.send(bytes(audio_buffer))
async def _recv_messages(self, ws: Any) -> None:
try:
async for msg in ws:
if isinstance(msg, bytes):
continue
data = json.loads(msg)
msg_type = data.get("type", "")
if msg_type == "transcript":
text = data.get("text", "")
if not text:
continue
self._event_ch.send_nowait(
stt.SpeechEvent(
type=stt.SpeechEventType.FINAL_TRANSCRIPT,
request_id=data.get("segment_id", ""),
alternatives=[
stt.SpeechData(
language=LanguageCode(self._opts.language),
text=text,
confidence=1.0,
)
],
)
)
elif msg_type in ("speech_start", "vad_start"):
self._event_ch.send_nowait(
stt.SpeechEvent(
type=stt.SpeechEventType.START_OF_SPEECH,
)
)
elif msg_type in ("speech_end", "vad_end"):
self._event_ch.send_nowait(
stt.SpeechEvent(
type=stt.SpeechEventType.END_OF_SPEECH,
)
)
elif msg_type == "processing":
pass
elif msg_type == "error":
error_msg = data.get("message", "Unknown error")
logger.error(f"Gnani STT stream error: {error_msg}")
raise APIStatusError(
message=f"Gnani STT stream error: {error_msg}",
status_code=500,
body=error_msg,
)
except asyncio.CancelledError:
raise
except (APIStatusError, APIConnectionError, APITimeoutError):
raise
except Exception as e:
raise APIConnectionError(f"Error receiving Gnani STT messages: {e}") from e