Module livekit.plugins.keyframe

def emotion_instructions(
    *, emotions: Sequence[Emotion] = EMOTIONS, style: MarkerStyle = "expr"
) -> str:
    """Prompt text that teaches the LLM the markers :func:`emotion_tts_node` reads.

    Args:
        emotions: The faces to offer.
        style: ``"expr"`` teaches ``<expr type="expression" label="..."/>``. With
            expressive mode (``AgentSession(expressive=...)``) pass it as
            ``tts_instructions_append`` so one marker steers voice and face; with any
            other TTS add it to the agent's instructions and it drives the face only.
            ``"brackets"`` teaches ``[label]`` tags, which ElevenLabs v3/v4 voices read
            natively; add it to the agent's instructions. Pass the same ``style`` to
            :func:`emotion_tts_node`.

    The text has no braces, so it is safe inside the expressive template.
    """
    names = ", ".join(emotions)
    if style == "expr":
        return (
            "Your face is animated from your expression markers: an "
            '<expr type="expression" label="..."/> marker also sets the expression you '
            "hold until the next one. Start every reply with one, and add one wherever "
            "the mood changes, right before the words it applies to. The face knows only "
            f"these expressions: {names}. Use only those names as expression labels; a "
            "marker with any other label leaves the face as it was."
        )
    if style == "brackets":
        return (
            "Your face is animated from emotion tags in your reply: a tag in square "
            "brackets sets the expression you hold until the next one. Start every reply "
            "with one, and add one wherever the mood changes, right before the words it "
            f"applies to. Available tags: {' '.join(f'[{e}]' for e in emotions)}. Use "
            "only these names in emotion tags. Other audio tags are fine alongside them, "
            "but they don't change your face."
        )
    raise ValueError(f"unknown marker style {style!r}")

def emotion_tts_node(
    agent: Agent,
    text: AsyncIterable[str],
    model_settings: ModelSettings,
    *,
    emotions: Sequence[Emotion] = EMOTIONS,
    style: MarkerStyle = "expr",
    max_cue_hold: float = DEFAULT_MAX_CUE_HOLD,
) -> AsyncGenerator[rtc.AudioFrame, None]:
    """``Agent.default.tts_node`` with the LLM's expression markers turned into timed cues.

    Use it as the agent's ``tts_node``::

        class MyAgent(Agent):
            def tts_node(self, text, model_settings):
                return keyframe.emotion_tts_node(self, text, model_settings)

    It is a drop-in for ``Agent.default.tts_node``, so a ``tts_node`` that processes the
    text before synthesis, or the frames after, can call it in the same place.

    Markers are ``<expr type="expression" label="..."/>`` (LiveKit's expressive markup).
    With expressive mode active the markup is passed through, so one marker steers both
    the voice and the face; otherwise it is stripped before the TTS. ``style="brackets"``
    also reads ``[tag]`` audio tags, which always go on to the TTS. They stay out of the
    transcript only with ``use_tts_aligned_transcript=True`` and a TTS that reports word
    timings; otherwise a warning is logged once per session.

    With a TTS that reports word timings, cues land on the words themselves; a
    non-streaming TTS gives sentence starts. A streaming TTS without timings moves the
    face only on a marker before a reply's first word; later markers are ignored, with a
    warning once per reply.

    Args:
        emotions: Emotions a marker may resolve to; a label naming none sends no cue.
        style: The markers the LLM writes: pass the ``style`` given to
            :func:`emotion_instructions`. In the default ``"expr"`` style, square
            brackets are text.
        max_cue_hold: Seconds of audio held back while a marker waits for its word
            timing, so the cue reaches the avatar before its audio.
    """

    activity = agent._get_activity_or_raise()
    session_tts = activity.tts
    expressive = activity._resolve_expressive_options() is not None
    aligned = session_tts is not None and bool(
        session_tts.capabilities.aligned_transcript or not session_tts.capabilities.streaming
    )
    # the transcript leaves [tag] markers out only when it is built from word timings
    if style == "brackets" and not (aligned and activity.use_tts_aligned_transcript):
        _warn_brackets_in_transcript(activity.session, aligned=aligned)

    def _inner(parsed: AsyncIterable[str]) -> AsyncIterable[rtc.AudioFrame]:
        return Agent.default.tts_node(agent, parsed, model_settings)

    return _emotion_frames(
        text,
        _inner,
        style=style,
        forward_expr=expressive,
        aligned=aligned,
        emotions=emotions,
        max_cue_hold=max_cue_hold,
    )

