add start_watchdog/reset_watchdog to tasks

This commit is contained in:
Aleix Conchillo Flaqué
2025-06-23 15:26:34 -07:00
parent 5a3457ba33
commit 076a8938f0
19 changed files with 86 additions and 10 deletions

View File

@@ -202,14 +202,18 @@ 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.
@@ -224,3 +228,5 @@ class ParallelPipeline(BasePipeline):
running = not (endframe_counter == 0 and isinstance(frame, EndFrame))
self._down_queue.task_done()
self.reset_watchdog()

View File

@@ -61,5 +61,7 @@ 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()

View File

@@ -412,6 +412,8 @@ class FrameProcessor(BaseObject):
(frame, direction, callback) = await self.__input_queue.get()
self.start_watchdog()
# Process the frame.
await self.process_frame(frame, direction)
@@ -421,6 +423,8 @@ class FrameProcessor(BaseObject):
self.__input_queue.task_done()
self.reset_watchdog()
def __create_push_task(self):
if not self.__push_frame_task:
self.__push_queue = asyncio.Queue()
@@ -434,5 +438,7 @@ 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()

View File

@@ -783,14 +783,18 @@ 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:

View File

@@ -190,6 +190,7 @@ class AssemblyAISTTService(STTService):
while self._connected:
try:
message = await self._websocket.recv()
self.start_watchdog()
data = json.loads(message)
await self._handle_message(data)
except websockets.exceptions.ConnectionClosedOK:
@@ -197,6 +198,8 @@ class AssemblyAISTTService(STTService):
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}")

View File

@@ -285,6 +285,9 @@ class AWSTranscribeSTTService(STTService):
try:
response = await self._ws_client.recv()
self.start_watchdog()
headers, payload = decode_event(response)
if headers.get(":message-type") == "event":
@@ -342,3 +345,5 @@ class AWSTranscribeSTTService(STTService):
except Exception as e:
logger.error(f"{self} Unexpected error in receive loop: {e}")
break
finally:
self.reset_watchdog()

View File

@@ -699,6 +699,8 @@ class AWSNovaSonicLLMService(LLMService):
output = await self._stream.await_output()
result = await output[1].receive()
self.start_watchdog()
if result.value and result.value.bytes_:
response_data = result.value.bytes_.decode("utf-8")
json_data = json.loads(response_data)
@@ -731,6 +733,8 @@ class AWSNovaSonicLLMService(LLMService):
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

View File

@@ -687,6 +687,8 @@ class GeminiMultimodalLiveLLMService(LLMService):
async def _receive_task_handler(self):
async for message in self._websocket:
self.start_watchdog()
evt = events.parse_server_event(message)
# logger.debug(f"Received event: {message[:500]}")
# logger.debug(f"Received event: {evt}")
@@ -708,8 +710,8 @@ class GeminiMultimodalLiveLLMService(LLMService):
await self._handle_evt_error(evt)
# errors are fatal, so exit the receive loop
return
else:
pass
self.reset_watchdog()
#
#

View File

