184 lines
6.9 KiB
Python
184 lines
6.9 KiB
Python
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import time
|
||
from typing import Awaitable, Callable, Optional
|
||
import numpy as np
|
||
|
||
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: int = 24) -> list[str]:
|
||
# Two-stage split:
|
||
# 1) hard sentence boundaries
|
||
# 2) soft commas / max length chunks for lower first-chunk latency
|
||
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:
|
||
cur = []
|
||
for ch in seg:
|
||
cur.append(ch)
|
||
cur_len = len(cur)
|
||
if ch in soft_seps and cur_len >= 8:
|
||
out.append("".join(cur).strip())
|
||
cur = []
|
||
elif cur_len >= max_len:
|
||
out.append("".join(cur).strip())
|
||
cur = []
|
||
if cur:
|
||
out.append("".join(cur).strip())
|
||
|
||
return [s for s in out if s]
|
||
|
||
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
|
||
if on_text_segment is not None:
|
||
await on_text_segment(seg, i, len(segments))
|
||
seg_audio, seg_sr = await self.tts.synthesize(seg)
|
||
sr = seg_sr
|
||
if i == 0:
|
||
first_chunk_ms = int((time.perf_counter() - t0) * 1000)
|
||
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 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)
|
||
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(),
|
||
)
|
||
await self.arbitrator.on_avatar_start()
|
||
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)
|
||
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(),
|
||
)
|
||
await self.arbitrator.on_avatar_start()
|
||
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
|