dnnai/asr_agent/conversation_worker.py

"""
Full Pipeline Worker: VAD → ASR → LLM → TTS
  VAD: dual-threshold energy-based + sliding window (xiaozhi-inspired)
  VP:  AEC/ANS/AGC handled by LiveKit WebRTC
  ASR: Qwen3-ASR (local model, batch on speech-end)
  LLM: streaming via OpenAI-compatible API, incremental TTS
  TTS: Mimo API (remote), per-sentence segmentation
"""

from __future__ import annotations

import asyncio
import base64
import json
import logging
import os
import re
import sys
import time
from collections import deque
from dataclasses import dataclass, field

import aiohttp
import jwt
import numpy as np
from livekit import rtc

_root = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "whisper-qt-client", "whisper_asr")
if not os.path.isdir(_root):
    _root = os.path.join(os.path.dirname(os.path.abspath(__file__)), "whisper_asr")
sys.path.insert(0, _root)
from audio_processor import AudioBuffer
from qwen_engine import QwenASREngine

logger = logging.getLogger("worker")

# ── Environment ──
LIVEKIT_URL = os.environ.get("LIVEKIT_URL", "ws://localhost:7880")
ASR_MODEL = os.environ.get("ASR_MODEL_PATH", "Qwen/Qwen3-ASR-0.6B")
VLLM_URL = os.environ.get("VLLM_URL", "http://127.0.0.1:8000/v1")
LLM_MODEL = os.environ.get("LLM_MODEL", "qwen3.6-35b-awq")
MIMO_KEY = os.environ.get("MIMO_KEY", "")

# ── VAD ──
VAD_THRESHOLD_HIGH = 0.015
VAD_THRESHOLD_LOW = 0.005
VAD_WINDOW_SIZE = 8
VAD_VOICE_RATIO = 0.5
MIN_SPEECH_S = 0.3
MIN_SILENCE_S = 0.8
MAX_SPEECH_S = 8.0

# ── LLM ──
LLM_MAX_TOKENS = 128
LLM_TEMPERATURE = 0.7
LLM_TIMEOUT = 20

# ── TTS ──
TTS_CHUNK_PATTERN = re.compile(r"[。！？；\n]")
SYSTEM_PROMPT = (
    "你是友好的中文语音助手，名字叫小智。"
    "回答简洁自然，2-3句即可，用口语化中文。"
    "不要使用 Markdown、代码块、表格或特殊符号。"
    "不要输出括号注释、不要使用英文缩写。"
)


# ── JWT ──
def _token(room, ident):
    n = int(time.time())
    return jwt.encode({
        "iss": "devkey", "sub": ident, "name": ident,
        "nbf": n - 60, "exp": n + 6 * 3600,
        "video": {"roomJoin": True, "room": room,
                   "canPublish": True, "canSubscribe": True, "canPublishData": True},
    }, "secretsecretsecretsecretsecret12", algorithm="HS256")


# ── LLM Streaming ──
async def _llm_stream(prompt: str, hist: list[dict]):
    """Stream LLM tokens via SSE, yields (delta_text, is_final)."""
    msgs = [{"role": "system", "content": SYSTEM_PROMPT}]
    msgs.extend(hist)
    msgs.append({"role": "user", "content": prompt})

    async with aiohttp.ClientSession() as s:
        async with s.post(
            f"{VLLM_URL}/chat/completions",
            json={"model": LLM_MODEL, "messages": msgs,
                  "max_tokens": LLM_MAX_TOKENS, "temperature": LLM_TEMPERATURE,
                  "stream": True},
            timeout=aiohttp.ClientTimeout(total=LLM_TIMEOUT),
        ) as r:
            full = ""
            async for line in r.content:
                line = line.decode().strip()
                if not line.startswith("data: "):
                    continue
                data = line[6:]
                if data == "[DONE]":
                    break
                try:
                    chunk = json.loads(data)
                    delta = chunk.get("choices", [{}])[0].get("delta", {})
                    content = delta.get("content", "")
                    if content:
                        full += content
                        yield content, False
                except Exception:
                    continue
    yield "", True


# ── TTS ──
async def _tts(text: str) -> np.ndarray | None:
    key = os.environ.get("MIMO_KEY", "")
    if not key:
        return None
    h = {"api-key": key, "Content-Type": "application/json"}
    async with aiohttp.ClientSession() as s:
        async with s.post(
            "https://token-plan-cn.xiaomimimo.com/v1/chat/completions",
            json={"model": "mimo-v2.5-tts", "messages": [
                {"role": "user", "content": "请用自然语速朗读。"},
                {"role": "assistant", "content": text},
            ], "audio": {"format": "wav", "voice": "Chloe"}},
            headers=h, timeout=aiohttp.ClientTimeout(total=20),
        ) as r:
            try:
                b64 = (await r.json())["choices"][0]["message"]["audio"]["data"]
                return np.frombuffer(base64.b64decode(b64), dtype=np.int16)
            except Exception:
                return None