@@ -502,6 +502,8 @@ class GladiaSTTService(STTService):
async def _receive_task_handler(self):
try:
async for message in self._websocket:
self.start_watchdog()
content = json.loads(message)
# Handle audio chunk acknowledgments
@@ -559,11 +561,15 @@ class GladiaSTTService(STTService):
translation, "", time_now_iso8601(), translated_language
)
)
self.reset_watchdog()
except websockets.exceptions.ConnectionClosed:
# Expected when closing the connection
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."""

View File

@@ -747,9 +747,12 @@ 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)
self.reset_watchdog()
continue
# Start bi-directional streaming
@@ -760,12 +763,13 @@ 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")
# Reset stream start time
self._stream_start_time = int(time.time() * 1000)
continue
else:
# Normal stream end
break
@@ -775,7 +779,8 @@ class GoogleSTTService(STTService):
await asyncio.sleep(1) # Brief delay before reconnecting
self._stream_start_time = int(time.time() * 1000)
continue
finally:
self.reset_watchdog()
except Exception as e:
logger.error(f"Error in streaming task: {e}")
@@ -800,12 +805,16 @@ class GoogleSTTService(STTService):
"""Process streaming recognition responses."""
try:
async for response in streaming_recognize:
self.start_watchdog()
# 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:
@@ -848,8 +857,10 @@ class GoogleSTTService(STTService):
)
)
self.reset_watchdog()
except Exception as e:
logger.error(f"Error processing Google STT responses: {e}")
# Re-raise the exception to let it propagate (e.g. in the case of a timeout, propagate to _stream_audio to reconnect)
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

View File

@@ -203,12 +203,11 @@ class ResponseCancelEvent(ClientEvent):
class ServerEvent(BaseModel):
model_config = ConfigDict(arbitrary_types_allowed=True)
event_id: str
type: str
class Config:
arbitrary_types_allowed = True
class SessionCreatedEvent(ServerEvent):
type: Literal["session.created"]

View File

@@ -370,6 +370,7 @@ class OpenAIRealtimeBetaLLMService(LLMService):
async def _receive_task_handler(self):
async for message in self._websocket:
self.start_watchdog()
evt = events.parse_server_event(message)
if evt.type == "session.created":
await self._handle_evt_session_created(evt)
@@ -400,6 +401,7 @@ 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):

View File

@@ -224,11 +224,13 @@ class RivaSTTService(STTService):
streaming_config=self._config,
)
for response in responses:
self.start_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:
@@ -283,7 +285,9 @@ 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()
async def run_stt(self, audio: bytes) -> AsyncGenerator[Frame, None]:
await self.start_ttfb_metrics()

View File

@@ -62,6 +62,7 @@ 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()
resampled_frames = self._pipecat_resampler.resample(audio_frame)
for resampled_frame in resampled_frames:
audio_array = resampled_frame.to_ndarray()
@@ -74,10 +75,12 @@ 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()
# Process the video frame
convertedFrame: OutputImageRawFrame = OutputImageRawFrame(
image=video_frame.to_rgb().to_image().tobytes(),
@@ -86,6 +89,7 @@ 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)

View File

@@ -217,5 +217,7 @@ class TavusVideoService(AIService):
async def _send_task_handler(self):
while True:
frame = await self._queue.get()
if isinstance(frame, OutputAudioRawFrame):
self.start_watchdog()
if isinstance(frame, OutputAudioRawFrame) and self._client:
await self._client.write_audio_frame(frame)
self.reset_watchdog()

View File

@@ -357,6 +357,8 @@ class BaseInputTransport(FrameProcessor):
while True:
frame: InputAudioRawFrame = await self._audio_in_queue.get()
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)
@@ -376,6 +378,8 @@ class BaseInputTransport(FrameProcessor):
self._audio_in_queue.task_done()
self.reset_watchdog()
async def _handle_prediction_result(self, result: MetricsData):
"""Handle a prediction result event from the turn analyzer.

View File

@@ -171,6 +171,8 @@ class FastAPIWebsocketInputTransport(BaseInputTransport):
if not self._params.serializer:
continue
self.start_watchdog()
frame = await self._params.serializer.deserialize(message)
if not frame:
@@ -180,9 +182,13 @@ 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):

View File

@@ -423,8 +423,10 @@ class SmallWebRTCInputTransport(BaseInputTransport):
async def _receive_audio(self):
try:
async for audio_frame in self._client.read_audio_frame():
self.start_watchdog()
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})")
@@ -432,6 +434,7 @@ class SmallWebRTCInputTransport(BaseInputTransport):
async def _receive_video(self):
try:
async for video_frame in self._client.read_video_frame():
self.start_watchdog()
if video_frame:
await self.push_video_frame(video_frame)
@@ -450,6 +453,7 @@ 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})")

View File

@@ -415,6 +415,7 @@ class LiveKitInputTransport(BaseInputTransport):
logger.info("Audio input task started")
while True:
audio_data = await self._client.get_next_audio_frame()
self.start_watchdog()
if audio_data:
audio_frame_event, participant_id = audio_data
pipecat_audio_frame = await self._convert_livekit_audio_to_pipecat(
@@ -427,6 +428,7 @@ 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