ai-video-fullstack/backend/services/pipecat/call_lifecycle.py

"""Shared call termination timing for prompt tools and workflow end nodes."""

from __future__ import annotations

import asyncio
from collections.abc import Awaitable, Callable
from dataclasses import dataclass

from loguru import logger
from pipecat.frames.frames import (
    BotStartedSpeakingFrame,
    BotStoppedSpeakingFrame,
    DataFrame,
    InterruptionFrame,
)
from pipecat.processors.frame_processor import FrameDirection, FrameProcessor


class SpeechPlaybackCompletion:
    """Awaitable completed by its exact marker at the transport output."""

    def __init__(self, coordinator: CallEndCoordinator):
        self._coordinator = coordinator
        self._future: asyncio.Future[None] = (
            asyncio.get_running_loop().create_future()
        )
        self._queued = False

    def __await__(self):
        return self._future.__await__()

    def done(self) -> bool:
        return self._future.done()

    @property
    def queued(self) -> bool:
        return self._queued

    def mark_queued(self) -> None:
        self._queued = True

    async def mark_played(self) -> None:
        await self._coordinator.complete_tracked_speech(self)

    def _resolve(self) -> None:
        if not self._future.done():
            self._future.set_result(None)


@dataclass
class FixedSpeechPlaybackMarkerFrame(DataFrame):
    """Ordered frame that identifies one fixed utterance at transport output."""

    completion: SpeechPlaybackCompletion


def playback_marker_for(
    completion: Awaitable[None] | None,
) -> FixedSpeechPlaybackMarkerFrame | None:
    """Build a marker only for the production call-end coordinator."""
    if not isinstance(completion, SpeechPlaybackCompletion):
        return None
    return FixedSpeechPlaybackMarkerFrame(completion=completion)


class CallEndCoordinator:
    """End immediately or after the currently armed closing speech finishes."""

    def __init__(self, queue_end: Callable[[str], Awaitable[None]]):
        self._queue_end = queue_end
        self._ending = False
        self._armed = False
        self._speaking = False
        self._speech_stopped = asyncio.Event()
        self._speech_stopped.set()
        self._response_speech_started = False
        self._tracked_speech_completions: set[SpeechPlaybackCompletion] = set()
        self._finish_after_tracked_speech = False
        self._finished = False
        self._reason = "completed"

    @property
    def ending(self) -> bool:
        return self._ending

    @property
    def speaking(self) -> bool:
        """Whether transport output is currently producing bot speech."""
        return self._speaking

    async def wait_until_silent(self) -> None:
        """Wait for the transport-owned bot speech boundary."""
        await self._speech_stopped.wait()

    def begin(self, reason: str) -> None:
        self._ending = True
        self._reason = reason or "completed"

    def begin_response(self) -> None:
        """Start tracking speech produced by one LLM response."""
        self._response_speech_started = False

    def arm_after_speech(self) -> None:
        """Wait for the next observed bot speech to finish."""
        self._armed = True

    def track_speech(self) -> SpeechPlaybackCompletion:
        """Register fixed speech and return its transport completion signal."""
        completion = SpeechPlaybackCompletion(self)
        self._tracked_speech_completions.add(completion)
        return completion

    async def complete_tracked_speech(
        self,
        completion: SpeechPlaybackCompletion,
    ) -> None:
        """Complete one fixed utterance when its marker reaches output."""
        if completion not in self._tracked_speech_completions:
            return
        self._tracked_speech_completions.remove(completion)
        completion._resolve()
        if (
            self._finish_after_tracked_speech
            and not self._tracked_speech_completions
        ):
            logger.info("所有工作流结束语播报完毕,挂断通话")
            await self.finish()

    async def arm_after_tracked_speech(self) -> None:
        """Finish after every already queued fixed utterance has played."""
        self._finish_after_tracked_speech = True
        if not self._tracked_speech_completions:
            await self.finish()

    async def finish_after_current_speech(self, *, has_text: bool) -> None:
        """Finish now if speech is absent/done, otherwise wait for its stop."""
        if not has_text:
            await self.finish()
            return
        if self._response_speech_started and not self._speaking:
            await self.finish()
            return
        self._armed = True

    async def finish(self) -> None:
        if self._finished:
            return
        self._finished = True
        await self._queue_end(self._reason)

    async def observe(self, frame) -> None:
        if isinstance(frame, InterruptionFrame):
            # Pipecat discards queued data frames on interruption, including
            # playback markers. Treat already queued fixed speech as stopped so
            # an interrupted Message cannot block a later EndNode forever.
            interrupted = tuple(
                completion
                for completion in self._tracked_speech_completions
                if completion.queued
            )
            for completion in interrupted:
                await self.complete_tracked_speech(completion)
        elif isinstance(frame, BotStartedSpeakingFrame):
            self._speaking = True
            self._speech_stopped.clear()
            self._response_speech_started = True
        elif isinstance(frame, BotStoppedSpeakingFrame) and self._speaking:
            self._speaking = False
            self._speech_stopped.set()
            if self._armed:
                logger.info("结束语播报完毕,挂断通话")
                await self.finish()


class EndCallAfterSpeechProcessor(FrameProcessor):
    def __init__(self, coordinator: CallEndCoordinator):
        super().__init__()
        self._coordinator = coordinator

    async def process_frame(self, frame, direction: FrameDirection):
        await super().process_frame(frame, direction)
        await self.push_frame(frame, direction)
        await self._coordinator.observe(frame)