gsadmin/meet
1
0
Fork 0
mirror of https://github.com/suitenumerique/meet.git synced 2026-07-26 11:58:53 +00:00
Code Issues Packages Projects Releases Wiki Activity

wip - single websockets

Browse Source
This commit is contained in:
leo
2026-05-28 16:31:14 +02:00
parent e0f9460606
commit 64bf50116b
1 changed files with 170 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/agents/voxtral_vllm_stt.py
+170 -44
View File
@@ -4,6 +4,8 @@ vLLM exposes Voxtral Realtime over a WebSocket that follows the OpenAI Realtime
API protocol (not Mistral's proprietary realtime protocol).
"""
from __future__ import annotations
import asyncio
import base64
import json
@@ -11,7 +13,7 @@ import logging
import os
import weakref
from collections import deque
from dataclasses import dataclass
from dataclasses import dataclass, field
import websockets
from livekit.agents import (
@@ -35,6 +37,12 @@ NUM_CHANNELS = 1
CHUNK_SAMPLES = 1600 # 100 ms @ 16 kHz mono
PREROLL_CHUNKS = 5 # keep 500 ms of audio before START_OF_SPEECH
# Reconnect policy: exponential backoff capped at MAX, give up after MAX_ATTEMPTS
# consecutive failures (a successful handshake resets the counter).
RECONNECT_BACKOFF_BASE_S = 0.5
RECONNECT_BACKOFF_MAX_S = 8.0
RECONNECT_MAX_ATTEMPTS = 5
@dataclass
class _STTOptions:
@@ -43,6 +51,22 @@ class _STTOptions:
api_key: str | None
@dataclass
class _PendingUtterance:
"""An utterance in flight on the shared websocket.
`sent_chunks` holds every chunk we have already enqueued for send on this
or a prior connection; on reconnect we replay them before resuming reads
from `queue`. vLLM concatenates `input_audio_buffer.append` events into a
single audio buffer per generation, so duplicates from a partial prior send
are harmless.
"""
queue: asyncio.Queue[bytes | None]
sent_chunks: list[bytes] = field(default_factory=list)
ended: bool = False # the None sentinel has been drained from `queue`
class STT(stt.STT):
"""LiveKit STT speaking the OpenAI Realtime protocol served by vLLM."""
@@ -146,11 +170,16 @@ class SpeechStream(stt.RecognizeStream):
super().__init__(stt=stt, conn_options=conn_options, sample_rate=SAMPLE_RATE)
self._opts = opts
self._vad = vad_instance
self._audio_q: asyncio.Queue[bytes | None] = asyncio.Queue()
self._utterance_q: asyncio.Queue[bytes | None] | None = None
self._speaking = False
self._preroll: deque[bytes] = deque(maxlen=PREROLL_CHUNKS)
self._idle = asyncio.Event()
self._idle.set()
# Voxtral realtime is strictly sequential: only one generation runs at a
# time, and a new `commit` is ignored while the previous one is still
# producing. We queue per-utterance audio buffers here and let the
# pipeline process them one by one on the shared websocket.
self._utterance_chan: asyncio.Queue[asyncio.Queue[bytes | None] | None] = (
asyncio.Queue()
)
@utils.log_exceptions(logger=logger)
async def _run(self) -> None:
@@ -182,51 +211,42 @@ class SpeechStream(stt.RecognizeStream):
elif ev.type == vad_module.VADEventType.END_OF_SPEECH:
self._on_end_of_speech()
headers: dict[str, str] = {}
if self._opts.api_key:
headers["Authorization"] = f"Bearer {self._opts.api_key}"
pipeline_t = asyncio.create_task(self._utterance_pipeline())
try:
async with websockets.connect(
self._opts.base_url,
additional_headers=headers,
open_timeout=self._conn_options.timeout,
) as ws:
request_id = await self._handshake(ws)
send_t = asyncio.create_task(self._send_audio(ws, self._audio_q))
recv_t = asyncio.create_task(self._receive_events(ws, request_id))
try:
await asyncio.gather(input_task(), vad_task())
# wait for any in-flight utterance to finish transcribing
await self._idle.wait()
finally:
for t in (send_t, recv_t):
if not t.done():
t.cancel()
for t in (send_t, recv_t):
try:
await t
except (asyncio.CancelledError, Exception): # noqa: BLE001
pass
await vad_stream.aclose()
await asyncio.gather(input_task(), vad_task())
# signal end-of-stream; pipeline finishes pending utterances first
self._utterance_chan.put_nowait(None)
await pipeline_t
except (APIStatusError, APIConnectionError, asyncio.CancelledError):
raise
except Exception as exc:
logger.exception("vLLM realtime stream failed")
raise APIConnectionError() from exc
finally:
if not pipeline_t.done():
pipeline_t.cancel()
try:
await pipeline_t
except asyncio.CancelledError:
pass
except Exception:
logger.exception("utterance pipeline failed during finalize")
await vad_stream.aclose()
def _handle_chunk(self, chunk: bytes) -> None:
self._preroll.append(chunk)
if self._speaking:
self._audio_q.put_nowait(chunk)
if self._speaking and self._utterance_q is not None:
self._utterance_q.put_nowait(chunk)
def _on_start_of_speech(self) -> None:
if self._speaking:
return
self._speaking = True
self._idle.clear()
q: asyncio.Queue[bytes | None] = asyncio.Queue()
for chunk in self._preroll:
self._audio_q.put_nowait(chunk)
q.put_nowait(chunk)
self._utterance_q = q
self._utterance_chan.put_nowait(q)
self._event_ch.send_nowait(
stt.SpeechEvent(type=stt.SpeechEventType.START_OF_SPEECH)
)
@@ -235,7 +255,9 @@ class SpeechStream(stt.RecognizeStream):
if not self._speaking:
return
self._speaking = False
self._audio_q.put_nowait(None)
if self._utterance_q is not None:
self._utterance_q.put_nowait(None)
self._utterance_q = None
self._event_ch.send_nowait(
stt.SpeechEvent(type=stt.SpeechEventType.END_OF_SPEECH)
)
@@ -249,20 +271,122 @@ class SpeechStream(stt.RecognizeStream):
body=created,
)
await ws.send(json.dumps({"type": "session.update", "model": self._opts.model}))
await ws.send(json.dumps({"type": "input_audio_buffer.commit"}))
return created.get("id", "")
def _auth_headers(self) -> dict[str, str]:
if self._opts.api_key:
return {"Authorization": f"Bearer {self._opts.api_key}"}
return {}
async def _utterance_pipeline(self) -> None:
# Owns the websocket lifecycle. On drop, reopens and resumes the
# in-flight utterance (if any) by replaying its already-sent chunks.
pending: _PendingUtterance | None = None
attempt = 0
while True:
try:
async with websockets.connect(
self._opts.base_url,
additional_headers=self._auth_headers(),
open_timeout=self._conn_options.timeout,
) as ws:
request_id = await self._handshake(ws)
attempt = 0
while True:
if pending is None:
q = await self._utterance_chan.get()
if q is None:
return
pending = _PendingUtterance(queue=q)
await self._process_utterance(ws, pending, request_id)
pending = None
except (websockets.WebSocketException, OSError, TimeoutError) as exc:
attempt += 1
if attempt > RECONNECT_MAX_ATTEMPTS:
logger.exception(
"vLLM realtime: giving up after %d reconnect attempts",
RECONNECT_MAX_ATTEMPTS,
)
raise APIConnectionError() from exc
backoff = min(
RECONNECT_BACKOFF_BASE_S * (2 ** (attempt - 1)),
RECONNECT_BACKOFF_MAX_S,
)
if pending is None:
logger.warning(
"vLLM WS connection lost between utterances "
"(attempt %d/%d): %s; retrying in %.1fs",
attempt,
RECONNECT_MAX_ATTEMPTS,
exc,
backoff,
)
else:
logger.warning(
"vLLM WS dropped mid-utterance (%d chunks buffered, "
"ended=%s, attempt %d/%d): %s; retrying in %.1fs",
len(pending.sent_chunks),
pending.ended,
attempt,
RECONNECT_MAX_ATTEMPTS,
exc,
backoff,
)
await asyncio.sleep(backoff)
async def _process_utterance(
self,
ws: websockets.ClientConnection,
pending: _PendingUtterance,
request_id: str,
) -> None:
# Start a fresh generation. Safe to send here: the previous utterance's
# transcription.done has already been received (we await it below), so
# the server-side generation_task is done and won't ignore this commit.
await ws.send(json.dumps({"type": "input_audio_buffer.commit"}))
send_t = asyncio.create_task(self._send_audio(ws, pending))
try:
await self._receive_one_transcription(ws, request_id)
finally:
if not send_t.done():
send_t.cancel()
try:
await send_t
except (asyncio.CancelledError, websockets.WebSocketException):
pass
except Exception:
logger.exception("send-audio task failed during finalize")
@staticmethod
async def _send_audio(
ws: websockets.ClientConnection, q: asyncio.Queue[bytes | None]
ws: websockets.ClientConnection, pending: _PendingUtterance
) -> None:
# Replay anything already sent on a previous (now-dead) connection.
# sent_chunks is appended before send, so a chunk that failed to send
# last time is still present and gets retried here.
for chunk in pending.sent_chunks:
await ws.send(
json.dumps(
{
"type": "input_audio_buffer.append",
"audio": base64.b64encode(chunk).decode("ascii"),
}
)
)
if pending.ended:
await ws.send(
json.dumps({"type": "input_audio_buffer.commit", "final": True})
)
return
while True:
chunk = await q.get()
chunk = await pending.queue.get()
if chunk is None:
pending.ended = True
await ws.send(
json.dumps({"type": "input_audio_buffer.commit", "final": True})
)
continue
return
pending.sent_chunks.append(chunk)
await ws.send(
json.dumps(
{
@@ -272,16 +396,19 @@ class SpeechStream(stt.RecognizeStream):
)
)
async def _receive_events(
async def _receive_one_transcription(
self, ws: websockets.ClientConnection, request_id: str
) -> None:
# Use recv() rather than `async for`: the latter swallows
# ConnectionClosed on close-mid-iteration, which would let a dropped
# WS look like a clean "no transcription" return.
current_text = ""
async for raw in ws:
while True:
raw = await ws.recv()
data = json.loads(raw)
event_type = data.get("type")
if event_type == "transcription.delta":
# TODO: fix this
current_text += data.get("delta", "")
elif event_type == "transcription.done":
final_text = data.get("text") or current_text
@@ -308,8 +435,7 @@ class SpeechStream(stt.RecognizeStream):
),
)
)
current_text = ""
self._idle.set()
return
elif event_type == "error":
err = data.get("error")
raise APIStatusError(str(err), status_code=500, body=data)