| |
| |
|
|
| import os |
| import json |
| import logging |
| import asyncio |
| import websockets |
| from typing import Optional |
|
|
| logger = logging.getLogger("stt_client") |
|
|
| class SttClient: |
| def __init__(self, ws_url: Optional[str] = None): |
| self.ws_url = ws_url or os.getenv("STT_SPACE_WS_URL", "wss://abid-official-stt.hf.space/ws/stt") |
| self.ws = None |
| self.connected = False |
| self.send_queue = asyncio.Queue() |
| self.receive_queue = asyncio.Queue() |
| self._running = False |
| self._tasks = [] |
| self._run_task = None |
|
|
| def start(self): |
| """Starts the upstream connection management task.""" |
| if self._running or (self._run_task and not self._run_task.done()): |
| logger.info("STT client is already running or start task is active.") |
| return self |
| self._running = True |
| self._run_task = asyncio.create_task(self.connect_and_run()) |
| return self |
|
|
| async def connect_and_run(self): |
| """Main connection loop managing auto-reconnect backoff and read/write loops.""" |
| backoff = 1.0 |
| while self._running: |
| try: |
| logger.info(f"Connecting to upstream STT Space: {self.ws_url}") |
| async with websockets.connect(self.ws_url) as ws: |
| self.ws = ws |
| self.connected = True |
| backoff = 1.0 |
| logger.info("Connected to upstream STT Space ✓") |
|
|
| |
| read_task = asyncio.create_task(self._read_loop()) |
| write_task = asyncio.create_task(self._write_loop()) |
| self._tasks = [read_task, write_task] |
|
|
| await asyncio.gather(read_task, write_task) |
| except asyncio.CancelledError: |
| logger.info("STT client run task cancelled.") |
| break |
| except Exception as e: |
| logger.error(f"STT upstream connection error: {e}") |
| self.connected = False |
| for t in self._tasks: |
| if not t.done(): |
| t.cancel() |
| self._tasks = [] |
| self.ws = None |
| |
| |
| await asyncio.sleep(backoff) |
| backoff = min(backoff * 2, 30.0) |
| |
| self.connected = False |
|
|
| async def _write_loop(self): |
| try: |
| while self._running: |
| data = await self.send_queue.get() |
| if self.ws and self.connected: |
| try: |
| await self.ws.send(data) |
| except Exception as e: |
| logger.error(f"STT client write error: {e}") |
| break |
| self.send_queue.task_done() |
| except asyncio.CancelledError: |
| pass |
|
|
| async def _read_loop(self): |
| try: |
| async for message in self.ws: |
| try: |
| event = json.loads(message) |
| await self.receive_queue.put(event) |
| except Exception as e: |
| logger.warning(f"Failed to parse STT message: {e} | Raw: {message[:120]}") |
| except asyncio.CancelledError: |
| pass |
| except Exception as e: |
| logger.error(f"Upstream STT read loop error: {e}") |
| self.connected = False |
|
|
| def send_pcm(self, pcm_bytes: bytes): |
| """Queue raw PCM bytes to send to STT Space.""" |
| if self._running: |
| self.send_queue.put_nowait(pcm_bytes) |
|
|
| async def get_event(self) -> dict: |
| """Fetch next transcription event from queue.""" |
| return await self.receive_queue.get() |
|
|
| async def close(self): |
| """Closes all queues, loops, and websocket connections.""" |
| self._running = False |
| self.connected = False |
| |
| if self._run_task and not self._run_task.done(): |
| self._run_task.cancel() |
| |
| for t in self._tasks: |
| if not t.done(): |
| t.cancel() |
| self._tasks = [] |
| |
| if self.ws: |
| try: |
| await self.ws.close() |
| except Exception: |
| pass |
| self.ws = None |
| logger.info("STT client connections closed.") |
|
|
