init
This commit is contained in:
@@ -0,0 +1,180 @@
|
||||
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
|
||||
|
||||
|
||||
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[Callable[[np.ndarray, int], Awaitable[None]]] = 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)
|
||||
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[Callable[[np.ndarray, int], Awaitable[None]]] = 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[Callable[[np.ndarray, int], 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]:
|
||||
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
|
||||
Reference in New Issue
Block a user