File size: 4,402 Bytes
3de49c7 e3d8805 3de49c7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 | # stt_client.py
# Upstream WebSocket client to STT Space
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 ✓")
# Spawn read and write loop tasks
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
# Exponential backoff
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.")
|