From eb5ecab1044782ba734e0ccb767a868b33e14686 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aleix=20Conchillo=20Flaqu=C3=A9?= Date: Tue, 24 Jun 2025 16:43:10 -0700 Subject: [PATCH] no need to call start_watchdog() only reset_watchdog() --- src/pipecat/pipeline/parallel_pipeline.py | 11 +- src/pipecat/pipeline/task.py | 19 ++-- src/pipecat/pipeline/task_observer.py | 9 +- .../processors/aggregators/dtmf_aggregator.py | 1 + .../processors/aggregators/llm_response.py | 1 + src/pipecat/processors/consumer_processor.py | 5 +- src/pipecat/processors/frame_processor.py | 20 ++-- src/pipecat/processors/frameworks/rtvi.py | 9 +- src/pipecat/processors/metrics/sentry.py | 12 ++- src/pipecat/processors/producer_processor.py | 5 +- src/pipecat/services/anthropic/llm.py | 3 +- src/pipecat/services/assemblyai/stt.py | 7 +- src/pipecat/services/aws/llm.py | 3 + src/pipecat/services/aws/stt.py | 8 +- src/pipecat/services/aws_nova_sonic/aws.py | 9 +- src/pipecat/services/cartesia/tts.py | 3 +- src/pipecat/services/elevenlabs/tts.py | 6 +- .../services/gemini_multimodal_live/gemini.py | 9 +- src/pipecat/services/gladia/stt.py | 7 +- src/pipecat/services/google/llm.py | 3 +- src/pipecat/services/google/llm_openai.py | 3 +- src/pipecat/services/google/stt.py | 16 +-- src/pipecat/services/openai/base_llm.py | 3 +- .../services/openai_realtime_beta/openai.py | 7 +- src/pipecat/services/riva/stt.py | 9 +- src/pipecat/services/sambanova/llm.py | 3 +- src/pipecat/services/simli/video.py | 11 +- src/pipecat/services/tavus/video.py | 8 +- src/pipecat/services/tts_service.py | 9 +- src/pipecat/transports/base_input.py | 2 - src/pipecat/transports/base_output.py | 7 +- .../transports/network/fastapi_websocket.py | 11 +- .../transports/network/small_webrtc.py | 11 +- src/pipecat/transports/services/daily.py | 15 ++- src/pipecat/transports/services/livekit.py | 12 +-- src/pipecat/utils/asyncio.py | 101 +++++------------- src/pipecat/utils/watchdog_async_iterator.py | 60 +++++++++++ src/pipecat/utils/watchdog_event.py | 31 ++++++ src/pipecat/utils/watchdog_priority_queue.py | 35 ++++++ src/pipecat/utils/watchdog_queue.py | 35 ++++++ src/pipecat/utils/watchdog_reseter.py | 13 +++ 41 files changed, 341 insertions(+), 211 deletions(-) create mode 100644 src/pipecat/utils/watchdog_async_iterator.py create mode 100644 src/pipecat/utils/watchdog_event.py create mode 100644 src/pipecat/utils/watchdog_priority_queue.py create mode 100644 src/pipecat/utils/watchdog_queue.py create mode 100644 src/pipecat/utils/watchdog_reseter.py diff --git a/src/pipecat/pipeline/parallel_pipeline.py b/src/pipecat/pipeline/parallel_pipeline.py index 3a08f44ed..4794ea708 100644 --- a/src/pipecat/pipeline/parallel_pipeline.py +++ b/src/pipecat/pipeline/parallel_pipeline.py @@ -21,6 +21,7 @@ from pipecat.frames.frames import ( from pipecat.pipeline.base_pipeline import BasePipeline from pipecat.pipeline.pipeline import Pipeline from pipecat.processors.frame_processor import FrameDirection, FrameProcessor, FrameProcessorSetup +from pipecat.utils.watchdog_queue import WatchdogQueue class ParallelPipelineSource(FrameProcessor): @@ -83,8 +84,8 @@ class ParallelPipeline(BasePipeline): self._up_task = None self._down_task = None - self._up_queue = asyncio.Queue() - self._down_queue = asyncio.Queue() + self._up_queue = WatchdogQueue(self) + self._down_queue = WatchdogQueue(self) self._pipelines = [] @@ -202,18 +203,14 @@ class ParallelPipeline(BasePipeline): async def _process_up_queue(self): while True: frame = await self._up_queue.get() - self.start_watchdog() await self._parallel_push_frame(frame, FrameDirection.UPSTREAM) self._up_queue.task_done() - self.reset_watchdog() async def _process_down_queue(self): running = True while running: frame = await self._down_queue.get() - self.start_watchdog() - endframe_counter = self._endframe_counter.get(frame.id, 0) # If we have a counter, decrement it. @@ -228,5 +225,3 @@ class ParallelPipeline(BasePipeline): running = not (endframe_counter == 0 and isinstance(frame, EndFrame)) self._down_queue.task_done() - - self.reset_watchdog() diff --git a/src/pipecat/pipeline/task.py b/src/pipecat/pipeline/task.py index 88a47189b..ad7a5cda6 100644 --- a/src/pipecat/pipeline/task.py +++ b/src/pipecat/pipeline/task.py @@ -41,6 +41,8 @@ from pipecat.processors.frame_processor import FrameDirection, FrameProcessor, F from pipecat.utils.asyncio import WATCHDOG_TIMEOUT, BaseTaskManager, TaskManager, TaskManagerParams from pipecat.utils.tracing.setup import is_tracing_available from pipecat.utils.tracing.turn_trace_observer import TurnTraceObserver +from pipecat.utils.watchdog_queue import WatchdogQueue +from pipecat.utils.watchdog_reseter import WatchdogReseter HEARTBEAT_SECONDS = 1.0 HEARTBEAT_MONITOR_SECONDS = HEARTBEAT_SECONDS * 10 @@ -131,7 +133,7 @@ class PipelineTaskSink(FrameProcessor): await self._down_queue.put(frame) -class PipelineTask(BasePipelineTask): +class PipelineTask(WatchdogReseter, BasePipelineTask): """Manages the execution of a pipeline, handling frame processing and task lifecycle. It has a couple of event handlers `on_frame_reached_upstream` and @@ -261,18 +263,18 @@ class PipelineTask(BasePipelineTask): self._cancelled = False # This queue receives frames coming from the pipeline upstream. - self._up_queue = asyncio.Queue() + self._up_queue = WatchdogQueue(self) # This queue receives frames coming from the pipeline downstream. - self._down_queue = asyncio.Queue() + self._down_queue = WatchdogQueue(self) # This queue is the queue used to push frames to the pipeline. - self._push_queue = asyncio.Queue() + self._push_queue = WatchdogQueue(self) # This is the heartbeat queue. When a heartbeat frame is received in the # down queue we add it to the heartbeat queue for processing. - self._heartbeat_queue = asyncio.Queue() + self._heartbeat_queue = WatchdogQueue(self) # This is the idle queue. When frames are received downstream they are # put in the queue. If no frame is received the pipeline is considered # idle. - self._idle_queue = asyncio.Queue() + self._idle_queue = WatchdogQueue(self) # This event is used to indicate a finalize frame (e.g. EndFrame, # StopFrame) has been received in the down queue. self._pipeline_end_event = asyncio.Event() @@ -424,6 +426,9 @@ class PipelineTask(BasePipelineTask): for frame in frames: await self.queue_frame(frame) + def reset_watchdog(self): + self._task_manager.reset_watchdog(asyncio.current_task()) + async def _cancel(self): if not self._cancelled: logger.debug(f"Canceling pipeline task {self}") @@ -526,8 +531,6 @@ class PipelineTask(BasePipelineTask): await self._pipeline.cleanup() await self._sink.cleanup() - await self._task_manager.cleanup() - async def _process_push_queue(self): """This is the task that runs the pipeline for the first time by sending a StartFrame and by pushing any other frames queued by the user. It runs diff --git a/src/pipecat/pipeline/task_observer.py b/src/pipecat/pipeline/task_observer.py index 950ddc43b..2f7450a75 100644 --- a/src/pipecat/pipeline/task_observer.py +++ b/src/pipecat/pipeline/task_observer.py @@ -12,6 +12,8 @@ from attr import dataclass from pipecat.observers.base_observer import BaseObserver, FramePushed from pipecat.utils.asyncio import BaseTaskManager +from pipecat.utils.watchdog_queue import WatchdogQueue +from pipecat.utils.watchdog_reseter import WatchdogReseter @dataclass @@ -26,7 +28,7 @@ class Proxy: observer: BaseObserver -class TaskObserver(BaseObserver): +class TaskObserver(WatchdogReseter, BaseObserver): """This is a pipeline frame observer that is meant to be used as a proxy to the user provided observers. That is, this is the observer that should be passed to the frame processors. Then, every time a frame is pushed this @@ -89,11 +91,14 @@ class TaskObserver(BaseObserver): for proxy in self._proxies.values(): await proxy.queue.put(data) + def reset_watchdog(self): + self._task_manager.reset_watchdog(asyncio.current_task()) + def _started(self) -> bool: return self._proxies is not None def _create_proxy(self, observer: BaseObserver) -> Proxy: - queue = asyncio.Queue() + queue = WatchdogQueue(self) task = self._task_manager.create_task( self._proxy_task_handler(queue, observer), f"TaskObserver::{observer}::_proxy_task_handler", diff --git a/src/pipecat/processors/aggregators/dtmf_aggregator.py b/src/pipecat/processors/aggregators/dtmf_aggregator.py index cc6218a1f..3008fb398 100644 --- a/src/pipecat/processors/aggregators/dtmf_aggregator.py +++ b/src/pipecat/processors/aggregators/dtmf_aggregator.py @@ -119,6 +119,7 @@ class DTMFAggregator(FrameProcessor): await asyncio.wait_for(self._digit_event.wait(), timeout=self._idle_timeout) self._digit_event.clear() except asyncio.TimeoutError: + self.reset_watchdog() if self._aggregation: await self._flush_aggregation() diff --git a/src/pipecat/processors/aggregators/llm_response.py b/src/pipecat/processors/aggregators/llm_response.py index 1984e3cf8..40016aaa0 100644 --- a/src/pipecat/processors/aggregators/llm_response.py +++ b/src/pipecat/processors/aggregators/llm_response.py @@ -470,6 +470,7 @@ class LLMUserContextAggregator(LLMContextResponseAggregator): ) self._emulating_vad = False finally: + self.reset_watchdog() self._aggregation_event.clear() async def _maybe_emulate_user_speaking(self): diff --git a/src/pipecat/processors/consumer_processor.py b/src/pipecat/processors/consumer_processor.py index 121dd7712..10cae11a3 100644 --- a/src/pipecat/processors/consumer_processor.py +++ b/src/pipecat/processors/consumer_processor.py @@ -10,6 +10,7 @@ from typing import Awaitable, Callable, Optional from pipecat.frames.frames import CancelFrame, EndFrame, Frame, StartFrame from pipecat.processors.frame_processor import FrameDirection, FrameProcessor from pipecat.processors.producer_processor import ProducerProcessor, identity_transformer +from pipecat.utils.watchdog_queue import WatchdogQueue class ConsumerProcessor(FrameProcessor): @@ -31,7 +32,7 @@ class ConsumerProcessor(FrameProcessor): super().__init__(**kwargs) self._transformer = transformer self._direction = direction - self._queue: asyncio.Queue = producer.add_consumer() + self._queue: WatchdogQueue = producer.add_consumer(self) self._consumer_task: Optional[asyncio.Task] = None async def process_frame(self, frame: Frame, direction: FrameDirection): @@ -61,7 +62,5 @@ class ConsumerProcessor(FrameProcessor): async def _consumer_task_handler(self): while True: frame = await self._queue.get() - self.start_watchdog() new_frame = await self._transformer(frame) await self.push_frame(new_frame, self._direction) - self.reset_watchdog() diff --git a/src/pipecat/processors/frame_processor.py b/src/pipecat/processors/frame_processor.py index bee5ce91c..d7723d868 100644 --- a/src/pipecat/processors/frame_processor.py +++ b/src/pipecat/processors/frame_processor.py @@ -31,6 +31,9 @@ from pipecat.observers.base_observer import BaseObserver, FramePushed from pipecat.processors.metrics.frame_processor_metrics import FrameProcessorMetrics from pipecat.utils.asyncio import BaseTaskManager from pipecat.utils.base_object import BaseObject +from pipecat.utils.watchdog_event import WatchdogEvent +from pipecat.utils.watchdog_queue import WatchdogQueue +from pipecat.utils.watchdog_reseter import WatchdogReseter class FrameDirection(Enum): @@ -45,13 +48,13 @@ class FrameProcessorSetup: observer: Optional[BaseObserver] = None -class FrameProcessor(BaseObject): +class FrameProcessor(WatchdogReseter, BaseObject): def __init__( self, *, name: Optional[str] = None, - metrics: Optional[FrameProcessorMetrics] = None, enable_watchdog_logging: Optional[bool] = None, + metrics: Optional[FrameProcessorMetrics] = None, watchdog_timeout_secs: Optional[float] = None, **kwargs, ): @@ -101,7 +104,7 @@ class FrameProcessor(BaseObject): # is called. To resume processing frames we need to call # `resume_processing_frames()` which will wake up the event. self.__should_block_frames = False - self.__input_event = asyncio.Event() + self.__input_event = WatchdogEvent(self) self.__input_frame_task: Optional[asyncio.Task] = None # Every processor in Pipecat should only output frames from a single @@ -210,9 +213,6 @@ class FrameProcessor(BaseObject): async def wait_for_task(self, task: asyncio.Task, timeout: Optional[float] = None): await self.get_task_manager().wait_for_task(task, timeout) - def start_watchdog(self): - self.get_task_manager().start_watchdog(asyncio.current_task()) - def reset_watchdog(self): self.get_task_manager().reset_watchdog(asyncio.current_task()) @@ -397,7 +397,7 @@ class FrameProcessor(BaseObject): if not self.__input_frame_task: self.__should_block_frames = False self.__input_event.clear() - self.__input_queue = asyncio.Queue() + self.__input_queue = WatchdogQueue(self) self.__input_frame_task = self.create_task(self.__input_frame_task_handler()) async def __cancel_input_task(self): @@ -416,7 +416,6 @@ class FrameProcessor(BaseObject): (frame, direction, callback) = await self.__input_queue.get() try: - self.start_watchdog() # Process the frame. await self.process_frame(frame, direction) # If this frame has an associated callback, call it now. @@ -427,11 +426,10 @@ class FrameProcessor(BaseObject): await self.push_error(ErrorFrame(str(e))) finally: self.__input_queue.task_done() - self.reset_watchdog() def __create_push_task(self): if not self.__push_frame_task: - self.__push_queue = asyncio.Queue() + self.__push_queue = WatchdogQueue(self) self.__push_frame_task = self.create_task(self.__push_frame_task_handler()) async def __cancel_push_task(self): @@ -442,7 +440,5 @@ class FrameProcessor(BaseObject): async def __push_frame_task_handler(self): while True: (frame, direction) = await self.__push_queue.get() - self.start_watchdog() await self.__internal_push_frame(frame, direction) self.__push_queue.task_done() - self.reset_watchdog() diff --git a/src/pipecat/processors/frameworks/rtvi.py b/src/pipecat/processors/frameworks/rtvi.py index 29e85e5d7..e35f72e0b 100644 --- a/src/pipecat/processors/frameworks/rtvi.py +++ b/src/pipecat/processors/frameworks/rtvi.py @@ -68,6 +68,7 @@ from pipecat.transports.base_input import BaseInputTransport from pipecat.transports.base_output import BaseOutputTransport from pipecat.transports.base_transport import BaseTransport from pipecat.utils.string import match_endofsentence +from pipecat.utils.watchdog_queue import WatchdogQueue RTVI_PROTOCOL_VERSION = "0.3.0" @@ -650,11 +651,11 @@ class RTVIProcessor(FrameProcessor): self._registered_services: Dict[str, RTVIService] = {} # A task to process incoming action frames. - self._action_queue = asyncio.Queue() + self._action_queue = WatchdogQueue(self) self._action_task: Optional[asyncio.Task] = None # A task to process incoming transport messages. - self._message_queue = asyncio.Queue() + self._message_queue = WatchdogQueue(self) self._message_task: Optional[asyncio.Task] = None self._register_event_handler("on_bot_started") @@ -783,18 +784,14 @@ class RTVIProcessor(FrameProcessor): async def _action_task_handler(self): while True: frame = await self._action_queue.get() - self.start_watchdog() await self._handle_action(frame.message_id, frame.rtvi_action_run) self._action_queue.task_done() - self.reset_watchdog() async def _message_task_handler(self): while True: message = await self._message_queue.get() - self.start_watchdog() await self._handle_message(message) self._message_queue.task_done() - self.reset_watchdog() async def _handle_transport_message(self, frame: TransportMessageUrgentFrame): try: diff --git a/src/pipecat/processors/metrics/sentry.py b/src/pipecat/processors/metrics/sentry.py index 20d854ffb..ac8ca2095 100644 --- a/src/pipecat/processors/metrics/sentry.py +++ b/src/pipecat/processors/metrics/sentry.py @@ -9,6 +9,8 @@ import asyncio from loguru import logger from pipecat.utils.asyncio import TaskManager +from pipecat.utils.watchdog_queue import WatchdogQueue +from pipecat.utils.watchdog_reseter import WatchdogReseter try: import sentry_sdk @@ -20,7 +22,7 @@ except ModuleNotFoundError as e: from pipecat.processors.metrics.frame_processor_metrics import FrameProcessorMetrics -class SentryMetrics(FrameProcessorMetrics): +class SentryMetrics(WatchdogReseter, FrameProcessorMetrics): def __init__(self): super().__init__() self._ttfb_metrics_tx = None @@ -28,13 +30,12 @@ class SentryMetrics(FrameProcessorMetrics): self._sentry_available = sentry_sdk.is_initialized() if not self._sentry_available: logger.warning("Sentry SDK not initialized. Sentry features will be disabled.") - self._sentry_queue = asyncio.Queue() self._sentry_task = None async def setup(self, task_manager: TaskManager): await super().setup(task_manager) if self._sentry_available: - self._sentry_queue = asyncio.Queue() + self._sentry_queue = WatchdogQueue(self) self._sentry_task = self.task_manager.create_task( self._sentry_task_handler(), name=f"{self}::_sentry_task_handler" ) @@ -48,6 +49,10 @@ class SentryMetrics(FrameProcessorMetrics): logger.trace(f"{self} Flushing Sentry metrics") sentry_sdk.flush(timeout=5.0) + def reset_watchdog(self): + if self._task_manager: + self._task_manager.reset_watchdog(asyncio.current_task()) + async def start_ttfb_metrics(self, report_only_initial_ttfb): await super().start_ttfb_metrics(report_only_initial_ttfb) @@ -93,3 +98,4 @@ class SentryMetrics(FrameProcessorMetrics): if tx: await self.task_manager.get_event_loop().run_in_executor(None, tx.finish) running = tx is not None + self._sentry_queue.task_done() diff --git a/src/pipecat/processors/producer_processor.py b/src/pipecat/processors/producer_processor.py index 6ada2ed83..6dd381a51 100644 --- a/src/pipecat/processors/producer_processor.py +++ b/src/pipecat/processors/producer_processor.py @@ -9,6 +9,7 @@ from typing import Awaitable, Callable, List from pipecat.frames.frames import Frame from pipecat.processors.frame_processor import FrameDirection, FrameProcessor +from pipecat.utils.watchdog_queue import WatchdogQueue async def identity_transformer(frame: Frame): @@ -36,14 +37,14 @@ class ProducerProcessor(FrameProcessor): self._passthrough = passthrough self._consumers: List[asyncio.Queue] = [] - def add_consumer(self): + def add_consumer(self, consumer: FrameProcessor): """ Adds a new consumer and returns its associated queue. Returns: asyncio.Queue: The queue for the newly added consumer. """ - queue = asyncio.Queue() + queue = WatchdogQueue(consumer) self._consumers.append(queue) return queue diff --git a/src/pipecat/services/anthropic/llm.py b/src/pipecat/services/anthropic/llm.py index 236f269fa..0d92e8b30 100644 --- a/src/pipecat/services/anthropic/llm.py +++ b/src/pipecat/services/anthropic/llm.py @@ -47,6 +47,7 @@ from pipecat.processors.aggregators.openai_llm_context import ( from pipecat.processors.frame_processor import FrameDirection from pipecat.services.llm_service import FunctionCallFromLLM, LLMService from pipecat.utils.tracing.service_decorators import traced_llm +from pipecat.utils.watchdog_async_iterator import WatchdogAsyncIterator try: from anthropic import NOT_GIVEN, AsyncAnthropic, NotGiven @@ -203,7 +204,7 @@ class AnthropicLLMService(LLMService): json_accumulator = "" function_calls = [] - async for event in response: + async for event in WatchdogAsyncIterator(response, reseter=self): # Aggregate streaming content, create frames, trigger events if event.type == "content_block_delta": diff --git a/src/pipecat/services/assemblyai/stt.py b/src/pipecat/services/assemblyai/stt.py index 09aaa4d25..452d4cfb6 100644 --- a/src/pipecat/services/assemblyai/stt.py +++ b/src/pipecat/services/assemblyai/stt.py @@ -189,17 +189,16 @@ class AssemblyAISTTService(STTService): try: while self._connected: try: - message = await self._websocket.recv() - self.start_watchdog() + message = await asyncio.wait_for(self._websocket.recv(), timeout=1.0) data = json.loads(message) await self._handle_message(data) + except asyncio.TimeoutError: + self.reset_watchdog() except websockets.exceptions.ConnectionClosedOK: break except Exception as e: logger.error(f"Error processing WebSocket message: {e}") break - finally: - self.reset_watchdog() except Exception as e: logger.error(f"Fatal error in receive handler: {e}") diff --git a/src/pipecat/services/aws/llm.py b/src/pipecat/services/aws/llm.py index dcec91463..d3217e7a1 100644 --- a/src/pipecat/services/aws/llm.py +++ b/src/pipecat/services/aws/llm.py @@ -711,6 +711,8 @@ class AWSBedrockLLMService(LLMService): function_calls = [] for event in response["stream"]: + self.reset_watchdog() + # Handle text content if "contentBlockDelta" in event: delta = event["contentBlockDelta"]["delta"] @@ -762,6 +764,7 @@ class AWSBedrockLLMService(LLMService): completion_tokens += usage.get("outputTokens", 0) cache_read_input_tokens += usage.get("cacheReadInputTokens", 0) cache_creation_input_tokens += usage.get("cacheWriteInputTokens", 0) + await self.run_function_calls(function_calls) except asyncio.CancelledError: # If we're interrupted, we won't get a complete usage report. So set our flag to use the diff --git a/src/pipecat/services/aws/stt.py b/src/pipecat/services/aws/stt.py index c9c1b32d1..c4170ebad 100644 --- a/src/pipecat/services/aws/stt.py +++ b/src/pipecat/services/aws/stt.py @@ -284,9 +284,7 @@ class AWSTranscribeSTTService(STTService): break try: - response = await self._ws_client.recv() - - self.start_watchdog() + response = await asyncio.wait_for(self._ws_client.recv(), timeout=1.0) headers, payload = decode_event(response) @@ -337,6 +335,8 @@ class AWSTranscribeSTTService(STTService): else: logger.debug(f"{self} Other message type received: {headers}") logger.debug(f"{self} Payload: {payload}") + except asyncio.TimeoutError: + self.reset_watchdog() except websockets.exceptions.ConnectionClosed as e: logger.error( f"{self} WebSocket connection closed in receive loop with code {e.code}: {e.reason}" @@ -345,5 +345,3 @@ class AWSTranscribeSTTService(STTService): except Exception as e: logger.error(f"{self} Unexpected error in receive loop: {e}") break - finally: - self.reset_watchdog() diff --git a/src/pipecat/services/aws_nova_sonic/aws.py b/src/pipecat/services/aws_nova_sonic/aws.py index 718e3dc50..93eb77e90 100644 --- a/src/pipecat/services/aws_nova_sonic/aws.py +++ b/src/pipecat/services/aws_nova_sonic/aws.py @@ -697,9 +697,9 @@ class AWSNovaSonicLLMService(LLMService): try: while self._stream and not self._disconnecting: output = await self._stream.await_output() - result = await output[1].receive() + result = await asyncio.wait_for(output[1].receive(), timeout=1.0) - self.start_watchdog() + self.reset_watchdog() if result.value and result.value.bytes_: response_data = result.value.bytes_.decode("utf-8") @@ -728,13 +728,12 @@ class AWSNovaSonicLLMService(LLMService): elif "completionEnd" in event_json: # Handle the LLM completion ending await self._handle_completion_end_event(event_json) - + except asyncio.TimeoutError: + self.reset_watchdog() except Exception as e: logger.error(f"{self} error processing responses: {e}") if self._wants_connection: await self.reset_conversation() - finally: - self.reset_watchdog() async def _handle_completion_start_event(self, event_json): pass diff --git a/src/pipecat/services/cartesia/tts.py b/src/pipecat/services/cartesia/tts.py index 22493f229..30f5c0754 100644 --- a/src/pipecat/services/cartesia/tts.py +++ b/src/pipecat/services/cartesia/tts.py @@ -30,6 +30,7 @@ from pipecat.transcriptions.language import Language from pipecat.utils.text.base_text_aggregator import BaseTextAggregator from pipecat.utils.text.skip_tags_aggregator import SkipTagsAggregator from pipecat.utils.tracing.service_decorators import traced_tts +from pipecat.utils.watchdog_async_iterator import WatchdogAsyncIterator # See .env.example for Cartesia configuration needed try: @@ -255,7 +256,7 @@ class CartesiaTTSService(AudioContextWordTTSService): self._context_id = None async def _receive_messages(self): - async for message in self._get_websocket(): + async for message in WatchdogAsyncIterator(self._get_websocket(), reseter=self): msg = json.loads(message) if not msg or not self.audio_context_available(msg["context_id"]): continue diff --git a/src/pipecat/services/elevenlabs/tts.py b/src/pipecat/services/elevenlabs/tts.py index ccd9b5b3f..d0bd29f5b 100644 --- a/src/pipecat/services/elevenlabs/tts.py +++ b/src/pipecat/services/elevenlabs/tts.py @@ -33,6 +33,7 @@ from pipecat.services.tts_service import ( ) from pipecat.transcriptions.language import Language from pipecat.utils.tracing.service_decorators import traced_tts +from pipecat.utils.watchdog_async_iterator import WatchdogAsyncIterator # See .env.example for ElevenLabs configuration needed try: @@ -394,7 +395,7 @@ class ElevenLabsTTSService(AudioContextWordTTSService): self._started = False async def _receive_messages(self): - async for message in self._get_websocket(): + async for message in WatchdogAsyncIterator(self._get_websocket(), reseter=self): msg = json.loads(message) received_ctx_id = msg.get("contextId") @@ -426,7 +427,8 @@ class ElevenLabsTTSService(AudioContextWordTTSService): async def _keepalive_task_handler(self): while True: - await asyncio.sleep(10) + self.reset_watchdog() + await asyncio.sleep(4) try: if self._websocket and self._websocket.open: if self._context_id: diff --git a/src/pipecat/services/gemini_multimodal_live/gemini.py b/src/pipecat/services/gemini_multimodal_live/gemini.py index d8a90fada..44829500f 100644 --- a/src/pipecat/services/gemini_multimodal_live/gemini.py +++ b/src/pipecat/services/gemini_multimodal_live/gemini.py @@ -60,7 +60,8 @@ from pipecat.services.openai.llm import ( from pipecat.transcriptions.language import Language from pipecat.utils.string import match_endofsentence from pipecat.utils.time import time_now_iso8601 -from pipecat.utils.tracing.service_decorators import traced_gemini_live, traced_stt, traced_tts +from pipecat.utils.tracing.service_decorators import traced_gemini_live, traced_stt +from pipecat.utils.watchdog_async_iterator import WatchdogAsyncIterator from . import events @@ -686,9 +687,7 @@ class GeminiMultimodalLiveLLMService(LLMService): # async def _receive_task_handler(self): - async for message in self._websocket: - self.start_watchdog() - + async for message in WatchdogAsyncIterator(self._websocket, reseter=self): evt = events.parse_server_event(message) # logger.debug(f"Received event: {message[:500]}") # logger.debug(f"Received event: {evt}") @@ -711,8 +710,6 @@ class GeminiMultimodalLiveLLMService(LLMService): # errors are fatal, so exit the receive loop return - self.reset_watchdog() - # # # diff --git a/src/pipecat/services/gladia/stt.py b/src/pipecat/services/gladia/stt.py index 885fe8dc2..ef73c9c97 100644 --- a/src/pipecat/services/gladia/stt.py +++ b/src/pipecat/services/gladia/stt.py @@ -27,6 +27,7 @@ from pipecat.services.stt_service import STTService from pipecat.transcriptions.language import Language from pipecat.utils.time import time_now_iso8601 from pipecat.utils.tracing.service_decorators import traced_stt +from pipecat.utils.watchdog_async_iterator import WatchdogAsyncIterator try: import websockets @@ -501,9 +502,7 @@ class GladiaSTTService(STTService): async def _receive_task_handler(self): try: - async for message in self._websocket: - self.start_watchdog() - + async for message in WatchdogAsyncIterator(self._websocket, reseter=self): content = json.loads(message) # Handle audio chunk acknowledgments @@ -568,8 +567,6 @@ class GladiaSTTService(STTService): pass except Exception as e: logger.error(f"Error in Gladia WebSocket handler: {e}") - finally: - self.reset_watchdog() async def _maybe_reconnect(self) -> bool: """Handle exponential backoff reconnection logic.""" diff --git a/src/pipecat/services/google/llm.py b/src/pipecat/services/google/llm.py index f983b7342..38fba45b9 100644 --- a/src/pipecat/services/google/llm.py +++ b/src/pipecat/services/google/llm.py @@ -48,6 +48,7 @@ from pipecat.services.openai.llm import ( OpenAIUserContextAggregator, ) from pipecat.utils.tracing.service_decorators import traced_llm +from pipecat.utils.watchdog_async_iterator import WatchdogAsyncIterator # Suppress gRPC fork warnings os.environ["GRPC_ENABLE_FORK_SUPPORT"] = "false" @@ -557,7 +558,7 @@ class GoogleLLMService(LLMService): ) function_calls = [] - async for chunk in response: + async for chunk in WatchdogAsyncIterator(response, reseter=self): # Stop TTFB metrics after the first chunk await self.stop_ttfb_metrics() if chunk.usage_metadata: diff --git a/src/pipecat/services/google/llm_openai.py b/src/pipecat/services/google/llm_openai.py index a497cb229..e7fcfdb0f 100644 --- a/src/pipecat/services/google/llm_openai.py +++ b/src/pipecat/services/google/llm_openai.py @@ -11,6 +11,7 @@ from openai import AsyncStream from openai.types.chat import ChatCompletionChunk from pipecat.services.llm_service import FunctionCallFromLLM +from pipecat.utils.watchdog_async_iterator import WatchdogAsyncIterator # Suppress gRPC fork warnings os.environ["GRPC_ENABLE_FORK_SUPPORT"] = "false" @@ -53,7 +54,7 @@ class GoogleLLMOpenAIBetaService(OpenAILLMService): context ) - async for chunk in chunk_stream: + async for chunk in WatchdogAsyncIterator(chunk_stream, reseter=self): if chunk.usage: tokens = LLMTokenUsage( prompt_tokens=chunk.usage.prompt_tokens, diff --git a/src/pipecat/services/google/stt.py b/src/pipecat/services/google/stt.py index 26186e6bf..067db7a3a 100644 --- a/src/pipecat/services/google/stt.py +++ b/src/pipecat/services/google/stt.py @@ -10,6 +10,7 @@ import os import time from pipecat.utils.tracing.service_decorators import traced_stt +from pipecat.utils.watchdog_async_iterator import WatchdogAsyncIterator # Suppress gRPC fork warnings os.environ["GRPC_ENABLE_FORK_SUPPORT"] = "false" @@ -747,8 +748,6 @@ class GoogleSTTService(STTService): try: while True: try: - self.start_watchdog() - if self._request_queue.empty(): # wait for 10ms in case we don't have audio await asyncio.sleep(0.01) @@ -763,8 +762,6 @@ class GoogleSTTService(STTService): # Process responses await self._process_responses(streaming_recognize) - self.reset_watchdog() - # If we're here, check if we need to reconnect if (int(time.time() * 1000) - self._stream_start_time) > self.STREAMING_LIMIT: logger.debug("Reconnecting stream after timeout") @@ -779,8 +776,6 @@ class GoogleSTTService(STTService): await asyncio.sleep(1) # Brief delay before reconnecting self._stream_start_time = int(time.time() * 1000) - finally: - self.reset_watchdog() except Exception as e: logger.error(f"Error in streaming task: {e}") @@ -804,17 +799,13 @@ class GoogleSTTService(STTService): async def _process_responses(self, streaming_recognize): """Process streaming recognition responses.""" try: - async for response in streaming_recognize: - self.start_watchdog() - + async for response in WatchdogAsyncIterator(streaming_recognize, reseter=self): # Check streaming limit if (int(time.time() * 1000) - self._stream_start_time) > self.STREAMING_LIMIT: logger.debug("Stream timeout reached in response processing") - self.reset_watchdog() break if not response.results: - self.reset_watchdog() continue for result in response.results: @@ -856,11 +847,8 @@ class GoogleSTTService(STTService): result=result, ) ) - - self.reset_watchdog() except Exception as e: logger.error(f"Error processing Google STT responses: {e}") - self.reset_watchdog() # Re-raise the exception to let it propagate (e.g. in the case of a # timeout, propagate to _stream_audio to reconnect) raise diff --git a/src/pipecat/services/openai/base_llm.py b/src/pipecat/services/openai/base_llm.py index 2badfed96..9651e0f99 100644 --- a/src/pipecat/services/openai/base_llm.py +++ b/src/pipecat/services/openai/base_llm.py @@ -36,6 +36,7 @@ from pipecat.processors.aggregators.openai_llm_context import ( from pipecat.processors.frame_processor import FrameDirection from pipecat.services.llm_service import FunctionCallFromLLM, LLMService from pipecat.utils.tracing.service_decorators import traced_llm +from pipecat.utils.watchdog_async_iterator import WatchdogAsyncIterator class BaseOpenAILLMService(LLMService): @@ -192,7 +193,7 @@ class BaseOpenAILLMService(LLMService): context ) - async for chunk in chunk_stream: + async for chunk in WatchdogAsyncIterator(chunk_stream, reseter=self): if chunk.usage: tokens = LLMTokenUsage( prompt_tokens=chunk.usage.prompt_tokens, diff --git a/src/pipecat/services/openai_realtime_beta/openai.py b/src/pipecat/services/openai_realtime_beta/openai.py index 7f459000a..49e3383e7 100644 --- a/src/pipecat/services/openai_realtime_beta/openai.py +++ b/src/pipecat/services/openai_realtime_beta/openai.py @@ -52,7 +52,8 @@ from pipecat.services.llm_service import FunctionCallFromLLM, LLMService from pipecat.services.openai.llm import OpenAIContextAggregatorPair from pipecat.transcriptions.language import Language from pipecat.utils.time import time_now_iso8601 -from pipecat.utils.tracing.service_decorators import traced_openai_realtime, traced_stt, traced_tts +from pipecat.utils.tracing.service_decorators import traced_openai_realtime, traced_stt +from pipecat.utils.watchdog_async_iterator import WatchdogAsyncIterator from . import events from .context import ( @@ -369,8 +370,7 @@ class OpenAIRealtimeBetaLLMService(LLMService): # async def _receive_task_handler(self): - async for message in self._websocket: - self.start_watchdog() + async for message in WatchdogAsyncIterator(self._websocket, reseter=self): evt = events.parse_server_event(message) if evt.type == "session.created": await self._handle_evt_session_created(evt) @@ -401,7 +401,6 @@ class OpenAIRealtimeBetaLLMService(LLMService): await self._handle_evt_error(evt) # errors are fatal, so exit the receive loop return - self.reset_watchdog() @traced_openai_realtime(operation="llm_setup") async def _handle_evt_session_created(self, evt): diff --git a/src/pipecat/services/riva/stt.py b/src/pipecat/services/riva/stt.py index 91c207c8b..16d9528c5 100644 --- a/src/pipecat/services/riva/stt.py +++ b/src/pipecat/services/riva/stt.py @@ -23,6 +23,7 @@ from pipecat.services.stt_service import SegmentedSTTService, STTService from pipecat.transcriptions.language import Language from pipecat.utils.time import time_now_iso8601 from pipecat.utils.tracing.service_decorators import traced_stt +from pipecat.utils.watchdog_queue import WatchdogQueue try: import riva.client @@ -198,7 +199,7 @@ class RivaSTTService(STTService): self._thread_task = self.create_task(self._thread_task_handler()) if not self._response_task: - self._response_queue = asyncio.Queue() + self._response_queue = WatchdogQueue(self) self._response_task = self.create_task(self._response_task_handler()) async def stop(self, frame: EndFrame): @@ -224,13 +225,12 @@ class RivaSTTService(STTService): streaming_config=self._config, ) for response in responses: - self.start_watchdog() + self.reset_watchdog() if not response.results: continue asyncio.run_coroutine_threadsafe( self._response_queue.put(response), self.get_event_loop() ) - self.reset_watchdog() async def _thread_task_handler(self): try: @@ -285,9 +285,8 @@ class RivaSTTService(STTService): async def _response_task_handler(self): while True: response = await self._response_queue.get() - self.start_watchdog() await self._handle_response(response) - self.reset_watchdog() + self._response_queue.task_done() async def run_stt(self, audio: bytes) -> AsyncGenerator[Frame, None]: await self.start_ttfb_metrics() diff --git a/src/pipecat/services/sambanova/llm.py b/src/pipecat/services/sambanova/llm.py index 01a8d294c..6c44215a0 100644 --- a/src/pipecat/services/sambanova/llm.py +++ b/src/pipecat/services/sambanova/llm.py @@ -19,6 +19,7 @@ from pipecat.processors.aggregators.openai_llm_context import OpenAILLMContext from pipecat.services.llm_service import FunctionCallFromLLM from pipecat.services.openai.llm import OpenAILLMService from pipecat.utils.tracing.service_decorators import traced_llm +from pipecat.utils.watchdog_async_iterator import WatchdogAsyncIterator class SambaNovaLLMService(OpenAILLMService): # type: ignore @@ -94,7 +95,7 @@ class SambaNovaLLMService(OpenAILLMService): # type: ignore context ) - async for chunk in chunk_stream: + async for chunk in WatchdogAsyncIterator(chunk_stream, reseter=self): if chunk.usage: tokens = LLMTokenUsage( prompt_tokens=chunk.usage.prompt_tokens, diff --git a/src/pipecat/services/simli/video.py b/src/pipecat/services/simli/video.py index 76381b9c6..19774ab4f 100644 --- a/src/pipecat/services/simli/video.py +++ b/src/pipecat/services/simli/video.py @@ -18,6 +18,7 @@ from pipecat.frames.frames import ( TTSAudioRawFrame, ) from pipecat.processors.frame_processor import FrameDirection, FrameProcessor, StartFrame +from pipecat.utils.watchdog_async_iterator import WatchdogAsyncIterator try: from av.audio.frame import AudioFrame @@ -61,8 +62,8 @@ class SimliVideoService(FrameProcessor): async def _consume_and_process_audio(self): await self._pipecat_resampler_event.wait() - async for audio_frame in self._simli_client.getAudioStreamIterator(): - self.start_watchdog() + audio_iterator = self._simli_client.getAudioStreamIterator() + async for audio_frame in WatchdogAsyncIterator(audio_iterator, reseter=self): resampled_frames = self._pipecat_resampler.resample(audio_frame) for resampled_frame in resampled_frames: audio_array = resampled_frame.to_ndarray() @@ -75,12 +76,11 @@ class SimliVideoService(FrameProcessor): num_channels=1, ), ) - self.reset_watchdog() async def _consume_and_process_video(self): await self._pipecat_resampler_event.wait() - async for video_frame in self._simli_client.getVideoStreamIterator(targetFormat="rgb24"): - self.start_watchdog() + video_iterator = self._simli_client.getVideoStreamIterator(targetFormat="rgb24") + async for video_frame in WatchdogAsyncIterator(video_iterator, reseter=self): # Process the video frame convertedFrame: OutputImageRawFrame = OutputImageRawFrame( image=video_frame.to_rgb().to_image().tobytes(), @@ -89,7 +89,6 @@ class SimliVideoService(FrameProcessor): ) convertedFrame.pts = video_frame.pts await self.push_frame(convertedFrame) - self.reset_watchdog() async def process_frame(self, frame: Frame, direction: FrameDirection): await super().process_frame(frame, direction) diff --git a/src/pipecat/services/tavus/video.py b/src/pipecat/services/tavus/video.py index 4ec744c51..41202d0e5 100644 --- a/src/pipecat/services/tavus/video.py +++ b/src/pipecat/services/tavus/video.py @@ -27,6 +27,7 @@ from pipecat.frames.frames import ( from pipecat.processors.frame_processor import FrameDirection, FrameProcessorSetup from pipecat.services.ai_service import AIService from pipecat.transports.services.tavus import TavusCallbacks, TavusParams, TavusTransportClient +from pipecat.utils.watchdog_queue import WatchdogQueue class TavusVideoService(AIService): @@ -71,7 +72,7 @@ class TavusVideoService(AIService): self._resampler = create_default_resampler() self._audio_buffer = bytearray() - self._queue = asyncio.Queue() + self._queue = WatchdogQueue(self) self._send_task: Optional[asyncio.Task] = None # This is the custom track destination expected by Tavus self._transport_destination: Optional[str] = "stream" @@ -188,7 +189,7 @@ class TavusVideoService(AIService): async def _create_send_task(self): if not self._send_task: - self._queue = asyncio.Queue() + self._queue = WatchdogQueue(self) self._send_task = self.create_task(self._send_task_handler()) async def _cancel_send_task(self): @@ -217,7 +218,6 @@ class TavusVideoService(AIService): async def _send_task_handler(self): while True: frame = await self._queue.get() - self.start_watchdog() if isinstance(frame, OutputAudioRawFrame) and self._client: await self._client.write_audio_frame(frame) - self.reset_watchdog() + self._queue.task_done() diff --git a/src/pipecat/services/tts_service.py b/src/pipecat/services/tts_service.py index 904f603a9..a623c9531 100644 --- a/src/pipecat/services/tts_service.py +++ b/src/pipecat/services/tts_service.py @@ -39,6 +39,7 @@ from pipecat.utils.text.base_text_aggregator import BaseTextAggregator from pipecat.utils.text.base_text_filter import BaseTextFilter from pipecat.utils.text.simple_text_aggregator import SimpleTextAggregator from pipecat.utils.time import seconds_to_nanoseconds +from pipecat.utils.watchdog_queue import WatchdogQueue class TTSService(AIService): @@ -315,6 +316,8 @@ class TTSService(AIService): if has_started: await self.push_frame(TTSStoppedFrame()) has_started = False + finally: + self.reset_watchdog() class WordTTSService(TTSService): @@ -327,7 +330,7 @@ class WordTTSService(TTSService): def __init__(self, **kwargs): super().__init__(**kwargs) self._initial_word_timestamp = -1 - self._words_queue = asyncio.Queue() + self._words_queue = WatchdogQueue(self) self._words_task = None self._llm_response_started: bool = False @@ -578,7 +581,7 @@ class AudioContextWordTTSService(WebsocketWordTTSService): def _create_audio_context_task(self): if not self._audio_context_task: - self._contexts_queue = asyncio.Queue() + self._contexts_queue = WatchdogQueue(self) self._contexts: Dict[str, asyncio.Queue] = {} self._audio_context_task = self.create_task(self._audio_context_task_handler()) @@ -620,10 +623,12 @@ class AudioContextWordTTSService(WebsocketWordTTSService): while running: try: frame = await asyncio.wait_for(queue.get(), timeout=AUDIO_CONTEXT_TIMEOUT) + self.reset_watchdog() if frame: await self.push_frame(frame) running = frame is not None except asyncio.TimeoutError: + self.reset_watchdog() # We didn't get audio, so let's consider this context finished. logger.trace(f"{self} time out on audio context {context_id}") break diff --git a/src/pipecat/transports/base_input.py b/src/pipecat/transports/base_input.py index d707c5969..93cd90825 100644 --- a/src/pipecat/transports/base_input.py +++ b/src/pipecat/transports/base_input.py @@ -368,8 +368,6 @@ class BaseInputTransport(FrameProcessor): self._audio_in_queue.get(), timeout=AUDIO_INPUT_TIMEOUT_SECS ) - self.start_watchdog() - # If an audio filter is available, run it before VAD. if self._params.audio_in_filter: frame.audio = await self._params.audio_in_filter.filter(frame.audio) diff --git a/src/pipecat/transports/base_output.py b/src/pipecat/transports/base_output.py index 386f223d7..2b37e7ddf 100644 --- a/src/pipecat/transports/base_output.py +++ b/src/pipecat/transports/base_output.py @@ -40,6 +40,7 @@ from pipecat.frames.frames import ( from pipecat.processors.frame_processor import FrameDirection, FrameProcessor from pipecat.transports.base_transport import TransportParams from pipecat.utils.time import nanoseconds_to_seconds +from pipecat.utils.watchdog_priority_queue import WatchdogPriorityQueue BOT_VAD_STOP_SECS = 0.35 @@ -441,8 +442,10 @@ class BaseOutputTransport(FrameProcessor): frame = await asyncio.wait_for( self._audio_queue.get(), timeout=vad_stop_secs ) + self._transport.reset_watchdog() yield frame except asyncio.TimeoutError: + self._transport.reset_watchdog() # Notify the bot stopped speaking upstream if necessary. await self._bot_stopped_speaking() @@ -452,11 +455,13 @@ class BaseOutputTransport(FrameProcessor): while True: try: frame = self._audio_queue.get_nowait() + self._transport.reset_watchdog() if isinstance(frame, OutputAudioRawFrame): frame.audio = await self._mixer.mix(frame.audio) last_frame_time = time.time() yield frame except asyncio.QueueEmpty: + self._transport.reset_watchdog() # Notify the bot stopped speaking upstream if necessary. diff_time = time.time() - last_frame_time if diff_time > vad_stop_secs: @@ -597,7 +602,7 @@ class BaseOutputTransport(FrameProcessor): def _create_clock_task(self): if not self._clock_task: - self._clock_queue = asyncio.PriorityQueue() + self._clock_queue = WatchdogPriorityQueue(self._transport) self._clock_task = self._transport.create_task(self._clock_task_handler()) async def _cancel_clock_task(self): diff --git a/src/pipecat/transports/network/fastapi_websocket.py b/src/pipecat/transports/network/fastapi_websocket.py index 7f6256b83..710c715a9 100644 --- a/src/pipecat/transports/network/fastapi_websocket.py +++ b/src/pipecat/transports/network/fastapi_websocket.py @@ -26,11 +26,12 @@ from pipecat.frames.frames import ( TransportMessageFrame, TransportMessageUrgentFrame, ) -from pipecat.processors.frame_processor import FrameDirection, FrameProcessorSetup +from pipecat.processors.frame_processor import FrameDirection from pipecat.serializers.base_serializer import FrameSerializer, FrameSerializerType from pipecat.transports.base_input import BaseInputTransport from pipecat.transports.base_output import BaseOutputTransport from pipecat.transports.base_transport import BaseTransport, TransportParams +from pipecat.utils.watchdog_async_iterator import WatchdogAsyncIterator try: from fastapi import WebSocket @@ -178,12 +179,10 @@ class FastAPIWebsocketInputTransport(BaseInputTransport): async def _receive_messages(self): try: - async for message in self._client.receive(): + async for message in WatchdogAsyncIterator(self._client.receive(), reseter=self): if not self._params.serializer: continue - self.start_watchdog() - frame = await self._params.serializer.deserialize(message) if not frame: @@ -193,13 +192,9 @@ class FastAPIWebsocketInputTransport(BaseInputTransport): await self.push_audio_frame(frame) else: await self.push_frame(frame) - - self.reset_watchdog() except Exception as e: logger.error(f"{self} exception receiving data: {e.__class__.__name__} ({e})") - self.reset_watchdog() - await self._client.trigger_client_disconnected() async def _monitor_websocket(self): diff --git a/src/pipecat/transports/network/small_webrtc.py b/src/pipecat/transports/network/small_webrtc.py index 5b36fec37..48ca9d53f 100644 --- a/src/pipecat/transports/network/small_webrtc.py +++ b/src/pipecat/transports/network/small_webrtc.py @@ -33,6 +33,7 @@ from pipecat.transports.base_input import BaseInputTransport from pipecat.transports.base_output import BaseOutputTransport from pipecat.transports.base_transport import BaseTransport, TransportParams from pipecat.transports.network.webrtc_connection import SmallWebRTCConnection +from pipecat.utils.watchdog_async_iterator import WatchdogAsyncIterator try: import cv2 @@ -422,19 +423,18 @@ class SmallWebRTCInputTransport(BaseInputTransport): async def _receive_audio(self): try: - async for audio_frame in self._client.read_audio_frame(): - self.start_watchdog() + audio_iterator = self._client.read_audio_frame() + async for audio_frame in WatchdogAsyncIterator(audio_iterator, reseter=self): if audio_frame: await self.push_audio_frame(audio_frame) - self.reset_watchdog() except Exception as e: logger.error(f"{self} exception receiving data: {e.__class__.__name__} ({e})") async def _receive_video(self): try: - async for video_frame in self._client.read_video_frame(): - self.start_watchdog() + video_iterator = self._client.read_video_frame() + async for video_frame in WatchdogAsyncIterator(video_iterator, reseter=self): if video_frame: await self.push_video_frame(video_frame) @@ -453,7 +453,6 @@ class SmallWebRTCInputTransport(BaseInputTransport): await self.push_video_frame(image_frame) # Remove from pending requests del self._image_requests[req_id] - self.reset_watchdog() except Exception as e: logger.error(f"{self} exception receiving data: {e.__class__.__name__} ({e})") diff --git a/src/pipecat/transports/services/daily.py b/src/pipecat/transports/services/daily.py index 776da1693..61afcb6c9 100644 --- a/src/pipecat/transports/services/daily.py +++ b/src/pipecat/transports/services/daily.py @@ -40,6 +40,8 @@ from pipecat.transports.base_input import BaseInputTransport from pipecat.transports.base_output import BaseOutputTransport from pipecat.transports.base_transport import BaseTransport, TransportParams from pipecat.utils.asyncio import BaseTaskManager +from pipecat.utils.watchdog_queue import WatchdogQueue +from pipecat.utils.watchdog_reseter import WatchdogReseter try: from daily import ( @@ -250,7 +252,7 @@ class DailyAudioTrack: track: CustomAudioTrack -class DailyTransportClient(EventHandler): +class DailyTransportClient(WatchdogReseter, EventHandler): """Core client for interacting with Daily's API. Manages the connection to Daily rooms and handles all low-level API interactions. @@ -320,9 +322,9 @@ class DailyTransportClient(EventHandler): # waits for it to finish using completions (and a future) we will # deadlock because completions use event handlers (which are holding the # GIL). - self._event_queue = asyncio.Queue() - self._audio_queue = asyncio.Queue() - self._video_queue = asyncio.Queue() + self._event_queue = WatchdogQueue(self) + self._audio_queue = WatchdogQueue(self) + self._video_queue = WatchdogQueue(self) self._event_task = None self._audio_task = None self._video_task = None @@ -395,6 +397,10 @@ class DailyTransportClient(EventHandler): if not frame.transport_destination and self._camera: self._camera.write_frame(frame.image) + def reset_watchdog(self): + if self._task_manager: + self._task_manager.reset_watchdog(asyncio.current_task()) + async def setup(self, setup: FrameProcessorSetup): if self._task_manager: return @@ -934,6 +940,7 @@ class DailyTransportClient(EventHandler): await self._joined_event.wait() (callback, *args) = await queue.get() await callback(*args) + queue.task_done() def _get_event_loop(self) -> asyncio.AbstractEventLoop: if not self._task_manager: diff --git a/src/pipecat/transports/services/livekit.py b/src/pipecat/transports/services/livekit.py index ecb2718c8..5a4c38069 100644 --- a/src/pipecat/transports/services/livekit.py +++ b/src/pipecat/transports/services/livekit.py @@ -28,6 +28,7 @@ from pipecat.transports.base_input import BaseInputTransport from pipecat.transports.base_output import BaseOutputTransport from pipecat.transports.base_transport import BaseTransport, TransportParams from pipecat.utils.asyncio import BaseTaskManager +from pipecat.utils.watchdog_async_iterator import WatchdogAsyncIterator try: from livekit import rtc @@ -341,8 +342,9 @@ class LiveKitTransportClient: logger.warning(f"Received unexpected event type: {type(event)}") async def get_next_audio_frame(self): - frame, participant_id = await self._audio_queue.get() - return frame, participant_id + while True: + frame, participant_id = await self._audio_queue.get() + yield frame, participant_id def __str__(self): return f"{self._transport_name}::LiveKitTransportClient" @@ -413,9 +415,8 @@ class LiveKitInputTransport(BaseInputTransport): async def _audio_in_task_handler(self): logger.info("Audio input task started") - while True: - audio_data = await self._client.get_next_audio_frame() - self.start_watchdog() + audio_iterator = self._client.get_next_audio_frame() + async for audio_data in WatchdogAsyncIterator(audio_iterator, reseter=self): if audio_data: audio_frame_event, participant_id = audio_data pipecat_audio_frame = await self._convert_livekit_audio_to_pipecat( @@ -428,7 +429,6 @@ class LiveKitInputTransport(BaseInputTransport): num_channels=pipecat_audio_frame.num_channels, ) await self.push_audio_frame(input_audio_frame) - self.reset_watchdog() async def _convert_livekit_audio_to_pipecat( self, audio_frame_event: rtc.AudioFrameEvent diff --git a/src/pipecat/utils/asyncio.py b/src/pipecat/utils/asyncio.py index 36479edd4..a2c75ceb6 100644 --- a/src/pipecat/utils/asyncio.py +++ b/src/pipecat/utils/asyncio.py @@ -27,10 +27,6 @@ class BaseTaskManager(ABC): def setup(self, params: TaskManagerParams): pass - @abstractmethod - async def cleanup(self): - pass - @abstractmethod def get_event_loop(self) -> asyncio.AbstractEventLoop: pass @@ -97,14 +93,6 @@ class BaseTaskManager(ABC): """Returns the list of currently created/registered tasks.""" pass - @abstractmethod - def start_watchdog(self, task: asyncio.Task): - """Starts the given task watchdog timer. If not reset, a warning will be - logged indicating the task is stalling. - - """ - pass - @abstractmethod def reset_watchdog(self, task: asyncio.Task): """Resets the given task watchdog timer. If not reset, a warning will be @@ -117,31 +105,21 @@ class BaseTaskManager(ABC): @dataclass class TaskData: task: asyncio.Task - watchdog_start: asyncio.Event watchdog_timer: asyncio.Event enable_watchdog_logging: bool watchdog_timeout: float + watchdog_task: Optional[asyncio.Task] class TaskManager(BaseTaskManager): def __init__(self) -> None: self._tasks: Dict[str, TaskData] = {} self._params: Optional[TaskManagerParams] = None - self._watchdog_tasks: List[asyncio.Task] = [] def setup(self, params: TaskManagerParams): if not self._params: self._params = params - async def cleanup(self): - for task in self._watchdog_tasks: - try: - task.cancel() - await task - except asyncio.CancelledError: - # This is expected, no need to re-raise. - pass - def get_event_loop(self) -> asyncio.AbstractEventLoop: if not self._params: raise Exception("TaskManager is not setup: unable to get event loop") @@ -189,7 +167,6 @@ class TaskManager(BaseTaskManager): self._add_task( TaskData( task=task, - watchdog_start=asyncio.Event(), watchdog_timer=asyncio.Event(), enable_watchdog_logging=( enable_watchdog_logging @@ -199,6 +176,7 @@ class TaskManager(BaseTaskManager): watchdog_timeout=( watchdog_timeout if watchdog_timeout else self._params.watchdog_timeout ), + watchdog_task=None, ) ) logger.trace(f"{name}: task created") @@ -231,7 +209,7 @@ class TaskManager(BaseTaskManager): except Exception as e: logger.exception(f"{name}: unexpected exception while stopping task: {e}") finally: - self._remove_task(task) + await self._remove_task(task) async def cancel_task(self, task: asyncio.Task, timeout: Optional[float] = None): """Cancels the given asyncio Task and awaits its completion with an @@ -265,36 +243,19 @@ class TaskManager(BaseTaskManager): logger.critical(f"{name}: fatal base exception while cancelling task: {e}") raise finally: - self._remove_task(task) + await self._remove_task(task) def current_tasks(self) -> Sequence[asyncio.Task]: """Returns the list of currently created/registered tasks.""" return [data.task for data in self._tasks.values()] - def start_watchdog(self, task: asyncio.Task): - """Starts the given task watchdog timer. If not reset, a warning will be - logged indicating the task is stalling. If the timer was already started - a warning will be logged. - - """ - name = task.get_name() - if name in self._tasks: - if self._tasks[name].watchdog_start.is_set(): - logger.warning(f"Watchdog timer for task {name} already started") - else: - self._tasks[name].watchdog_timer.clear() - self._tasks[name].watchdog_start.set() - else: - logger.warning(f"Unable to start watchdog timer: task {name} does not exist") - def reset_watchdog(self, task: asyncio.Task): - """Resets the given task watchdog timer. If not reset, a warning will be - logged indicating the task is stalling. + """Resets the given task watchdog timer. If not reset on time, a warning + will be logged indicating the task is stalling. """ name = task.get_name() if name in self._tasks: - self._tasks[name].watchdog_start.clear() self._tasks[name].watchdog_timer.set() else: logger.warning(f"Unable to reset watchdog timer: task {name} does not exist") @@ -302,44 +263,40 @@ class TaskManager(BaseTaskManager): def _add_task(self, task_data: TaskData): name = task_data.task.get_name() self._tasks[name] = task_data - watchdog_task = self.get_event_loop().create_task( - self._watchdog_task_handler(self._tasks[name]) - ) - self._watchdog_tasks.append(watchdog_task) + watchdog_task = self.get_event_loop().create_task(self._watchdog_task_handler(task_data)) + task_data.watchdog_task = watchdog_task - def _remove_task(self, task: asyncio.Task): + async def _remove_task(self, task: asyncio.Task): name = task.get_name() try: + task_data = self._tasks[name] + if task_data.watchdog_task: + try: + task_data.watchdog_task.cancel() + await task_data.watchdog_task + except asyncio.CancelledError: + pass + task_data.watchdog_task = None del self._tasks[name] except KeyError as e: logger.trace(f"{name}: unable to remove task (already removed?): {e}") async def _watchdog_task_handler(self, task_data: TaskData): name = task_data.task.get_name() - start = task_data.watchdog_start timer = task_data.watchdog_timer enable_watchdog_logging = task_data.enable_watchdog_logging watchdog_timeout = task_data.watchdog_timeout - async def wait_for_reset(): - waiting = True - while waiting: - try: - start_time = time.time() - await asyncio.wait_for(timer.wait(), timeout=watchdog_timeout) - total_time = time.time() - start_time - if enable_watchdog_logging: - logger.debug(f"{name} task processing time: {total_time:.20f}") - waiting = False - except asyncio.TimeoutError: - logger.warning( - f"{name}: task is taking too long {WATCHDOG_TIMEOUT} second(s) (forgot to reset watchdog?)" - ) - finally: - timer.clear() - while True: - # Wait for the user to start the watchdog timer. - await start.wait() - # Now, waiting for the task to finish. - await wait_for_reset() + try: + start_time = time.time() + await asyncio.wait_for(timer.wait(), timeout=watchdog_timeout) + total_time = time.time() - start_time + if enable_watchdog_logging: + logger.debug(f"{name} task processing time: {total_time:.20f}") + except asyncio.TimeoutError: + logger.warning( + f"{name}: task is taking too long {WATCHDOG_TIMEOUT} second(s) (forgot to reset watchdog?)" + ) + finally: + timer.clear() diff --git a/src/pipecat/utils/watchdog_async_iterator.py b/src/pipecat/utils/watchdog_async_iterator.py new file mode 100644 index 000000000..db126a363 --- /dev/null +++ b/src/pipecat/utils/watchdog_async_iterator.py @@ -0,0 +1,60 @@ +# +# Copyright (c) 2024–2025, Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + +import asyncio +from typing import AsyncIterator, Optional + +from pipecat.utils.watchdog_reseter import WatchdogReseter + + +class WatchdogAsyncIterator: + """An asynchronous iterator that monitors activity and resets the current + task watchdog timer. This is necessary to avoid task watchdog timers to + expire while we are waiting to get an item from the iterator. + + """ + + def __init__(self, async_iterable, *, reseter: WatchdogReseter, timeout: float = 2.0): + self._async_iterable = async_iterable + self._reseter = reseter + self._timeout = timeout + self._iter: Optional[AsyncIterator] = None + self._current_anext_task: Optional[asyncio.Task] = None + + def __aiter__(self): + return self + + async def __anext__(self): + if not self._iter: + self._iter = await self._ensure_async_iterator(self._async_iterable) + + while True: + try: + if not self._current_anext_task: + self._current_anext_task = asyncio.create_task(self._iter.__anext__()) + + item = await asyncio.wait_for( + asyncio.shield(self._current_anext_task), + timeout=self._timeout, + ) + + self._reseter.reset_watchdog() + + # The task has finish, so we will create a new one for th next item. + self._current_anext_task = None + + return item + except asyncio.TimeoutError: + self._reseter.reset_watchdog() + except StopAsyncIteration: + self._current_anext_task = None + raise + + async def _ensure_async_iterator(self, obj) -> AsyncIterator: + aiter = obj.__aiter__() + if asyncio.iscoroutine(aiter): + aiter = await aiter + return aiter diff --git a/src/pipecat/utils/watchdog_event.py b/src/pipecat/utils/watchdog_event.py new file mode 100644 index 000000000..3165c31df --- /dev/null +++ b/src/pipecat/utils/watchdog_event.py @@ -0,0 +1,31 @@ +# +# Copyright (c) 2024–2025, Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + +import asyncio + +from pipecat.utils.watchdog_reseter import WatchdogReseter + + +class WatchdogEvent(asyncio.Event): + """An asynchronous event that resets the current task watchdog timer. This + is necessary to avoid task watchdog timers to expire while we are waiting on + the event. + + """ + + def __init__(self, reseter: WatchdogReseter, timeout: float = 2.0) -> None: + super().__init__() + self._reseter = reseter + self._timeout = timeout + + async def wait(self): + while True: + try: + await asyncio.wait_for(super().wait(), timeout=self._timeout) + self._reseter.reset_watchdog() + return True + except asyncio.TimeoutError: + self._reseter.reset_watchdog() diff --git a/src/pipecat/utils/watchdog_priority_queue.py b/src/pipecat/utils/watchdog_priority_queue.py new file mode 100644 index 000000000..34782fbf8 --- /dev/null +++ b/src/pipecat/utils/watchdog_priority_queue.py @@ -0,0 +1,35 @@ +# +# Copyright (c) 2024–2025, Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + +import asyncio + +from pipecat.utils.watchdog_reseter import WatchdogReseter + + +class WatchdogPriorityQueue(asyncio.PriorityQueue): + """An asynchronous priority queue that resets the current task watchdog + timer. This is necessary to avoid task watchdog timers to expire while we + are waiting to get an item from the queue. + + """ + + def __init__(self, reseter: WatchdogReseter, maxsize: int = 0, timeout: float = 2.0) -> None: + super().__init__(maxsize) + self._reseter = reseter + self._timeout = timeout + + async def get(self): + while True: + try: + item = await asyncio.wait_for(super().get(), timeout=self._timeout) + self._reseter.reset_watchdog() + return item + except asyncio.TimeoutError: + self._reseter.reset_watchdog() + + def task_done(self): + self._reseter.reset_watchdog() + super().task_done() diff --git a/src/pipecat/utils/watchdog_queue.py b/src/pipecat/utils/watchdog_queue.py new file mode 100644 index 000000000..0d9ffd7a8 --- /dev/null +++ b/src/pipecat/utils/watchdog_queue.py @@ -0,0 +1,35 @@ +# +# Copyright (c) 2024–2025, Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + +import asyncio + +from pipecat.utils.watchdog_reseter import WatchdogReseter + + +class WatchdogQueue(asyncio.Queue): + """An asynchronous queue that resets the current task watchdog timer. This + is necessary to avoid task watchdog timers to expire while we are waiting to + get an item from the queue. + + """ + + def __init__(self, reseter: WatchdogReseter, maxsize: int = 0, timeout: float = 2.0) -> None: + super().__init__(maxsize) + self._reseter = reseter + self._timeout = timeout + + async def get(self): + while True: + try: + item = await asyncio.wait_for(super().get(), timeout=self._timeout) + self._reseter.reset_watchdog() + return item + except asyncio.TimeoutError: + self._reseter.reset_watchdog() + + def task_done(self): + self._reseter.reset_watchdog() + super().task_done() diff --git a/src/pipecat/utils/watchdog_reseter.py b/src/pipecat/utils/watchdog_reseter.py new file mode 100644 index 000000000..ee70207b3 --- /dev/null +++ b/src/pipecat/utils/watchdog_reseter.py @@ -0,0 +1,13 @@ +# +# Copyright (c) 2024–2025, Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + +from abc import ABC, abstractmethod + + +class WatchdogReseter(ABC): + @abstractmethod + def reset_watchdog(self): + pass