Universal-3 Pro Streaming on LiveKit

from __future__ import annotations

import logging
import string
import time
from collections.abc import AsyncIterable

from livekit import rtc
from livekit.agents import Agent, stt
from livekit.agents.voice import ModelSettings


# "yes" / "no" deliberately omitted - in a booking flow a bare "yes"
# is a real confirmation. Edit for your domain.
BACKCHANNELS = frozenset({
    "mhm", "mm", "mmhm", "mmhmm",
    "uh", "uhhuh", "huh",
    "um", "umm", "uhm",
    "er", "erm",
    "hmm", "hm",
    "ah", "oh",
    "yeah", "yep", "yup",
    "okay", "ok",
    "right", "alright", "gotcha",
})

_TRANSCRIPT_TYPES = {
    stt.SpeechEventType.INTERIM_TRANSCRIPT,
    stt.SpeechEventType.PREFLIGHT_TRANSCRIPT,
    stt.SpeechEventType.FINAL_TRANSCRIPT,
}

_PUNCT_STRIP = str.maketrans("", "", string.punctuation)

log = logging.getLogger("backchannel_stt_filter")


def _is_all_backchannel(text: str) -> bool:
    """Return True only when every token is a known backchannel."""
    tokens = text.lower().translate(_PUNCT_STRIP).split()
    return bool(tokens) and all(tok in BACKCHANNELS for tok in tokens)


class BackchannelSTTFilterMixin:
    """Drop backchannel-only transcripts while the agent is speaking."""

    _FILTER_GRACE_S: float = 1.0
    _last_speaking_at: float = 0.0

    async def stt_node(
        self, audio: AsyncIterable[rtc.AudioFrame], model_settings: ModelSettings
    ):
        async for ev in Agent.default.stt_node(self, audio, model_settings):
            if self._should_drop(ev):
                text = ev.alternatives[0].text if ev.alternatives else ""
                log.info(
                    "event_filtered transcript=%r ev_type=%s agent_state=%s",
                    text, ev.type, self.session.agent_state,
                )
                continue
            yield ev

    def _should_drop(self, ev: stt.SpeechEvent) -> bool:
        now = time.monotonic()
        if self.session.agent_state == "speaking":
            self._last_speaking_at = now
        elif now - self._last_speaking_at > self._FILTER_GRACE_S:
            return False

        if ev.type not in _TRANSCRIPT_TYPES:
            return False

        text = ev.alternatives[0].text if ev.alternatives else ""
        return _is_all_backchannel(text)

from __future__ import annotations

import logging

from livekit.agents import AgentSession


log = logging.getLogger("short_utterance_buffer_filter")


def install_short_utterance_filter(session: AgentSession) -> None:
    """Clear transcript buffers when a short utterance arrives during agent speech."""

    @session.on("user_input_transcribed")
    def _on_user_input_transcribed(ev) -> None:
        word_count = len(ev.transcript.split())
        min_words = session.options.interruption["min_words"]

        if session.agent_state != "speaking":
            return

        if word_count >= min_words:
            return

        activity = getattr(session, "_activity", None)
        recognition = getattr(activity, "_audio_recognition", None) if activity else None
        if recognition is None:
            return

        # Wipe all three transcript buffers so short utterances
        # don't accumulate past the interrupt threshold.
        recognition._audio_transcript = ""
        recognition._audio_interim_transcript = ""
        recognition._audio_preflight_transcript = ""

        # Best-effort: abort any in-flight preemptive LLM call
        # triggered by this short utterance.
        cancel = getattr(activity, "_cancel_preemptive_generation", None)
        if callable(cancel):
            try:
                cancel()
            except Exception:
                log.debug("_cancel_preemptive_generation failed", exc_info=True)

        log.info(
            "buffer_cleared transcript=%r words=%d is_final=%s",
            ev.transcript, word_count, ev.is_final,
        )

session = AgentSession(
    stt=assemblyai.STT(
        model="u3-rt-pro",
        min_turn_silence=100,
        max_turn_silence=1000,
        vad_threshold=0.3,
    ),
    # llm=your_llm_plugin(),
    # tts=your_tts_plugin(),
    vad=None,  # Recommended: disable VAD so only STT drives interruption
    turn_handling={
        "turn_detection": "stt",
        "endpointing": {"min_delay": 1.0, "max_delay": 4.0},
        "interruption": {
            "enabled": True,
            "resume_false_interruption": True,
            "false_interruption_timeout": 1.5,
            "min_words": 2,  # Required for the buffer-clearing filter
        },
    },
)

from dotenv import load_dotenv
from livekit import agents
from livekit.agents import AgentSession, Agent, TurnHandlingOptions
from livekit.plugins import (
    assemblyai,
    silero,
)
# For MultilingualModel, uncomment the following:
# from livekit.plugins.turn_detector.multilingual import MultilingualModel

load_dotenv()


class Assistant(Agent):
    def __init__(self) -> None:
        super().__init__(instructions="You are a helpful voice AI assistant.")


async def entrypoint(ctx: agents.JobContext):
    await ctx.connect()

    session = AgentSession(
        stt=assemblyai.STT(
            model="u3-rt-pro",
            min_turn_silence=100,
            max_turn_silence=1000,  # When turn_detection="stt", override plugin default of 100.
            # If using MultilingualModel(), the plugin defaults (min: 100, max: 100) work well. Omit min_turn_silence and max_turn_silence above if preferred.
            vad_threshold=0.3,  # Match Silero's activation_threshold
        ),
        # llm=your_llm_plugin(),  # Add your LLM provider here
        # tts=your_tts_plugin(),  # Add your TTS provider here
        vad=silero.VAD.load(
            activation_threshold=0.3,  # Match AssemblyAI's internal VAD threshold
        ),
        turn_handling=TurnHandlingOptions(
            turn_detection="stt",
            # To use LiveKit's turn detection instead, replace the line above with:
            # turn_detection=MultilingualModel(),
            endpointing={"min_delay": 0},  # Avoid additive delay in STT mode
            # If using MultilingualModel(), set these instead:
            # endpointing={"min_delay": 0.5, "max_delay": 3.0},
        ),
    )

    await session.start(
        room=ctx.room,
        agent=Assistant(),
    )

    await session.generate_reply(
        instructions="Greet the user and offer your assistance."
    )


if __name__ == "__main__":
    agents.cli.run_app(agents.WorkerOptions(entrypoint_fnc=entrypoint))