# ── Enhanced VAD (dual-threshold + sliding window) ──

@dataclass
class _VADState:
    window: deque = field(default_factory=lambda: deque(maxlen=VAD_WINDOW_SIZE))
    in_speech: bool = False
    speech_start_sample: int = 0
    silence_counter: int = 0
    total_samples: int = 0

    def reset(self) -> None:
        self.window.clear()
        self.in_speech = False
        self.speech_start_sample = 0
        self.silence_counter = 0


def _vad_process(state: _VADState, frame: np.ndarray, sample_rate: int) -> bool:
    n = len(frame)
    energy = float(np.sqrt(np.mean(frame ** 2)))
    if energy > VAD_THRESHOLD_HIGH:
        is_voice = True
    elif energy < VAD_THRESHOLD_LOW:
        is_voice = False
    else:
        is_voice = state.window[-1] if state.window else False
    state.window.append(is_voice)
    voice_ratio = sum(state.window) / max(len(state.window), 1)
    have_voice = voice_ratio >= VAD_VOICE_RATIO

    if have_voice and not state.in_speech:
        logger.info("VAD START")
        state.in_speech = True
        state.speech_start_sample = state.total_samples
        state.silence_counter = 0
    elif not have_voice and state.in_speech:
        state.silence_counter += n
    elif have_voice:
        state.silence_counter = 0
    state.total_samples += n
    return state.in_speech


def _vad_should_end(state: _VADState, sample_rate: int) -> bool:
    return state.in_speech and state.silence_counter >= int(MIN_SILENCE_S * sample_rate)


def _vad_min_speech_met(state: _VADState, sample_rate: int) -> bool:
    return (state.total_samples - state.speech_start_sample) / sample_rate >= MIN_SPEECH_S


# ── Worker ──

