dnnai/asr_agent/conversation_worker.py

"""
Full Pipeline Worker: VAD -> ASR -> LLM -> TTS
  VAD: Silero VAD (ONNX) with dual-threshold + sliding window
  VP:  AEC/ANS/AGC handled by LiveKit WebRTC
  ASR: Qwen3-ASR (local model, batch on speech-end)
  LLM: streaming via OpenAI-compatible API (vLLM), <think> tag filtering
  TTS: Mimo API (remote), xiaozhi-style two-tier 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 numpy as np
try:
    import onnxruntime
    _HAS_ONNX = True
except ImportError:
    _HAS_ONNX = False

import aiohttp
import jwt
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")

# ── Env ──
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", "")


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

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

# ── TTS segmentation (xiaozhi-style two-tier) ──
FIRST_SENTENCE_PUNCT = "，、,。！？；:：\n"
SENTENCE_END_PUNCT   = "。！？!?\n"
MIN_TTS_CHARS        = 2


# ═══════════════════════════════════════════════════════════════
#  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, <think> filtering)
# ═══════════════════════════════════════════════════════════════

async def _llm_stream(prompt: str, hist: list[dict]):
    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,
                  "chat_template_kwargs": {"enable_thinking": False}},
            timeout=aiohttp.ClientTimeout(total=LLM_TIMEOUT),
        ) as r:
            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", {})
                    text = delta.get("content", "")
                    if text:
                        if "</think>" in text:
                            text = text.split("</think>")[-1]
                        if "<think>" in text:
                            text = text.split("<think>")[0]
                        if text.strip():
                            yield text, 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


def _clean_tts_text(text: str) -> str:
    text = re.sub(r"\*{1,3}(.*?)\*{1,3}", r"\1", text)
    text = re.sub(r"`{1,3}.*?`{1,3}", "", text)
    text = re.sub(r"\[([^\]]+)]\([^)]+\)", r"\1", text)
    text = re.sub(r"#{1,6}\s*", "", text)
    text = re.sub(r"[>\-]\s", "", text)
    text = re.sub(r"[\U0001F300-\U0001F9FF]", "", text)
    return text.strip()


# ═══════════════════════════════════════════════════════════════
#  TTS Segmenter (xiaozhi-style: two-tier, processed_chars)
# ═══════════════════════════════════════════════════════════════

@dataclass
class _TTSSegmenter:
    buffer: str = ""
    processed: int = 0
    is_first: bool = True

    def feed(self, text: str) -> list[str]:
        self.buffer += text
        segments = []
        puncts = FIRST_SENTENCE_PUNCT if self.is_first else SENTENCE_END_PUNCT
        unprocessed = self.buffer[self.processed:]

        last_pos = -1
        for ch in puncts:
            pos = unprocessed.rfind(ch)
            if pos > last_pos:
                last_pos = pos

        if last_pos >= 0:
            seg_raw = unprocessed[:last_pos + 1]
            seg_clean = _clean_tts_text(seg_raw)
            if len(seg_clean) >= MIN_TTS_CHARS:
                segments.append(seg_clean)
                self.processed += len(seg_raw)
                if self.is_first:
                    self.is_first = False
        return segments

    def flush(self) -> str:
        remaining = self.buffer[self.processed:].strip()
        if remaining:
            remaining = _clean_tts_text(remaining)
        self.processed = len(self.buffer)
        return remaining



# ── VAD ──
VAD_MODEL_PATH = os.environ.get("VAD_MODEL_PATH", "/data/silero_vad/silero_vad.onnx")
VAD_THRESHOLD_HIGH = float(os.environ.get("VAD_THRESHOLD_HIGH", "0.5"))
VAD_THRESHOLD_LOW  = float(os.environ.get("VAD_THRESHOLD_LOW", "0.2"))
VAD_WINDOW_SIZE    = int(os.environ.get("VAD_WINDOW_SIZE", "5"))
VAD_VOICE_RATIO    = float(os.environ.get("VAD_VOICE_RATIO", "0.5"))
MIN_SPEECH_S       = float(os.environ.get("MIN_SPEECH_S", "0.3"))
MIN_SILENCE_S      = float(os.environ.get("MIN_SILENCE_S", "0.8"))
MAX_SPEECH_S       = float(os.environ.get("MAX_SPEECH_S", "8.0"))

# Global ONNX session (shared across participants)
_vad_session = None


def _get_vad_session():
    global _vad_session
    if _vad_session is not None:
        return _vad_session
    if not _HAS_ONNX or not os.path.exists(VAD_MODEL_PATH):
        return None
    opts = onnxruntime.SessionOptions()
    opts.inter_op_num_threads = 1
    opts.intra_op_num_threads = 1
    _vad_session = onnxruntime.InferenceSession(
        VAD_MODEL_PATH, providers=["CPUExecutionProvider"], sess_options=opts,
    )
    logger.info("Silero VAD loaded (%s)", VAD_MODEL_PATH)
    return _vad_session


# ═══════════════════════════════════════════════════════════════
#  Silero VAD (ONNX, per-participant state)
# ═══════════════════════════════════════════════════════════════

class _SileroVAD:
    """Xiaozhi-style Silero VAD with dual-threshold + sliding window."""

    def __init__(self):
        self._sess = _get_vad_session()
        self._state = np.zeros((2, 1, 128), dtype=np.float32)
        self._context = np.zeros((1, 64), dtype=np.float32)
        self._window: deque = deque(maxlen=max(VAD_WINDOW_SIZE, 1))
        self.in_speech = False
        self.speech_start_sample = 0
        self.silence_counter = 0
        self.total_samples = 0
        self._sr = np.array(16000, dtype=np.int64)

    def process(self, frame: np.ndarray) -> bool:
        """Process 160-sample (10ms) frame. Call every frame. Returns in_speech."""
        n = len(frame)

        if self._sess is not None:
            # ── Silero ONNX inference ──
            # Silero processes 512-sample windows. Accumulate frames.
            if not hasattr(self, '_buf'):
                self._buf = np.array([], dtype=np.float32)
            self._buf = np.concatenate([self._buf, frame]).astype(np.float32)

            is_voice = False
            while len(self._buf) >= 512:
                chunk = self._buf[:512]
                self._buf = self._buf[512:]
                audio_in = chunk.reshape(1, -1)
                inp = np.concatenate([self._context, audio_in], axis=1).astype(np.float32)
                out, self._state = self._sess.run(
                    None, {"input": inp, "state": self._state, "sr": self._sr},
                )
                self._context = inp[:, -64:]
                prob = out.item()
                is_voice = prob >= VAD_THRESHOLD_HIGH or (
                    prob > VAD_THRESHOLD_LOW and (
                        self._window[-1] if self._window else False
                    )
                )
                self._window.append(is_voice)
        else:
            # ── Fallback: energy-based ──
            energy = float(np.sqrt(np.mean(frame ** 2)))
            if energy > 0.015:
                is_voice = True
            elif energy < 0.005:
                is_voice = False
            else:
                is_voice = self._window[-1] if self._window else False
            self._window.append(is_voice)

        have_voice = sum(self._window) / max(len(self._window), 1) >= VAD_VOICE_RATIO

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

    def should_end(self) -> bool:
        return self.in_speech and self.silence_counter >= int(MIN_SILENCE_S * 16000)

    def min_speech_met(self) -> bool:
        return (self.total_samples - self.speech_start_sample) / 16000 >= MIN_SPEECH_S

    def reset(self):
        self._window.clear()
        self.in_speech = False
        self.speech_start_sample = 0
        self.silence_counter = 0
        self._state = np.zeros((2, 1, 128), dtype=np.float32)
        self._context = np.zeros((1, 64), dtype=np.float32)
        if hasattr(self, '_buf'):
            self._buf = np.array([], dtype=np.float32)


# ═══════════════════════════════════════════════════════════════
#  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", p.identity)
        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 = _SileroVAD()
        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 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(arr)

                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() and vad.min_speech_met():
                    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 failed", sid)
                    busy = False
                    vad.reset()
                    vad.total_samples = len(buf.buffer)

            logger.info("[%s] stream end fc=%d", 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 + xiaozhi-style TTS ──
        reply_full = ""
        seg = _TTSSegmenter()
        tts_task = None

        try:
            async for delta, is_final in _llm_stream(txt, self.hist):
                if delta:
                    reply_full += delta
                    await self._send({"type": "reply_partial", "text": reply_full, "seq": 0})
                    for s in seg.feed(delta):
                        if tts_task:
                            await tts_task
                        logger.info("TTS seg(%d): %s", len(s), s[:40])
                        tts_task = asyncio.create_task(_tts(s))

                if is_final:
                    remaining = seg.flush()
                    if remaining:
                        if tts_task:
                            await tts_task
                        tts_task = asyncio.create_task(_tts(remaining))
                    break
        except Exception:
            logger.exception("LLM stream failed")
            reply_full = "抱歉，我暂时无法回答。"

        if tts_task:
            try:
                wav = await asyncio.wait_for(tts_task, timeout=15)
                if wav is not None and len(wav) > 0:
                    await self._play(wav, 24000)
            except (asyncio.TimeoutError, Exception):
                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:
            pass


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