realtime_voice_service/core/session.py

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,
        )