""" LiveKit ASR Worker — subscribes to participant audio, runs Qwen ASR, and sends transcriptions back as room data messages. """ from __future__ import annotations import asyncio import json import logging import os import sys import time import traceback import jwt import numpy as np from livekit import rtc # Import whisper_asr modules (same directory or adjacent). _asr_root = os.path.join(os.path.dirname(__file__), "whisper_asr") sys.path.insert(0, _asr_root) from audio_processor import AudioBuffer # noqa: E402 from qwen_engine import QwenASREngine # noqa: E402 from transcript_processor import TranscriptProcessor # noqa: E402 logger = logging.getLogger("asr-worker") LIVEKIT_URL = os.environ.get("LIVEKIT_URL", "ws://localhost:10003") API_KEY = "devkey" API_SECRET = "secretsecretsecretsecretsecret12" ASR_MODEL_PATH = os.environ.get("ASR_MODEL_PATH", "/data/models/Qwen3-ASR") TTS_MODEL_PATH = os.environ.get("TTS_MODEL_PATH", "/data/models/Qwen3-TTS") TRANSCRIBE_INTERVAL = 2.0 BUFFER_DURATION = 5.0 SILENCE_REPEATS = 3 SILENCE_SECONDS = 3.0 def create_token(room_name: str, identity: str) -> str: now = int(time.time()) return jwt.encode( { "iss": API_KEY, "sub": identity, "name": identity, "nbf": now - 60, "exp": now + 6 * 3600, "video": { "roomJoin": True, "room": room_name, "canPublish": True, "canSubscribe": True, "canPublishData": True, }, }, API_SECRET, algorithm="HS256", ) class ASRWorker: def __init__(self, room_name: str, identity: str = "asr-bot"): self._room_name = room_name self._identity = identity self._room = rtc.Room() self._engine: QwenASREngine | None = None self._processors: dict[str, TranscriptProcessor] = {} self._buffers: dict[str, AudioBuffer] = {} self._tasks: dict[str, asyncio.Task] = {} async def run(self): logger.info("Loading Qwen ASR model...") self._engine = QwenASREngine( model_id=ASR_MODEL_PATH, language=None ) self._room.on("track_subscribed", self._on_track_subscribed) self._room.on( "participant_disconnected", self._on_participant_disconnected ) self._room.on("participant_connected", lambda p: None) self._room.on("track_published", lambda pub, p: None) token = create_token(self._room_name, self._identity) logger.info(f"Connecting to room: {self._room_name}") await self._room.connect(LIVEKIT_URL, token) logger.info(f"ASR worker ready. Participants: {len(self._room.remote_participants)}") # Subscribe to any tracks that already exist in the room. for p in self._room.remote_participants.values(): logger.info(f"Existing participant: {p.identity}, tracks: {len(p.track_publications)}") for pub in p.track_publications.values(): logger.info(f" pub kind={pub.kind}, source={pub.source.name if pub.source else 'N/A'}, track={type(pub.track).__name__ if pub.track else 'None'}, subscribed={pub.subscribed}") if pub.track is not None and pub.kind == rtc.TrackKind.KIND_AUDIO: self._on_track_subscribed(pub.track, pub, p) logger.info(f"Subscribed tracks: {list(self._buffers.keys())}") try: await asyncio.Future() finally: await self._room.disconnect() def _on_track_subscribed( self, track: rtc.RemoteAudioTrack, publication: rtc.RemoteTrackPublication, participant: rtc.RemoteParticipant, ): sid = participant.sid if sid in self._buffers: return logger.info(f"Audio track subscribed: {participant.identity} ({sid})") logger.info(f" track kind={track.kind}, muted={track.muted}, stream_state={track.stream_state}") self._buffers[sid] = AudioBuffer( max_duration=BUFFER_DURATION * 2, sample_rate=16000, ) self._processors[sid] = TranscriptProcessor( silence_repeats=SILENCE_REPEATS, silence_seconds=SILENCE_SECONDS, ) self._tasks[sid] = asyncio.create_task(self._process(sid, track)) def _on_participant_disconnected(self, participant: rtc.RemoteParticipant): sid = participant.sid t = self._tasks.pop(sid, None) if t: t.cancel() self._buffers.pop(sid, None) self._processors.pop(sid, None) async def _process(self, sid: str, track: rtc.RemoteAudioTrack): buf = self._buffers[sid] proc = self._processors[sid] async def _accumulate(): try: stream = rtc.AudioStream(track, sample_rate=16000, num_channels=1) count = 0 async for ev in stream: frame = ev.frame data = frame.data count += 1 if count == 1: logger.info(f"[{sid}] first audio frame: {len(data)} bytes, sr={frame.sample_rate}, ch={frame.num_channels}") arr = ( np.frombuffer(data, dtype=np.int16).astype(np.float32) / 32768.0 ) buf.append(arr) logger.info(f"[{sid}] stream ended after {count} frames") except Exception: logger.exception(f"[{sid}] accumulate failed") async def _transcribe(): tick = 0 last_pos = 0 while True: await asyncio.sleep(TRANSCRIBE_INTERVAL) tick += 1 total = len(buf.buffer) if total - last_pos < 3200: continue context = int(0.5 * 16000) audio = buf.buffer[max(0, last_pos - context):] if len(audio) <= 1600: continue loop = asyncio.get_running_loop() try: result = await loop.run_in_executor( None, self._engine.transcribe_array, audio ) except Exception: logger.exception("transcription failed") continue new_samples = total - last_pos last_pos = total raw = result.get("text", "").strip() if tick % 5 == 0: logger.info(f"[{sid}] transcribe tick={tick} raw='{raw}' new={new_samples}") if not raw: continue for msg in proc.feed(raw): logger.info(f"[{sid}] msg: {msg['type']} '{msg['text'][:50]}'") await self._send(msg) t = asyncio.create_task(_transcribe()) try: await _accumulate() finally: t.cancel() async def _send(self, message: dict): try: await self._room.local_participant.publish_data( json.dumps(message, ensure_ascii=False).encode(), reliable=True, topic="transcription", ) except Exception: logger.exception("publish_data failed") async def main(): import argparse parser = argparse.ArgumentParser() parser.add_argument("--room", required=True) parser.add_argument("--identity", default="asr-bot") args = parser.parse_args() logging.basicConfig(level=logging.INFO) worker = ASRWorker(room_name=args.room, identity=args.identity) await worker.run() if __name__ == "__main__": asyncio.run(main())