Răsfoiți Sursa

feat(llm-tts): streaming LLM + incremental TTS segmentation (xiaozhi pattern)

- LLM: streaming via SSE, yield text tokens as they arrive
- TTS: split LLM output at punctuation (。!?;) and send incrementally
- Client: receive reply_partial for live text display
- Prompt: professional system prompt (identity, constraints, anti-markdown)
- Per-sentence TTS: lower first-word latency
- TTS crash isolation: each segment in its own task
wenhongquan 3 săptămâni în urmă
părinte
comite
ea42df876d
1 a modificat fișierele cu 113 adăugiri și 93 ștergeri
  1. 113 93
      asr_agent/conversation_worker.py

+ 113 - 93
asr_agent/conversation_worker.py

@@ -3,8 +3,8 @@ 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: vLLM streaming (local, or Mimo fallback)
-  TTS: Mimo API (remote)
+  LLM: streaming via OpenAI-compatible API, incremental TTS
+  TTS: Mimo API (remote), per-sentence segmentation
 """
 
 from __future__ import annotations
@@ -14,10 +14,11 @@ import base64
 import json
 import logging
 import os
+import re
 import sys
 import time
-from dataclasses import dataclass, field
 from collections import deque
+from dataclasses import dataclass, field
 
 import aiohttp
 import jwt
@@ -33,49 +34,62 @@ 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 parameters (dual-threshold hysteresis + sliding window) ──
-VAD_THRESHOLD_HIGH = 0.015   # above this → voice
-VAD_THRESHOLD_LOW = 0.005    # below this → silence; between → hold state
-VAD_WINDOW_SIZE = 8          # sliding window frame count
-VAD_VOICE_RATIO = 0.5        # frames in window must be ≥ this ratio
-MIN_SPEECH_S = 0.3            # minimum speech duration for utterance
-MIN_SILENCE_S = 0.8           # silence to end utterance
-
-# ── ASR ──
-MAX_SPEECH_S = 8.0           # force transcribe after this much continuous speech
+# ── 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_MODEL = os.environ.get("LLM_MODEL", "qwen3.6-35b-awq")
 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},
+        "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]):
-    """Streaming LLM via OpenAI-compatible API, yields text chunks."""
-    msgs = [{"role": "system", "content": "你是友好的中文语音助手,回答简洁自然,2-3句即可。"}]
+    """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=20),
+            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:
@@ -86,31 +100,18 @@ async def _llm_stream(prompt: str, hist: list[dict]):
                 if data == "[DONE]":
                     break
                 try:
-                    delta = json.loads(data)["choices"][0]["delta"]
-                    c = delta.get("content", "")
-                    if c:
-                        full += c
-                        yield c
+                    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
-    return
-
-
-async def _llm(prompt: str, hist: list[dict]) -> str:
-    """Non-streaming fallback."""
-    msgs = [{"role": "system", "content": "你是友好的中文语音助手,回答简洁自然,2-3句即可。"}]
-    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},
-            timeout=aiohttp.ClientTimeout(total=15),
-        ) as r:
-            return (await r.json())["choices"][0]["message"]["content"]
+    yield "", True
 
 
+# ── TTS ──
 async def _tts(text: str) -> np.ndarray | None:
     key = os.environ.get("MIMO_KEY", "")
     if not key:
@@ -132,11 +133,10 @@ async def _tts(text: str) -> np.ndarray | None:
                 return None
 
 
-# ── Enhanced VAD: dual-threshold + sliding window ──
+# ── Enhanced VAD (dual-threshold + sliding window) ──
 
 @dataclass
 class _VADState:
-    """Per-participant VAD state."""
     window: deque = field(default_factory=lambda: deque(maxlen=VAD_WINDOW_SIZE))
     in_speech: bool = False
     speech_start_sample: int = 0
@@ -151,27 +151,16 @@ class _VADState:
 
 
 def _vad_process(state: _VADState, frame: np.ndarray, sample_rate: int) -> bool:
-    """
-    Process one audio frame (float32, 10ms = 160 samples) through VAD.
-    Returns True if currently in-speech after processing.
-
-    Dual-threshold hysteresis + sliding window (inspired by xiaozhi SileroVAD).
-    """
     n = len(frame)
     energy = float(np.sqrt(np.mean(frame ** 2)))