class AvatarSession(BaseAvatarSession):
    """A Keyframe avatar session for the LiveKit Agents framework."""

    def __init__(
        self,
        *,
        persona_id: NotGivenOr[str] = NOT_GIVEN,
        persona_slug: NotGivenOr[str] = NOT_GIVEN,
        api_url: NotGivenOr[str] = NOT_GIVEN,
        api_key: NotGivenOr[str] = NOT_GIVEN,
        avatar_participant_identity: NotGivenOr[str] = NOT_GIVEN,
        avatar_participant_name: NotGivenOr[str] = NOT_GIVEN,
        conn_options: APIConnectOptions = DEFAULT_API_CONNECT_OPTIONS,
    ) -> None:
        super().__init__()
        self._conn_options = conn_options
        self._avatar_participant_identity = (
            avatar_participant_identity
            if utils.is_given(avatar_participant_identity)
            else _AVATAR_AGENT_IDENTITY
        )
        self._avatar_participant_name = (
            avatar_participant_name
            if utils.is_given(avatar_participant_name)
            else _AVATAR_AGENT_NAME
        )

        # Resolve persona
        has_id = utils.is_given(persona_id)
        has_slug = utils.is_given(persona_slug)
        if has_id == has_slug:
            raise KeyframeException("Provide exactly one of persona_id or persona_slug")

        self._persona_id: str | None = None
        self._persona_slug: str | None = None
        if utils.is_given(persona_id):
            self._persona_id = persona_id
        if utils.is_given(persona_slug):
            self._persona_slug = persona_slug

        # Resolve API config
        api_url_val = (
            api_url if utils.is_given(api_url) else os.getenv("KEYFRAME_API_URL", DEFAULT_API_URL)
        )
        api_key_val = api_key if utils.is_given(api_key) else os.getenv("KEYFRAME_API_KEY")

        if not api_key_val:
            raise KeyframeException(
                "KEYFRAME_API_KEY must be set by arguments or environment variables"
            )

        self._api_url = api_url_val
        self._api_key = api_key_val
        self._audio_output: KeyframeAudioOutput | None = None
        self._emotions: tuple[Emotion, ...] = EMOTIONS

    @property
    def avatar_identity(self) -> str:
        return self._avatar_participant_identity

    @property
    def provider(self) -> str:
        return "keyframe"

    async def start(
        self,
        agent_session: AgentSession,
        room: rtc.Room,
        *,
        livekit_url: NotGivenOr[str] = NOT_GIVEN,
        livekit_api_key: NotGivenOr[str] = NOT_GIVEN,
        livekit_api_secret: NotGivenOr[str] = NOT_GIVEN,
    ) -> None:
        await super().start(agent_session, room)

        livekit_url = livekit_url or (os.getenv("LIVEKIT_URL") or NOT_GIVEN)
        livekit_api_key = livekit_api_key or (os.getenv("LIVEKIT_API_KEY") or NOT_GIVEN)
        livekit_api_secret = livekit_api_secret or (os.getenv("LIVEKIT_API_SECRET") or NOT_GIVEN)
        if not livekit_url or not livekit_api_key or not livekit_api_secret:
            raise KeyframeException(
                "livekit_url, livekit_api_key, and livekit_api_secret must be set "
                "by arguments or environment variables"
            )

        job_ctx = get_job_context()
        local_participant_identity = job_ctx.local_participant_identity

        # Mint a LiveKit token for the avatar worker with publish_on_behalf
        livekit_token = (
            api.AccessToken(
                api_key=livekit_api_key,
                api_secret=livekit_api_secret,
            )
            .with_kind("agent")
            .with_identity(self._avatar_participant_identity)
            .with_name(self._avatar_participant_name)
            .with_grants(
                api.VideoGrants(
                    room_join=True,
                    room=room.name,
                    can_publish=True,
                    can_subscribe=True,
                    # after joining, the avatar sets a `kfl.protocol` attribute on its own
                    # participant to advertise the audio framing it accepts; LiveKit
                    # silently drops a participant's attribute updates without this grant
                    can_update_own_metadata=True,
                )
            )
            .with_attributes({ATTRIBUTE_PUBLISH_ON_BEHALF: local_participant_identity})
            .to_jwt()
        )

        # Call API to create a reservation and dispatch
        async with KeyframeAPI(
            api_key=self._api_key,
            api_url=self._api_url,
            conn_options=self._conn_options,
        ) as keyframe_api:
            result = await keyframe_api.create_plugin_session(
                persona_id=self._persona_id,
                persona_slug=self._persona_slug,
                room_name=room.name,
                livekit_url=livekit_url,
                livekit_token=livekit_token,
                source_participant_identity=local_participant_identity,
            )
            model_id = result.get("model_id")
            logger.debug(
                "Keyframe plugin session created",
                extra={"lk.pii.session_id": result.get("reservation_id"), "model_id": model_id},
            )

        if model_id is None:
            logger.warning(
                "the Keyframe API did not report the persona's model; emotions will be untimed"
            )

        self._room = room

        timed = model_id is not None and model_id.startswith(_TIMED_EMOTION_MODEL_PREFIX)
        # with no model reported, allow every name: the avatar shows the ones it lacks as neutral
        self._emotions = EMOTIONS if timed or model_id is None else _LEGACY_EMOTIONS
        self._audio_output = KeyframeAudioOutput(
            room,
            destination_identity=self._avatar_participant_identity,
            framed=timed,
        )
        agent_session.output.replace_audio_tail(self._audio_output)

    async def set_emotion(self, emotion: Emotion) -> None:
        """Set the avatar's emotional expression.

        The avatar switches from the next video frame it renders. While it is idle, or
        when the call comes before a reply's audio (as a tool call the LLM makes before
        answering does), the change shows as soon as it can. In the middle of a reply the
        avatar has usually rendered ahead of playback, so the change shows when playback
        reaches that point, which with a fast TTS and renderer can be the end of the reply.
        An avatar that predates protocol 2 applies a mid-reply change after the audio it
        has already received instead, which is usually the end of the reply too.

        For a change on a particular word, use ``emotion_tts_node``: it times the LLM's
        expression markers against the synthesized audio in-band, so they are exact
        however fast the TTS and the renderer run.

        Args:
            emotion: One of :data:`EMOTIONS`. A ``persona-1.5-live`` model renders only
                neutral, happy, sad and angry. A name the persona's model lacks is ignored
                with a warning.
        """
        if emotion not in self._emotions:
            logger.warning(
                "ignoring emotion %r: this persona renders %s", emotion, ", ".join(self._emotions)
            )
            return

        if self._audio_output is None:
            logger.warning("set_emotion() called before start()")
            return

        await self._audio_output.set_emotion(emotion)

@dataclass(frozen=True)
class EmotionCue:
    """An emotion change timed against the audio frame that carries it.

    Attributes:
        emotion: The expression to switch to.
        offset: Seconds from the frame's first sample. May be negative (the change
            belongs to audio already sent) or run past the frame's end.
    """

    emotion: Emotion
    offset: float = 0.0

class KeyframeException(Exception):
    """Base exception for Keyframe plugin errors."""

    pass