feat(whisper): Track connections individually not by lookup (#227)

* feat(whisper): Track connections in a list, not keyed by meeting_id

* updated logic to ensure connection is handled properly

* update vox connection manager to match new connection-only scheme

* rename disconnect, remove send

* cleanup naming

* proper close

* fix awaits and disconnect

* handle exceptions on disconnections

Author: aaronkvanmeerten

skynet/modules/stt/streaming_whisper/app.py

@@ -1,4 +1,4 @@
-from fastapi import FastAPI, WebSocket, WebSocketDisconnect
+from fastapi import FastAPI, WebSocket, WebSocketDisconnect, WebSocketException
 
 from skynet.logs import get_logger
 from skynet.modules.stt.streaming_whisper.connection_manager import ConnectionManager
@@ -12,20 +12,25 @@
 
 @app.websocket('/ws/{meeting_id}')
 async def websocket_endpoint(websocket: WebSocket, meeting_id: str, auth_token: str | None = None):
-    await ws_connection_manager.connect(websocket, meeting_id, auth_token)
-    try:
-        while True:
+    connection = await ws_connection_manager.connect(websocket, meeting_id, auth_token)
+    if connection:
+        while connection.connected:
             try:
                 chunk = await websocket.receive_bytes()
+            except WebSocketDisconnect as dc:
+                log.info(f'Meeting {connection.meeting_id} has ended')
+                await ws_connection_manager.disconnect(connection, already_closed=True)
+                break 
+            except WebSocketException as wserr:
+                log.warning(f'Error on websocket {connection.meeting_id}. Error {wserr.__class__}: \n{wserr}')
+                await ws_connection_manager.disconnect(connection)
+                break
             except Exception as err:
-                log.warning(f'Expected bytes, received something else, disconnecting {meeting_id}. Error: \n{err}')
-                ws_connection_manager.disconnect(meeting_id)
+                log.warning(f'Expected bytes, received something else, disconnecting {connection.meeting_id}. Error {err.__class__}: \n{err}')
+                await ws_connection_manager.disconnect(connection)
                 break
             if len(chunk) == 1 and ord(b'' + chunk) == 0:
-                log.info(f'Received disconnect message for {meeting_id}')
-                ws_connection_manager.disconnect(meeting_id)
+                log.info(f'Received disconnect message for {connection.meeting_id}')
+                await ws_connection_manager.disconnect(connection)
                 break
-            await ws_connection_manager.process(meeting_id, chunk, utils.now())
-    except WebSocketDisconnect:
-        ws_connection_manager.disconnect(meeting_id)
-        log.info(f'Meeting {meeting_id} has ended')
+            await ws_connection_manager.process(connection, chunk, utils.now())

skynet/modules/stt/streaming_whisper/connection_manager.py

@@ -14,11 +14,11 @@
 
 
 class ConnectionManager:
-    connections: dict[str, MeetingConnection]
+    connections: list[MeetingConnection]
     flush_audio_task: Task | None
 
     def __init__(self):
-        self.connections: dict[str, MeetingConnection] = {}
+        self.connections: list[MeetingConnection] = []
         self.flush_audio_task = None
 
     async def connect(self, websocket: WebSocket, meeting_id: str, auth_token: str | None):
@@ -29,37 +29,48 @@ async def connect(self, websocket: WebSocket, meeting_id: str, auth_token: str |
                 await websocket.close(401, 'Bad JWT token')
                 return
         await websocket.accept()
-        self.connections[meeting_id] = MeetingConnection(websocket)
+        connection = MeetingConnection(websocket, meeting_id)
+        self.connections.append(connection)
         if self.flush_audio_task is None:
             loop = asyncio.get_running_loop()
             self.flush_audio_task = loop.create_task(self.flush_working_audio_worker())
         inc_ws_conn_count()
         log.info(f'Meeting with id {meeting_id} started. Ongoing meetings {len(self.connections)}')
 
-    async def process(self, meeting_id: str, chunk: bytes, chunk_timestamp: int):
-        log.debug(f'Processing chunk for meeting {meeting_id}')
-        if meeting_id not in self.connections:
-            log.warning(f'No such meeting id {meeting_id}, the connection was probably closed.')
-            return
-        results = await self.connections[meeting_id].process(chunk, chunk_timestamp)
-        await self.send(meeting_id, results)
+        return connection
 
-    async def send(self, meeting_id: str, results: list[utils.TranscriptionResponse] | None):
+    async def process(self, connection: MeetingConnection, chunk: bytes, chunk_timestamp: int):
+        log.debug(f'Processing chunk for meeting {connection.meeting_id}')
+        
+        try:
+            results = await connection.process(chunk, chunk_timestamp)
+            await self.send(connection, results)
+        except Exception as e:
+            log.error(f'Error processing chunk for meeting {connection.meeting_id}: {e}')
+            await self.disconnect(connection)
+
+    async def send(self, connection: MeetingConnection, results: list[utils.TranscriptionResponse] | None):
         if results is not None:
             for result in results:
                 try:
-                    await self.connections[meeting_id].ws.send_json(result.model_dump())
+                    await connection.ws.send_json(result.model_dump())
                 except WebSocketDisconnect as e:
-                    log.warning(f'Meeting {meeting_id}: the connection was closed before sending all results: {e}')
-                    self.disconnect(meeting_id)
+                    log.warning(f'Meeting {connection.meeting_id}: the connection was closed before sending all results: {e}')
+                    await self.disconnect(connection, True)
+                    break
                 except Exception as ex:
-                    log.error(f'Meeting {meeting_id}: exception while sending transcription results {ex}')
+                    log.error(f'Meeting {connection.meeting_id}: exception while sending transcription results {ex}')
 
-    def disconnect(self, meeting_id: str):
+    async def disconnect(self, connection: MeetingConnection, already_closed = False):
         try:
-            del self.connections[meeting_id]
-        except KeyError:
-            log.warning(f'The meeting {meeting_id} doesn\'t exist anymore.')
+            self.connections.remove(connection)
+        except ValueError:
+            log.warning(f'The connection for meeting {connection.meeting_id} doesn\'t exist in the list anymore.')
+        if not already_closed:
+            await connection.close()
+        else:
+            # mark connection as disconnected
+            connection.disconnect()
         dec_ws_conn_count()
 
     async def flush_working_audio_worker(self):
@@ -69,15 +80,15 @@ async def flush_working_audio_worker(self):
         to the next utterance when the participant resumes speaking.
         """
         while True:
-            for meeting_id in self.connections:
-                for participant in self.connections[meeting_id].participants:
-                    state = self.connections[meeting_id].participants[participant]
+            for connection in self.connections:
+                for participant in connection.participants:
+                    state = connection.participants[participant]
                     diff = utils.now() - state.last_received_chunk
                     log.debug(
-                        f'Participant {participant} in meeting {meeting_id} has been silent for {diff} ms and has {len(state.working_audio)} bytes of audio'
+                        f'Participant {participant} in meeting {connection.meeting_id} has been silent for {diff} ms and has {len(state.working_audio)} bytes of audio'
                     )
                     if diff > whisper_flush_interval and len(state.working_audio) > 0 and not state.is_transcribing:
-                        log.info(f'Forcing a transcription in meeting {meeting_id} for {participant}')
-                        results = await self.connections[meeting_id].force_transcription(participant)
-                        await self.send(meeting_id, results)
+                        log.info(f'Forcing a transcription in meeting {connection.meeting_id} for {participant}')
+                        results = await connection.force_transcription(participant)
+                        await self.send(connection, results)
             await asyncio.sleep(1)

skynet/modules/stt/streaming_whisper/meeting_connection.py

@@ -22,14 +22,19 @@ class MeetingConnection:
     previous_transcription_store: List[List[int]]
     tokenizer: Tokenizer | None
     meeting_language: str | None
+    meeting_id: str
+    ws: WebSocket
+    connected: True
 
-    def __init__(self, ws: WebSocket):
+    def __init__(self, ws: WebSocket, meeting_id: str):
         self.participants = {}
         self.ws = ws
+        self.meeting_id = meeting_id
         self.previous_transcription_tokens = []
         self.previous_transcription_store = []
         self.meeting_language = None
         self.tokenizer = None
+        self.connected = True
 
     async def update_initial_prompt(self, previous_payloads: list[utils.TranscriptionResponse]):
         for payload in previous_payloads:
@@ -68,3 +73,10 @@ async def force_transcription(self, participant_id: str):
                 await self.update_initial_prompt(payloads)
             return payloads
         return None
+
+    def disconnect(self):
+        self.connected = False
+
+    async def close(self):
+        await self.ws.close()
+        self.disconnect()

skynet/modules/stt/vox/app.py

@@ -23,7 +23,7 @@ async def websocket_endpoint(websocket: WebSocket, auth_token: str | None = None
     decoder = PcmaDecoder()
     resampler = PcmResampler()
     session_id = utils.Uuid7().get()
-    await ws_connection_manager.connect(websocket, session_id, auth_token)
+    connection = await ws_connection_manager.connect(websocket, session_id, auth_token)
 
     data_map = dict()
     resampler = None
@@ -64,7 +64,7 @@ async def websocket_endpoint(websocket: WebSocket, auth_token: str | None = None
 
                     task = asyncio.create_task(
                         ws_connection_manager.process(
-                            session_id, participant['header'] + decoded_raw, media['timestamp']
+                            connection, participant['header'] + decoded_raw, media['timestamp']
                         )
                     )
 
@@ -75,7 +75,7 @@ async def websocket_endpoint(websocket: WebSocket, auth_token: str | None = None
                     participant['raw'] = b''
 
         except WebSocketDisconnect:
-            ws_connection_manager.disconnect(session_id)
+            ws_connection_manager.disconnect(connection, True)
             data_map.clear()
             log.info(f'Session {session_id} has ended')
             break

skynet/modules/stt/vox/connection_manager.py

@@ -1,21 +1,21 @@
 from fastapi import WebSocketDisconnect
 
 from skynet.logs import get_logger
-from skynet.modules.stt.streaming_whisper.connection_manager import ConnectionManager as BaseConnectionManager
+from skynet.modules.stt.streaming_whisper.connection_manager import ConnectionManager as BaseConnectionManager, MeetingConnection
 from skynet.modules.stt.streaming_whisper.utils.utils import TranscriptionResponse
 
 log = get_logger(__name__)
 
 
 class ConnectionManager(BaseConnectionManager):
-    async def send(self, session_id: str, results: list[TranscriptionResponse] | None):
+    async def send(self, connection: MeetingConnection, results: list[TranscriptionResponse] | None):
         if results is None:
             return
 
         final_results = [r for r in results if r.type == 'final']
         for result in final_results:
             try:
-                await self.connections[session_id].ws.send_json(
+                await connection.ws.send_json(
                     {
                         'timestamp': result.ts,
                         'tag': result.participant_id,
@@ -25,7 +25,7 @@ async def send(self, session_id: str, results: list[TranscriptionResponse] | Non
                 )
                 log.debug(f'Participant {result.participant_id} result: {result.text}')
             except WebSocketDisconnect as e:
-                log.warning(f'Session {session_id}: the connection was closed before sending all results: {e}')
-                self.disconnect(session_id)
+                log.warning(f'Session {connection.meeting_id}: the connection was closed before sending all results: {e}')
+                self.disconnect(connection, True)
             except Exception as ex:
-                log.error(f'Session {session_id}: exception while sending transcription results {ex}')
+                log.error(f'Session {connection.meeting_id}: exception while sending transcription results {ex}')