-
-    # ── dual-threshold decision ──
     if energy > VAD_THRESHOLD_HIGH:
         is_voice = True
     elif energy < VAD_THRESHOLD_LOW:
         is_voice = False
     else:
-        # hold previous state
         is_voice = state.window[-1] if state.window else False
-
     state.window.append(is_voice)
-    voice_count = sum(state.window)
-    voice_ratio = voice_count / max(len(state.window), 1)
+    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:
@@ -183,20 +172,16 @@ def _vad_process(state: _VADState, frame: np.ndarray, sample_rate: int) -> bool:
         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:
-    """Check if speech should end based on silence duration."""
     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:
-    """Check if minimum speech duration has been met."""
-    speech_duration = (state.total_samples - state.speech_start_sample) / sample_rate
-    return speech_duration >= MIN_SPEECH_S
+    return (state.total_samples - state.speech_start_sample) / sample_rate >= MIN_SPEECH_S
 
 
 # ── Worker ──
@@ -245,11 +230,7 @@ class Worker:
         logger.info("_run %s", sid)
         buf = AudioBuffer(max_duration=MAX_SPEECH_S + 4, sample_rate=16000)
         sr = 16000
-
-        # ── VAD state (per participant) ──
         vad = _VADState()
-
-        # ── pipeline state ──
         busy = False
         cur_tts = None
 
@@ -260,27 +241,22 @@ class Worker:
             async for ev in stream:
                 fc += 1
                 if fc == 1:
-                    logger.info("[%s] first audio frame: %d bytes", sid, len(ev.frame.data))
-
+                    logger.info("[%s] first audio frame", sid)
                 if busy:
                     continue
-
-                data = ev.frame.data
-                arr = np.frombuffer(data, dtype=np.int16).astype(np.float32) / 32768.0
+                arr = np.frombuffer(ev.frame.data, dtype=np.int16).astype(np.float32) / 32768.0
                 buf.append(arr)
 
-                # ── VAD processing ──
                 was_speech = vad.in_speech
                 _vad_process(vad, arr, sr)
 
                 if fc % 100 == 0:
-                    logger.info("[%s] af#%d energy=%.4f vad=%s blen=%d",
+                    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))
 
-                # ── Speech ended (VAD silence threshold met) ──
                 if was_speech and _vad_should_end(vad, sr) and _vad_min_speech_met(vad, sr):
-                    logger.info("[%s] VAD END transcribe blen=%d", sid, len(buf.buffer))
+                    logger.info("[%s] VAD END blen=%d", sid, len(buf.buffer))
                     if cur_tts:
                         cur_tts = None
                     busy = True
@@ -290,11 +266,10 @@ class Worker:
                         logger.exception("[%s] transcribe failed", sid)
                     busy = False
                     vad.reset()
-                    vad.total_samples = len(buf.buffer)  # preserve buffer position after reset
+                    vad.total_samples = len(buf.buffer)
 
-                # ── Force transcribe on max speech duration ──
                 if vad.in_speech and vad.total_samples >= int(MAX_SPEECH_S * sr):
-                    logger.info("[%s] VAD MAX transcribe blen=%d", sid, len(buf.buffer))
+                    logger.info("[%s] VAD MAX blen=%d", sid, len(buf.buffer))
                     busy = True
                     try:
                         await self._transcribe_and_respond(buf)
@@ -327,24 +302,69 @@ class Worker:
 
         logger.info("ASR: %s", txt[:80])
         await self._send({"type": "utterance", "text": txt, "seq": 0})
-        await self._reply(txt)
 
-    async def _reply(self, text: str):
+        # ── 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:
-            rep = await _llm(text, self.hist)
+            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 failed")
-            rep = "抱歉,我暂时无法回答。"
-        logger.info("LLM: %s", rep[:80])
-        await self._send({"type": "reply", "text": rep})
-        try:
-            a = await _tts(rep)
-            if a is not None and len(a) > 0:
-                await self._play(a, 24000)
-        except (asyncio.CancelledError, Exception):
-            logger.exception("TTS/play failed")
-        self.hist.extend([{"role": "user", "content": text}, {"role": "assistant", "content": rep}])
-        self.hist[:] = self.hist[-20:]
+            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)