dnnai/asr_agent/worker.py

"""
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())