from __future__ import annotations import asyncio import contextlib import logging import time from collections.abc import Iterable from enum import Enum from realtime_voice_service.core.vad import BaseVAD, SileroVADDetector from realtime_voice_service.providers.base import BaseLLM, BaseSTT, BaseTTS, MockLLM, MockSTT, MockTTS from realtime_voice_service.transports.base import BaseMediaTransport LOGGER = logging.getLogger("uvicorn.error") class SessionState(str, Enum): LISTENING = "LISTENING" USER_SPEAKING = "USER_SPEAKING" ASSISTANT_THINKING = "ASSISTANT_THINKING" ASSISTANT_SPEAKING = "ASSISTANT_SPEAKING" class GenerationInterrupted(RuntimeError): pass class TextChunker: _BOUNDARY_CHARS = ".!?\n" def __init__(self) -> None: self._buffer = "" def feed(self, fragment: str) -> list[str]: if not fragment: return [] self._buffer += fragment return self._extract_ready_chunks() def flush(self) -> list[str]: chunks = self._extract_ready_chunks() tail = self._buffer.strip() self._buffer = "" if tail: chunks.append(tail) return chunks def reset(self) -> None: self._buffer = "" def _extract_ready_chunks(self) -> list[str]: chunks: list[str] = [] while True: next_chunk = self._pop_next_chunk() if next_chunk is None: break chunks.append(next_chunk) return chunks def _pop_next_chunk(self) -> str | None: for index, char in enumerate(self._buffer): if char not in self._BOUNDARY_CHARS: continue end_index = index + 1 while end_index < len(self._buffer) and self._buffer[end_index] in self._BOUNDARY_CHARS: end_index += 1 while end_index < len(self._buffer) and self._buffer[end_index].isspace(): end_index += 1 chunk = self._buffer[:end_index].strip() self._buffer = self._buffer[end_index:] if chunk: return chunk return None class CallSession: def __init__( self, *, session_id: str, transport: BaseMediaTransport, vad: BaseVAD | None = None, stt: BaseSTT | None = None, llm: BaseLLM | None = None, tts: BaseTTS | None = None, ) -> None: self.session_id = session_id self.transport = transport self.state = SessionState.LISTENING self.generation_epoch = 0 self.interruptions: list[str] = [] self.last_latency_ms: dict[str, int] = {} self._vad = vad or SileroVADDetector(sample_rate_hz=transport.sample_rate_hz) self._stt = stt or MockSTT(sample_rate_hz=transport.sample_rate_hz) self._llm = llm or MockLLM() self._tts = tts or MockTTS(sample_rate_hz=transport.sample_rate_hz) self._conversation: list[tuple[str, str]] = [] self._assistant_task: asyncio.Task[None] | None = None self._sentence_queue: asyncio.Queue[str | None] | None = None self._closed = False @property def conversation(self) -> tuple[tuple[str, str], ...]: return tuple(self._conversation) async def run(self) -> None: if self._assistant_task is not None: raise RuntimeError("CallSession.run() can only be called once per session") try: await self.media_loop() finally: await self.stop() async def media_loop(self) -> None: while not self._closed: audio_chunk = await self.transport.receive_audio() if audio_chunk is None: LOGGER.info("realtime session %s transport closed", self.session_id) break if not audio_chunk: continue vad_result = self._vad.feed(audio_chunk) if vad_result.is_speech and self.state != SessionState.USER_SPEAKING: if self.state in {SessionState.ASSISTANT_THINKING, SessionState.ASSISTANT_SPEAKING}: self.interrupt("barge-in") self._set_state( SessionState.USER_SPEAKING, reason=f"speech detected prob={vad_result.speech_probability:.3f}", ) if vad_result.speech_ended and vad_result.utterance_audio: speech_end_monotonic = time.perf_counter() self._set_state(SessionState.ASSISTANT_THINKING, reason="speech end detected") self._start_assistant_turn( epoch=self.generation_epoch, utterance_audio=vad_result.utterance_audio, speech_end_monotonic=speech_end_monotonic, ) def interrupt(self, reason: str = "interrupt") -> int: self.generation_epoch += 1 self.interruptions.append(reason) self._clear_sentence_queue() if self._assistant_task is not None and not self._assistant_task.done(): self._assistant_task.cancel() LOGGER.info( "realtime session %s interrupted: epoch=%s reason=%s", self.session_id, self.generation_epoch, reason, ) return self.generation_epoch async def stop(self) -> None: if self._closed: return self._closed = True self._vad.reset() self._clear_sentence_queue() if self._assistant_task is not None and not self._assistant_task.done(): self._assistant_task.cancel() with contextlib.suppress(asyncio.CancelledError): await self._assistant_task await self.transport.close() def _start_assistant_turn( self, *, epoch: int, utterance_audio: bytes, speech_end_monotonic: float, ) -> None: if self._assistant_task is not None and not self._assistant_task.done(): self._assistant_task.cancel() self._assistant_task = asyncio.create_task( self._run_assistant_turn( epoch=epoch, utterance_audio=utterance_audio, speech_end_monotonic=speech_end_monotonic, ), name=f"{self.session_id}-assistant-{epoch}", ) async def _run_assistant_turn( self, *, epoch: int, utterance_audio: bytes, speech_end_monotonic: float, ) -> None: sentence_queue: asyncio.Queue[str | None] | None = None playback_task: asyncio.Task[None] | None = None assistant_fragments: list[str] = [] try: transcript = await self._stt.transcribe(utterance_audio) self._ensure_generation(epoch) self._log_latency( "stt_latency", speech_end_monotonic, epoch=epoch, message="speech end -> transcript ready", ) transcript = transcript.strip() if not transcript: LOGGER.info("realtime session %s produced empty transcript", self.session_id) self._set_state(SessionState.LISTENING, reason="empty transcript") return self._conversation.append(("user", transcript)) sentence_queue = asyncio.Queue() self._sentence_queue = sentence_queue playback_task = asyncio.create_task( self._stream_tts_pipeline(epoch=epoch, sentence_queue=sentence_queue), name=f"{self.session_id}-playback-{epoch}", ) llm_started_monotonic = time.perf_counter() chunker = TextChunker() first_token_seen = False async for token in self._llm.generate_stream(transcript, list(self._conversation)): self._ensure_generation(epoch) if not token: continue if not first_token_seen: first_token_seen = True self._log_latency( "ttft", llm_started_monotonic, epoch=epoch, message="llm request -> first token", ) assistant_fragments.append(token) await self._enqueue_chunks( epoch=epoch, sentence_queue=sentence_queue, chunks=chunker.feed(token), ) assistant_text = "".join(assistant_fragments).strip() if not assistant_text: await self._finish_sentence_queue(sentence_queue) LOGGER.warning("realtime session %s llm produced no text", self.session_id) self._set_state(SessionState.LISTENING, reason="empty llm response") return await self._enqueue_chunks( epoch=epoch, sentence_queue=sentence_queue, chunks=chunker.flush(), ) await self._finish_sentence_queue(sentence_queue) if playback_task is not None: await playback_task self._ensure_generation(epoch) self._conversation.append(("assistant", assistant_text)) self._set_state(SessionState.LISTENING, reason="assistant turn completed") except GenerationInterrupted: LOGGER.info( "realtime session %s ignored stale generation %s", self.session_id, epoch, ) except asyncio.CancelledError: LOGGER.info( "realtime session %s cancelled generation %s", self.session_id, epoch, ) raise except Exception: LOGGER.exception( "realtime session %s generation %s failed", self.session_id, epoch, ) if epoch == self.generation_epoch and not self._closed: self._set_state(SessionState.LISTENING, reason="assistant generation failed") finally: self._clear_sentence_queue(sentence_queue) if playback_task is not None and not playback_task.done(): playback_task.cancel() with contextlib.suppress(asyncio.CancelledError): await playback_task current_task = asyncio.current_task() if self._assistant_task is current_task: self._assistant_task = None def _ensure_generation(self, epoch: int) -> None: if self._closed or epoch != self.generation_epoch: raise GenerationInterrupted(f"stale generation {epoch}") def _log_latency( self, metric_name: str, started_monotonic: float, *, epoch: int, message: str, ) -> None: latency_ms = max(int((time.perf_counter() - started_monotonic) * 1000.0), 0) self.last_latency_ms[metric_name] = latency_ms LOGGER.info( "realtime session %s %s=%sms epoch=%s (%s)", self.session_id, metric_name, latency_ms, epoch, message, ) async def _enqueue_chunks( self, *, epoch: int, sentence_queue: asyncio.Queue[str | None], chunks: Iterable[str], ) -> None: for chunk in chunks: self._ensure_generation(epoch) normalized = chunk.strip() if not normalized: continue await sentence_queue.put(normalized) async def _finish_sentence_queue(self, sentence_queue: asyncio.Queue[str | None]) -> None: await sentence_queue.put(None) async def _stream_tts_pipeline( self, *, epoch: int, sentence_queue: asyncio.Queue[str | None], ) -> None: first_audio_seen = False first_tts_started_monotonic: float | None = None while True: self._ensure_generation(epoch) sentence = await sentence_queue.get() if sentence is None: return normalized = sentence.strip() if not normalized: continue if first_tts_started_monotonic is None: first_tts_started_monotonic = time.perf_counter() async for audio_chunk in self._tts.synthesize_stream(normalized): self._ensure_generation(epoch) if not audio_chunk: continue if not first_audio_seen: first_audio_seen = True self._log_latency( "ttfa", first_tts_started_monotonic or time.perf_counter(), epoch=epoch, message="tts request -> first audio", ) self._set_state(SessionState.ASSISTANT_SPEAKING, reason="assistant playback started") await self.transport.send_audio(audio_chunk) def _clear_sentence_queue(self, sentence_queue: asyncio.Queue[str | None] | None = None) -> None: queue = sentence_queue if sentence_queue is not None else self._sentence_queue if queue is None: return while True: try: queue.get_nowait() except asyncio.QueueEmpty: break with contextlib.suppress(asyncio.QueueFull): queue.put_nowait(None) if queue is self._sentence_queue: self._sentence_queue = None def _set_state(self, state: SessionState, *, reason: str | None = None) -> None: if self.state == state: return previous_state = self.state self.state = state suffix = f" ({reason})" if reason else "" LOGGER.info( "realtime session %s state %s -> %s%s", self.session_id, previous_state.value, state.value, suffix, )