Skip to content

fix(gradium): STT model delay and clean session end (1/2) #7870

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We鈥檒l occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
hhugo-gradium wants to merge 4 commits into
base: main
Choose a base branch
from
Draft

fix(gradium): STT model delay and clean session end (1/2) #7870

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
b0faaba
fix(gradium): read the model delay from ready's delay_in_frames
hhugo-gradium Oct 8, 2026
c865ab0
fix(gradium): end the STT session cleanly and surface server errors
hhugo-gradium Oct 8, 2026
b3f2f61
test(gradium): mark the STT tests as unit tests
hhugo-gradium Oct 16, 2026
fd770b6
fix(gradium): keep the STT transcript on every way a session can end
hhugo-gradium Oct 16, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
105 changes: 78 additions & 27 deletions livekit-plugins/livekit-plugins-gradium/livekit/plugins/gradium/stt.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@
from livekit.agents import (
DEFAULT_API_CONNECT_OPTIONS,
APIConnectOptions,
APIError,
APIStatusError,
LanguageCode,
stt,
Expand All @@ -51,6 +52,16 @@
SUPPORTED_SAMPLE_RATE = 24000


def _api_error(msg_data: dict[str, Any]) -> APIError:
"""Build the error for a `{"type": "error"}` frame.

The server sends it right before closing the socket, so the message is the only
place the actual failure reason (unsupported language, no credits, ...) shows up.
"""
message = msg_data.get("message") or "unknown error"
return APIError(f"Gradium returned error: {message}", body=msg_data)


@dataclass
class STTOptions:
sample_rate: int = SUPPORTED_SAMPLE_RATE
Expand Down Expand Up @@ -189,9 +200,6 @@ def update_options(


class SpeechStream(stt.SpeechStream):
# Used to close websocket
_CLOSE_MSG: str = json.dumps({"terminate_session": True})

def __init__(
self,
*,
Expand All @@ -216,9 +224,10 @@ def __init__(
self._ready_msg: dict[str, Any] | None = None

@property
def delay_in_tokens(self) -> int:
def delay_in_frames(self) -> int:
"""Steps the model lags behind the audio, from the server's ready message."""
if self._ready_msg is not None:
return int(self._ready_msg.get("delay_in_tokens", 6))
return int(self._ready_msg.get("delay_in_frames", 6))
return 6

@property
Expand Down Expand Up @@ -246,6 +255,7 @@ async def _run(self) -> None:
closing_ws = False

async def send_task(ws: aiohttp.ClientWebSocketResponse) -> None:
nonlocal closing_ws
samples_per_buffer = 1920

audio_bstream = utils.audio.AudioByteStream(
Expand All @@ -271,17 +281,50 @@ async def send_task(ws: aiohttp.ClientWebSocketResponse) -> None:
}
await ws.send_str(json.dumps(audio_msg))

# Input ended: tell the server, then let recv_task drain the model's last
# words until the server replies with its own end_of_stream.
closing_ws = True
try:
await ws.send_str(json.dumps({"type": "end_of_stream"}))
except (aiohttp.ClientError, ConnectionError):
# the server closed first; recv_task reports what it got
return

async def recv_task(ws: aiohttp.ClientWebSocketResponse) -> None:
nonlocal closing_ws
buffered_text: list[str] = []
speaking = False
remaining_vad_steps: int | None = None

def _finalize() -> None:
nonlocal speaking, remaining_vad_steps, buffered_text
if not speaking:
return
speaking = False
remaining_vad_steps = None
self._event_ch.send_nowait(
stt.SpeechEvent(
type=stt.SpeechEventType.FINAL_TRANSCRIPT,
alternatives=[
stt.SpeechData(
text=" ".join(buffered_text),
language=self._opts.language,
)
],
)
)
buffered_text = []
self._event_ch.send_nowait(stt.SpeechEvent(type=stt.SpeechEventType.END_OF_SPEECH))

while True:
try:
msg = await asyncio.wait_for(ws.receive(), timeout=5)
except asyncio.TimeoutError:
if closing_ws:
break
# the server did not answer the end_of_stream in time: keep
# what the model already produced
_finalize()
return
continue

if msg.type in (
Expand All @@ -290,6 +333,8 @@ async def recv_task(ws: aiohttp.ClientWebSocketResponse) -> None:
aiohttp.WSMsgType.CLOSING,
):
if closing_ws:
# a close instead of the server's end_of_stream: same thing
_finalize()
return
raise APIStatusError(
"Gradium connection closed unexpectedly",
Expand All @@ -303,9 +348,30 @@ async def recv_task(ws: aiohttp.ClientWebSocketResponse) -> None:

try:
data = json.loads(msg.data)
except json.JSONDecodeError:
logger.exception("Failed to decode message from Gradium")
continue

type_ = data.get("type", "")
type_ = data.get("type", "")
if type_ == "error":
raise _api_error(data)
if type_ == "end_of_stream":
# The model has consumed all the audio: whatever is still buffered
# is the end of the last utterance.
_finalize()
if closing_ws:
return
# The server ended the session while audio is still coming (a
# duration limit, for instance) and will close the socket; raise
# so the base class reconnects instead of the next audio write
# failing with a raw connection error.
raise APIStatusError(
"Gradium ended the session while input was still open",
status_code=ws.close_code or -1,
body=data,
)

try:
if type_ == "text":
if speaking is False:
speaking = True
Expand Down Expand Up @@ -334,35 +400,20 @@ async def recv_task(ws: aiohttp.ClientWebSocketResponse) -> None:
)
if positive_vad:
if remaining_vad_steps is None:
remaining_vad_steps = self.delay_in_tokens
remaining_vad_steps = self.delay_in_frames
if self._opts.vad_flush:
samples_per_channel = self.frame_size * self.delay_in_tokens
samples_per_channel = self.frame_size * self.delay_in_frames
zeros = AudioFrame.create(
sample_rate=self._opts.sample_rate,
num_channels=1,
samples_per_channel=samples_per_channel,
)
await self._input_ch.send(zeros)
if not self._input_ch.closed:
await self._input_ch.send(zeros)
else:
remaining_vad_steps -= 1
if remaining_vad_steps <= 0:
speaking = False
remaining_vad_steps = None
event = stt.SpeechEvent(
type=stt.SpeechEventType.FINAL_TRANSCRIPT,
alternatives=[
stt.SpeechData(
text=" ".join(buffered_text),
language=self._opts.language,
)
],
)
self._event_ch.send_nowait(event)

buffered_text = []
self._event_ch.send_nowait(
stt.SpeechEvent(type=stt.SpeechEventType.END_OF_SPEECH)
)
_finalize()
else:
remaining_vad_steps = None

Expand Down
Loading
Loading