Просмотр исходного кода

rewrite(conversation-worker): full pipeline rewrite based on xiaozhi-esp32-server patterns

Major improvements:
- VAD: dual-threshold hysteresis (0.015/0.005) + sliding window (8f/50%)
- Fix busy/cur_tts nonlocal bug
- Remove dead _loop; VAD end-of-speech trigger only
- Max speech duration force-transcribe (8s)
- Streaming LLM support
- Per-participant VAD state (_VADState dataclass)
- Audio frame logging every 100 frames
wenhongquan 3 недель назад
Родитель
Сommit
1e97164777
1 измененных файлов с 250 добавлено и 83 удалено
  1. 250 83
      asr_agent/conversation_worker.py

+ 250 - 83
asr_agent/conversation_worker.py

@@ -1,9 +1,9 @@
 """
 Full Pipeline Worker: VAD → ASR → LLM → TTS
-  VAD: energy-based voice activity detection (local)
+  VAD: dual-threshold energy-based + sliding window (xiaozhi-inspired)
   VP:  AEC/ANS/AGC handled by LiveKit WebRTC
-  ASR: Qwen3-ASR (local model)
-  LLM: vLLM (local, or Mimo fallback)
+  ASR: Qwen3-ASR (local model, batch on speech-end)
+  LLM: vLLM streaming (local, or Mimo fallback)
   TTS: Mimo API (remote)
 """
 
@@ -16,6 +16,8 @@ import logging
 import os
 import sys
 import time
+from dataclasses import dataclass, field
+from collections import deque
 
 import aiohttp
 import jwt
@@ -26,55 +28,178 @@ _root = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "whisper-
 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, VADProcessor
+from audio_processor import AudioBuffer
 from qwen_engine import QwenASREngine
 
 logger = logging.getLogger("worker")
 
-LIVEKIT_URL = os.environ.get("LIVEKIT_URL", "ws://localhost:7888")
+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")
 MIMO_KEY = os.environ.get("MIMO_KEY", "")
 
-# VAD params
-VAD_THRESHOLD = 0.015       # energy threshold
-MIN_SPEECH_S = 0.3           # minimum speech before considering utterance
-MIN_SILENCE_S = 1.2          # silence to end utterance
+# ── 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
+
+# ── LLM ──
+LLM_MAX_TOKENS = 128
+LLM_TEMPERATURE = 0.7
 
 
 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")
+    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")
+
+
+async def _llm_stream(prompt: str, hist: list[dict]):
+    """Streaming LLM via OpenAI-compatible API, yields text chunks."""
+    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": "default", "messages": msgs, "max_tokens": LLM_MAX_TOKENS,
+                  "temperature": LLM_TEMPERATURE, "stream": True},
+            timeout=aiohttp.ClientTimeout(total=20),
+        ) 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:
+                    delta = json.loads(data)["choices"][0]["delta"]
+                    c = delta.get("content", "")
+                    if c:
+                        full += c
+                        yield c
+                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})
+    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": "default", "messages": msgs, "max_tokens": 128, "temperature": 0.7},
-            timeout=aiohttp.ClientTimeout(total=15)) as r:
+        async with s.post(
+            f"{VLLM_URL}/chat/completions",
+            json={"model": "default", "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"]
 
 
 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"}
+    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",
+        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:
+                {"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: return None
+            except Exception:
+                return None
+
+
+# ── 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
+    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:
+    """
+    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)
+    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:
+    """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
+
+
+# ── Worker ──
+
 class Worker:
     def __init__(self, room, identity="asr-bot"):
         self.rn, self.id = room, identity
@@ -98,94 +223,125 @@ class Worker:
             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()
+        try:
+            await asyncio.Future()
+        finally:
+            await self.room.disconnect()
 
     def _on_track(self, t, _, p):
         s = p.sid
-        if s in self.tasks: return
+        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()
+        if x:
+            x.cancel()
 
     async def _run(self, sid, track):
-        buf = AudioBuffer(max_duration=10, sample_rate=16000)
-        vad = VADProcessor(
-            threshold=VAD_THRESHOLD,
-            min_speech_duration=MIN_SPEECH_S,
-            min_silence_duration=MIN_SILENCE_S,
-            sample_rate=16000,
-        )
+        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
 
         async def _acc():
-            async for ev in rtc.AudioStream(track, sample_rate=16000, num_channels=1):
-                if not busy:
-                    data = ev.frame.data
-                    arr = np.frombuffer(data, dtype=np.int16).astype(np.float32) / 32768.0
-                    buf.append(arr)
-
-                    # Feed VAD with audio frames
-                    in_speech = vad.in_speech
-                    segments = vad.process(arr)
-                    # Speech just ended
-                    if not vad.in_speech and in_speech and len(buf.buffer) > 1600 * MIN_SPEECH_S:
-                        nonlocal cur_tts
-                        if cur_tts: cur_tts = None
-                        busy = True
+            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: %d bytes", sid, len(ev.frame.data))
+
+                if busy:
+                    continue
+
+                data = ev.frame.data
+                arr = np.frombuffer(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",
+                                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))
+                    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)  # preserve buffer position after reset
+
+                # ── 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))
+                    busy = True
+                    try:
                         await self._transcribe_and_respond(buf)
-                        busy = False
-
-        async def _loop():
-            # Fallback: periodic transcription when VAD detects speech
-            last_len = 0
-            while True:
-                await asyncio.sleep(1.5)
-                if busy: continue
-                if len(buf.buffer) <= 1600: continue
-                if not vad.in_speech: continue
-                # Only transcribe if buffer grew (new speech happened)
-                if len(buf.buffer) == last_len: continue
-                last_len = len(buf.buffer)
-
-                a = buf.get_last(4.0)
-                if len(a) <= 1600: continue
-                loop = asyncio.get_running_loop()
-                result = await loop.run_in_executor(None, self.asr.transcribe_array, a)
-                txt = result.get("text", "").strip()
-                if txt:
-                    await self._send({"type": "partial", "text": txt, "seq": 0})
-
-        t = asyncio.create_task(_loop())
-        try: await _acc()
-        finally: t.cancel()
+                    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
+        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
+        if not txt:
+            return
 
+        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):
-        try: rep = await _llm(text, self.hist)
-        except: rep = "抱歉。"
+        try:
+            rep = await _llm(text, self.hist)
+        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):
-            pass
+            logger.exception("TTS/play failed")
         self.hist.extend([{"role": "user", "content": text}, {"role": "assistant", "content": rep}])
         self.hist[:] = self.hist[-20:]
 
@@ -194,21 +350,32 @@ class Worker:
         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)))
+            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: pass
+        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()
+    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())