222 lines
8.3 KiB
Python
222 lines
8.3 KiB
Python
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
|