from __future__ import annotations import asyncio import time from typing import Awaitable, Callable, Optional import numpy as np from config import settings from core.state_machine import Arbitrator from services.asr import ASRService from services.llm import LLMService from services.tts import TTSService AudioChunkCallback = Callable[[np.ndarray, int, str, int, int], Awaitable[None]] class ChatPipeline: def __init__( self, arbitrator: Arbitrator, asr: ASRService, llm: LLMService, tts: TTSService, ) -> None: self.arbitrator = arbitrator self.asr = asr self.llm = llm self.tts = tts async def _finish_avatar_after_playback(self, playback_s: float) -> None: t_end = time.perf_counter() + playback_s while time.perf_counter() < t_end: if self.arbitrator.cancel_avatar_event.is_set(): break await asyncio.sleep(0.05) await self.arbitrator.on_avatar_done() @staticmethod def _split_sentences(text: str, max_len: Optional[int] = None) -> list[str]: max_len = max(16, int(max_len or settings.tts_sentence_max_len)) hard_seps = "。！？!?；;\n" soft_seps = "，、,：:" stage1 = [] buf = [] for ch in text: buf.append(ch) if ch in hard_seps: seg = "".join(buf).strip() if seg: stage1.append(seg) buf = [] tail = "".join(buf).strip() if tail: stage1.append(tail) if not stage1: stage1 = [text] out = [] for seg in stage1: seg = seg.strip() if len(seg) <= max_len: out.append(seg) continue start = 0 while start < len(seg): end = min(len(seg), start + max_len) if end >= len(seg): out.append(seg[start:].strip()) break cut = end probe = seg[start:end] soft_cut = max(18, max_len // 2) for idx in range(len(probe) - 1, soft_cut - 1, -1): if probe[idx] in soft_seps: cut = start + idx + 1 break out.append(seg[start:cut].strip()) start = cut return [s for s in out if s] @staticmethod def _append_segment_pause(audio: np.ndarray, sr: int, is_last: bool) -> np.ndarray: if is_last or audio.size == 0: return audio pause_ms = max(0, int(settings.tts_segment_pause_ms)) if pause_ms <= 0: return audio pad = np.zeros(int(sr * pause_ms / 1000.0), dtype=np.float32) if pad.size == 0: return audio return np.concatenate([audio, pad]) async def _synthesize_sentence_first( self, reply_text: str, on_audio_chunk: Optional[AudioChunkCallback] = None, on_text_segment: Optional[Callable[[str, int, int], Awaitable[None]]] = None, should_interrupt: Optional[Callable[[], bool]] = None, ) -> tuple[np.ndarray, int, int, int]: segments = self._split_sentences(reply_text) audio_chunks = [] sr = 24000 t0 = time.perf_counter() first_chunk_ms = 0 for i, seg in enumerate(segments): if should_interrupt and should_interrupt(): break seg_audio, seg_sr = await self.tts.synthesize(seg) if should_interrupt and should_interrupt(): break sr = seg_sr seg_audio = self._append_segment_pause(seg_audio, seg_sr, is_last=(i == len(segments) - 1)) if i == 0: first_chunk_ms = int((time.perf_counter() - t0) * 1000) if on_text_segment is not None: await on_text_segment(seg, i, len(segments)) if on_audio_chunk is not None: await on_audio_chunk(seg_audio, seg_sr, seg, i, len(segments)) audio_chunks.append(seg_audio) tts_total_ms = int((time.perf_counter() - t0) * 1000) audio = np.concatenate(audio_chunks) if audio_chunks else np.zeros(1, dtype=np.float32) return audio, sr, first_chunk_ms, tts_total_ms async def process_turn( self, audio_16k: np.ndarray, on_audio_chunk: Optional[AudioChunkCallback] = None, on_user_text: Optional[Callable[[str], Awaitable[None]]] = None, on_reply_text: Optional[Callable[[str], Awaitable[None]]] = None, on_reply_segment: Optional[Callable[[str, int, int], Awaitable[None]]] = None, ) -> tuple[dict, np.ndarray, int]: await self.arbitrator.on_speech_end() asr_t0 = time.perf_counter() user_text = await self.asr.transcribe(audio_16k) asr_latency_ms = int((time.perf_counter() - asr_t0) * 1000) if not user_text.strip(): await self.arbitrator.on_avatar_done() payload = { "user_text": "", "reply_text": "", "audio_samples": 0, "sample_rate": 24000, "asr_latency_ms": asr_latency_ms, "llm_latency_ms": 0, "llm_source": "skipped", "tts_latency_ms": 0, "tts_first_chunk_ms": 0, } return payload, np.zeros(1, dtype=np.float32), 24000 if on_user_text is not None: await on_user_text(user_text) llm_t0 = time.perf_counter() reply_text, llm_source = await self.llm.reply_with_meta(user_text) if on_reply_text is not None: await on_reply_text(reply_text) llm_latency_ms = int((time.perf_counter() - llm_t0) * 1000) await self.arbitrator.on_avatar_start() audio, sr, tts_first_chunk_ms, tts_latency_ms = await self._synthesize_sentence_first( reply_text, on_audio_chunk=on_audio_chunk, on_text_segment=on_reply_segment, should_interrupt=lambda: self.arbitrator.cancel_avatar_event.is_set(), ) playback_s = max(0.2, min(5.0, float(audio.shape[0]) / float(sr))) # Do not block request completion on playback duration. asyncio.create_task(self._finish_avatar_after_playback(playback_s)) payload = { "user_text": user_text, "reply_text": reply_text, "audio_samples": int(audio.shape[0]), "sample_rate": sr, "asr_latency_ms": asr_latency_ms, "llm_latency_ms": llm_latency_ms, "llm_source": llm_source, "tts_latency_ms": tts_latency_ms, "tts_first_chunk_ms": tts_first_chunk_ms, } return payload, audio, sr async def process_text_turn( self, user_text: str, on_audio_chunk: Optional[AudioChunkCallback] = None, on_reply_text: Optional[Callable[[str], Awaitable[None]]] = None, on_reply_segment: Optional[Callable[[str, int, int], Awaitable[None]]] = None, ) -> tuple[dict, np.ndarray, int]: llm_t0 = time.perf_counter() reply_text, llm_source = await self.llm.reply_with_meta(user_text) if on_reply_text is not None: await on_reply_text(reply_text) llm_latency_ms = int((time.perf_counter() - llm_t0) * 1000) await self.arbitrator.on_avatar_start() audio, sr, tts_first_chunk_ms, tts_latency_ms = await self._synthesize_sentence_first( reply_text, on_audio_chunk=on_audio_chunk, on_text_segment=on_reply_segment, should_interrupt=lambda: self.arbitrator.cancel_avatar_event.is_set(), ) playback_s = max(0.2, min(5.0, float(audio.shape[0]) / float(sr))) asyncio.create_task(self._finish_avatar_after_playback(playback_s)) payload = { "user_text": user_text, "reply_text": reply_text, "audio_samples": int(audio.shape[0]), "sample_rate": sr, "llm_latency_ms": llm_latency_ms, "llm_source": llm_source, "tts_latency_ms": tts_latency_ms, "tts_first_chunk_ms": tts_first_chunk_ms, } return payload, audio, sr async def run_once(self, audio_16k: np.ndarray) -> dict: payload, _, _ = await self.process_turn(audio_16k) return payload