class Worker:
    def __init__(self, room, identity="asr-bot"):
        self.rn, self.id = room, identity
        self.room = rtc.Room()
        self.tasks: dict = {}
        self.hist: list[dict] = []

    async def run(self):
        logger.info("Loading ASR model...")
        self.asr = QwenASREngine(model_id=ASR_MODEL, language=None)
        logger.info("ASR ready")

        self.room.on("track_subscribed", self._on_track)
        self.room.on("participant_disconnected", self._on_off)
        self.room.on("participant_connected", lambda p: None)
        self.room.on("track_published", lambda pub, p: None)
        await self.room.connect(LIVEKIT_URL, _token(self.rn, self.id))
        logger.info("Worker ready (VAD+ASR+LLM+TTS)")

        for p in self.room.remote_participants.values():
            for pub in p.track_publications.values():
                if pub.track and pub.kind == rtc.TrackKind.KIND_AUDIO:
                    self._on_track(pub.track, pub, p)
        try:
            await asyncio.Future()
        finally:
            await self.room.disconnect()

    def _on_track(self, t, _, p):
        s = p.sid
        if s in self.tasks:
            return
        logger.info("TRACK %s k=%s", p.identity, t.kind)
        self.tasks[s] = asyncio.create_task(self._run(s, t))

    def _on_off(self, p):
        x = self.tasks.pop(p.sid, None)
        if x:
            x.cancel()

    async def _run(self, sid, track):
        logger.info("_run %s", sid)
        buf = AudioBuffer(max_duration=MAX_SPEECH_S + 4, sample_rate=16000)
        sr = 16000
        vad = _VADState()
        busy = False
        cur_tts = None

        async def _acc():
            nonlocal busy, cur_tts
            stream = rtc.AudioStream(track, sample_rate=sr, num_channels=1)
            fc = 0
            async for ev in stream:
                fc += 1
                if fc == 1:
                    logger.info("[%s] first audio frame", sid)
                if busy:
                    continue
                arr = np.frombuffer(ev.frame.data, dtype=np.int16).astype(np.float32) / 32768.0
                buf.append(arr)

                was_speech = vad.in_speech
                _vad_process(vad, arr, sr)

                if fc % 100 == 0:
                    logger.info("[%s] af#%d e=%.4f vad=%s blen=%d",
                                sid, fc, float(np.sqrt(np.mean(arr ** 2))),
                                vad.in_speech, len(buf.buffer))

                if was_speech and _vad_should_end(vad, sr) and _vad_min_speech_met(vad, sr):
                    logger.info("[%s] VAD END blen=%d", sid, len(buf.buffer))
                    if cur_tts:
                        cur_tts = None
                    busy = True
                    try:
                        await self._transcribe_and_respond(buf)
                    except Exception:
                        logger.exception("[%s] transcribe failed", sid)
                    busy = False
                    vad.reset()
                    vad.total_samples = len(buf.buffer)

                if vad.in_speech and vad.total_samples >= int(MAX_SPEECH_S * sr):
                    logger.info("[%s] VAD MAX blen=%d", sid, len(buf.buffer))
                    busy = True
                    try:
                        await self._transcribe_and_respond(buf)
                    except Exception:
                        logger.exception("[%s] max transcribe failed", sid)
                    busy = False
                    vad.reset()
                    vad.total_samples = len(buf.buffer)

            logger.info("[%s] stream ended after %d frames", sid, fc)

        try:
            await _acc()
        except asyncio.CancelledError:
            pass
        except Exception:
            logger.exception("[%s] _acc failed", sid)

    async def _transcribe_and_respond(self, buf: AudioBuffer):
        all_audio = buf.get_all()
        buf.clear()
        if len(all_audio) <= 1600:
            return

        loop = asyncio.get_running_loop()
        result = await loop.run_in_executor(None, self.asr.transcribe_array, all_audio)
        txt = result.get("text", "").strip()
        if not txt:
            return

        logger.info("ASR: %s", txt[:80])
        await self._send({"type": "utterance", "text": txt, "seq": 0})

        # ── Streaming LLM + incremental TTS ──
        reply_full = ""
        tts_pending = ""
        tts_task = None

        async def _play_segment(text: str):
            try:
                a = await _tts(text)
                if a is not None and len(a) > 0:
                    await self._play(a, 24000)
            except Exception:
                logger.exception("TTS segment failed")

        try:
            async for delta, is_final in _llm_stream(txt, self.hist):
                if delta:
                    reply_full += delta
                    tts_pending += delta
                    # Send partial reply to client for display
                    await self._send({"type": "reply_partial", "text": reply_full, "seq": 0})

                    # Split at punctuation for incremental TTS
                    while True:
                        m = TTS_CHUNK_PATTERN.search(tts_pending)
                        if not m:
                            break
                        pos = m.end()
                        segment = tts_pending[:pos].strip()
                        tts_pending = tts_pending[pos:].lstrip()
                        if segment and len(segment) >= 2:
                            logger.info("TTS seg: %s", segment[:40])
                            # Play previous segment if still running
                            if tts_task:
                                await tts_task
                            tts_task = asyncio.create_task(_play_segment(segment))

                if is_final:
                    # Flush remaining text
                    if tts_pending.strip():
                        if tts_task:
                            await tts_task
                        tts_task = asyncio.create_task(_play_segment(tts_pending.strip()))
                    break

        except Exception:
            logger.exception("LLM stream failed")
            reply_full = "抱歉，我暂时无法回答。"

        # Wait for final TTS segment to finish
        if tts_task:
            try:
                await asyncio.wait_for(tts_task, timeout=15)
            except asyncio.TimeoutError:
                pass

        if reply_full:
            await self._send({"type": "reply", "text": reply_full})
            self.hist.extend([
                {"role": "user", "content": txt},
                {"role": "assistant", "content": reply_full},
            ])
            self.hist[:] = self.hist[-20:]

    async def _play(self, audio: np.ndarray, sr: int):
        src = rtc.AudioSource(sr, 1)
        tk = rtc.LocalAudioTrack.create_audio_track("tts", src)
        await self.room.local_participant.publish_track(tk)
        for i in range(0, len(audio), sr // 50):
            c = audio[i: i + sr // 50]
            await src.capture_frame(rtc.AudioFrame(
                data=c.tobytes(), sample_rate=sr,
                num_channels=1, samples_per_channel=len(c),
            ))
            await asyncio.sleep(0.018)
        return tk

    async def _send(self, msg: dict):
        try:
            await self.room.local_participant.publish_data(
                json.dumps(msg, ensure_ascii=False).encode(),
                reliable=True, topic="transcription",
            )
        except Exception:
            logger.exception("send failed")


async def main():
    import argparse
    p = argparse.ArgumentParser()
    p.add_argument("--room", required=True)
    a = p.parse_args()
    logging.basicConfig(level=logging.INFO)
    await Worker(a.room).run()


if __name__ == "__main__":
    asyncio.run(main())