wx44wx/pipecat
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #1118 from pipecat-ai/aleix/task-manager

Browse Source 
introduce TaskManager
This commit is contained in:
Aleix Conchillo Flaqué
2025-01-31 10:24:52 -08:00
committed by GitHub
parent a17243bc1e 002699f16c
commit bb9a2560c3
29 changed files with 498 additions and 336 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
CHANGELOG.md
View File

@@ -9,6 +9,19 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
### Added
- It is now possible to specify the asyncio event loop that a `PipelineTask` and
all the processors should run on by passing it as a new argument to the
`PipelineRunner`. This could allow running pipelines in multiple threads each
one with its own event loop.
- Added a new `utils.TaskManager`. Instead of a global task manager we now have
a task manager per `PipelineTask`. In the previous version the task manager
was global, so running multiple simultaneous `PipelineTask`s could result in
dangling task warnings which were not actually true. In order, for all the
processors to know about the task manager, we pass it through the
`StartFrame`. This means that processors should create tasks when they receive
a `StartFrame` but not before (because they don't have a task manager yet).
- Added `TelnyxFrameSerializer` to support Telnyx calls. A full running example
has also been added to `examples/telnyx-chatbot`.

2
examples/foundational/06-listen-and-respond.py
View File

@@ -38,6 +38,8 @@ logger.add(sys.stderr, level="DEBUG")
class MetricsLogger(FrameProcessor):
async def process_frame(self, frame: Frame, direction: FrameDirection):
await super().process_frame(frame, direction)
if isinstance(frame, MetricsFrame):
for d in frame.data:
if isinstance(d, TTFBMetricsData):

5
examples/foundational/22b-natural-conversation-proposal.py
View File

@@ -117,12 +117,15 @@ class CompletenessCheck(FrameProcessor):
async def process_frame(self, frame: Frame, direction: FrameDirection):
await super().process_frame(frame, direction)
if isinstance(frame, TextFrame) and frame.text == "YES":
logger.debug("Completeness check YES")
await self.push_frame(UserStoppedSpeakingFrame())
await self._notifier.notify()
elif isinstance(frame, TextFrame) and frame.text == "NO":
logger.debug("Completeness check NO")
else:
await self.push_frame(frame, direction)
class OutputGate(FrameProcessor):
@@ -166,7 +169,7 @@ class OutputGate(FrameProcessor):
async def _start(self):
self._frames_buffer = []
self._gate_task = self.get_event_loop().create_task(self._gate_task_handler())
self._gate_task = self.create_task(self._gate_task_handler())
async def _stop(self):
await self.cancel_task(self._gate_task)

4
examples/foundational/22c-natural-conversation-mixed-llms.py
View File

@@ -328,6 +328,8 @@ class CompletenessCheck(FrameProcessor):
await self._notifier.notify()
elif isinstance(frame, TextFrame) and frame.text == "NO":
logger.debug("!!! Completeness check NO")
else:
await self.push_frame(frame, direction)
class OutputGate(FrameProcessor):
@@ -371,7 +373,7 @@ class OutputGate(FrameProcessor):
async def _start(self):
self._frames_buffer = []
self._gate_task = self.get_event_loop().create_task(self._gate_task_handler())
self._gate_task = self.create_task(self._gate_task_handler())
async def _stop(self):
await self.cancel_task(self._gate_task)

6
examples/foundational/22d-natural-conversation-gemini-audio.py
View File

@@ -455,7 +455,9 @@ class CompletenessCheck(FrameProcessor):
else:
# logger.debug("!!! CompletenessCheck idle wait START")
self._wakeup_time = time.time() + self.wait_time
self._idle_task = self.get_event_loop().create_task(self._idle_task_handler())
self._idle_task = self.create_task(self._idle_task_handler())
else:
await self.push_frame(frame, direction)
async def _idle_task_handler(self):
try:
@@ -597,7 +599,7 @@ class OutputGate(FrameProcessor):
async def _start(self):
self._frames_buffer = []
self._gate_task = self.get_event_loop().create_task(self._gate_task_handler())
self._gate_task = self.create_task(self._gate_task_handler())
async def _stop(self):
await self.cancel_task(self._gate_task)

2
examples/foundational/25-google-audio-in.py
View File

@@ -212,7 +212,7 @@ class InputTranscriptionFrameEmitter(FrameProcessor):
elif isinstance(frame, LLMFullResponseEndFrame):
await self.push_frame(LLMDemoTranscriptionFrame(text=self._aggregation.strip()))
self._aggregation = ""
elif isinstance(frame, MetricsFrame):
else:
await self.push_frame(frame, direction)

2
src/pipecat/frames/frames.py
View File

@@ -12,6 +12,7 @@ from pipecat.audio.vad.vad_analyzer import VADParams
from pipecat.clocks.base_clock import BaseClock
from pipecat.metrics.metrics import MetricsData
from pipecat.transcriptions.language import Language
from pipecat.utils.asyncio import TaskManager
from pipecat.utils.time import nanoseconds_to_str
from pipecat.utils.utils import obj_count, obj_id
@@ -412,6 +413,7 @@ class StartFrame(SystemFrame):
"""This is the first frame that should be pushed down a pipeline."""
clock: BaseClock
task_manager: TaskManager
allow_interruptions: bool = False
enable_metrics: bool = False
enable_usage_metrics: bool = False

22
src/pipecat/pipeline/base_task.py
View File

@@ -4,6 +4,7 @@
# SPDX-License-Identifier: BSD 2-Clause License
#
import asyncio
from abc import ABC, abstractmethod
from typing import AsyncIterable, Iterable
@@ -23,6 +24,11 @@ class BaseTask(ABC):
"""Returns the name of this task."""
pass
@abstractmethod
def set_event_loop(self, loop: asyncio.AbstractEventLoop):
"""Sets the event loop that this task will run on."""
pass
@abstractmethod
def has_finished(self) -> bool:
"""Indicates whether the tasks has finished. That is, all processors
@@ -41,28 +47,20 @@ class BaseTask(ABC):
@abstractmethod
async def cancel(self):
"""
Stops the running pipeline immediately.
"""
"""Stops the running pipeline immediately."""
pass
@abstractmethod
async def run(self):
"""
Starts running the given pipeline.
"""
"""Starts running the given pipeline."""
pass
@abstractmethod
async def queue_frame(self, frame: Frame):
"""
Queue a frame to be pushed down the pipeline.
"""
"""Queue a frame to be pushed down the pipeline."""
pass
@abstractmethod
async def queue_frames(self, frames: Iterable[Frame] | AsyncIterable[Frame]):
"""
Queues multiple frames to be pushed down the pipeline.
"""
"""Queues multiple frames to be pushed down the pipeline."""
pass

26
src/pipecat/pipeline/runner.py
View File

@@ -5,22 +5,32 @@
#
import asyncio
import gc
import signal
from typing import Optional
from loguru import logger
from pipecat.pipeline.task import PipelineTask
from pipecat.utils.asyncio import current_tasks
from pipecat.utils.utils import obj_count, obj_id
class PipelineRunner:
def __init__(self, *, name: str | None = None, handle_sigint: bool = True):
def __init__(
self,
*,
name: str | None = None,
handle_sigint: bool = True,
force_gc: bool = False,
loop: Optional[asyncio.AbstractEventLoop] = None,
):
self.id: int = obj_id()
self.name: str = name or f"{self.__class__.__name__}#{obj_count(self)}"
self._tasks = {}
self._sig_task = None
self._force_gc = force_gc
self._loop = loop or asyncio.get_running_loop()
if handle_sigint:
self._setup_sigint()
@@ -28,13 +38,15 @@ class PipelineRunner:
async def run(self, task: PipelineTask):
logger.debug(f"Runner {self} started running {task}")
self._tasks[task.name] = task
task.set_event_loop(self._loop)
await task.run()
del self._tasks[task.name]
# If we are cancelling through a signal, make sure we wait for it so
# everything gets cleaned up nicely.
if self._sig_task:
await self._sig_task
self._print_dangling_tasks()
if self._force_gc:
self._gc_collect()
logger.debug(f"Runner {self} finished running {task}")
async def stop_when_done(self):
@@ -58,10 +70,10 @@ class PipelineRunner:
logger.warning(f"Interruption detected. Canceling runner {self}")
await self.cancel()
def _print_dangling_tasks(self):
tasks = [t.get_name() for t in current_tasks()]
if tasks:
logger.warning(f"Dangling tasks detected: {tasks}")
def _gc_collect(self):
collected = gc.collect()
logger.debug(f"Garbage collector: collected {collected} objects.")
logger.debug(f"Garbage collector: uncollectable objects {gc.garbage}")
def __str__(self):
return self.name

12
src/pipecat/pipeline/sync_parallel_pipeline.py
View File

@@ -78,16 +78,18 @@ class SyncParallelPipeline(BasePipeline):
down_queue = asyncio.Queue()
source = SyncParallelPipelineSource(up_queue)
sink = SyncParallelPipelineSink(down_queue)
processors: List[FrameProcessor] = [source] + processors + [sink]
# Create pipeline
pipeline = Pipeline(processors)
source.link(pipeline)
pipeline.link(sink)
self._pipelines.append(pipeline)
# Keep track of sources and sinks. We also keep the output queue of
# the source and the sinks so we can use it later.
self._sources.append({"processor": source, "queue": down_queue})
self._sinks.append({"processor": sink, "queue": up_queue})
# Create pipeline
pipeline = Pipeline(processors)
self._pipelines.append(pipeline)
logger.debug(f"Finished creating {self} pipelines")
#
@@ -103,7 +105,9 @@ class SyncParallelPipeline(BasePipeline):
async def cleanup(self):
await super().cleanup()
await asyncio.gather(*[s["processor"].cleanup() for s in self._sources])
await asyncio.gather(*[p.cleanup() for p in self._pipelines])
await asyncio.gather(*[s["processor"].cleanup() for s in self._sinks])
async def process_frame(self, frame: Frame, direction: FrameDirection):
await super().process_frame(frame, direction)

60
src/pipecat/pipeline/task.py
View File

@@ -30,7 +30,7 @@ from pipecat.pipeline.base_pipeline import BasePipeline
from pipecat.pipeline.base_task import BaseTask
from pipecat.pipeline.task_observer import TaskObserver
from pipecat.processors.frame_processor import FrameDirection, FrameProcessor
from pipecat.utils.asyncio import cancel_task, create_task, wait_for_task
from pipecat.utils.asyncio import TaskManager
from pipecat.utils.utils import obj_count, obj_id
HEARTBEAT_SECONDS = 1.0
@@ -122,7 +122,9 @@ class PipelineTask(BaseTask):
self._sink = PipelineTaskSink(self._down_queue)
pipeline.link(self._sink)
self._observer = TaskObserver(params.observers)
self._task_manager = TaskManager()
self._observer = TaskObserver(observers=params.observers, task_manager=self._task_manager)
@property
def id(self) -> int:
@@ -134,6 +136,9 @@ class PipelineTask(BaseTask):
"""Returns the name of this task."""
return self._name
def set_event_loop(self, loop: asyncio.AbstractEventLoop):
self._task_manager.set_event_loop(loop)
def has_finished(self) -> bool:
"""Indicates whether the tasks has finished. That is, all processors
have stopped.
@@ -159,15 +164,17 @@ class PipelineTask(BaseTask):
# we want to cancel right away.
await self._source.push_frame(CancelFrame())
# Only cancel the push task. Everything else will be cancelled in run().
await cancel_task(self._process_push_task)
await self._task_manager.cancel_task(self._process_push_task)
async def run(self):
"""
Starts running the given pipeline.
"""
if self.has_finished():
return
try:
push_task = self._create_tasks()
await wait_for_task(push_task)
push_task = await self._create_tasks()
await self._task_manager.wait_for_task(push_task)
except asyncio.CancelledError:
# We are awaiting on the push task and it might be cancelled
# (e.g. Ctrl-C). This means we will get a CancelledError here as
@@ -176,6 +183,7 @@ class PipelineTask(BaseTask):
pass
await self._cancel_tasks()
await self._cleanup()
self._print_dangling_tasks()
self._finished = True
async def queue_frame(self, frame: Frame):
@@ -195,42 +203,42 @@ class PipelineTask(BaseTask):
for frame in frames:
await self.queue_frame(frame)
def _create_tasks(self):
loop = asyncio.get_running_loop()
self._process_up_task = create_task(
loop, self._process_up_queue(), f"{self}::_process_up_queue"
async def _create_tasks(self):
self._process_up_task = self._task_manager.create_task(
self._process_up_queue(), f"{self}::_process_up_queue"
)
self._process_down_task = create_task(
loop, self._process_down_queue(), f"{self}::_process_down_queue"
self._process_down_task = self._task_manager.create_task(
self._process_down_queue(), f"{self}::_process_down_queue"
)
self._process_push_task = create_task(
loop, self._process_push_queue(), f"{self}::_process_push_queue"
self._process_push_task = self._task_manager.create_task(
self._process_push_queue(), f"{self}::_process_push_queue"
)
await self._observer.start()
return self._process_push_task
def _maybe_start_heartbeat_tasks(self):
if self._params.enable_heartbeats:
loop = asyncio.get_running_loop()
self._heartbeat_push_task = create_task(
loop, self._heartbeat_push_handler(), f"{self}::_heartbeat_push_handler"
self._heartbeat_push_task = self._task_manager.create_task(
self._heartbeat_push_handler(), f"{self}::_heartbeat_push_handler"
)
self._heartbeat_monitor_task = create_task(
loop, self._heartbeat_monitor_handler(), f"{self}::_heartbeat_monitor_handler"
self._heartbeat_monitor_task = self._task_manager.create_task(
self._heartbeat_monitor_handler(), f"{self}::_heartbeat_monitor_handler"
)
async def _cancel_tasks(self):
await self._maybe_cancel_heartbeat_tasks()
await cancel_task(self._process_up_task)
await cancel_task(self._process_down_task)
await self._task_manager.cancel_task(self._process_up_task)
await self._task_manager.cancel_task(self._process_down_task)
await self._observer.stop()
async def _maybe_cancel_heartbeat_tasks(self):
if self._params.enable_heartbeats:
await cancel_task(self._heartbeat_push_task)
await cancel_task(self._heartbeat_monitor_task)
await self._task_manager.cancel_task(self._heartbeat_push_task)
await self._task_manager.cancel_task(self._heartbeat_monitor_task)
def _initial_metrics_frame(self) -> MetricsFrame:
processors = self._pipeline.processors_with_metrics()
@@ -260,12 +268,13 @@ class PipelineTask(BaseTask):
self._maybe_start_heartbeat_tasks()
start_frame = StartFrame(
clock=self._clock,
task_manager=self._task_manager,
allow_interruptions=self._params.allow_interruptions,
enable_metrics=self._params.enable_metrics,
enable_usage_metrics=self._params.enable_usage_metrics,
report_only_initial_ttfb=self._params.report_only_initial_ttfb,
observer=self._observer,
clock=self._clock,
)
await self._source.queue_frame(start_frame, FrameDirection.DOWNSTREAM)
@@ -357,5 +366,10 @@ class PipelineTask(BaseTask):
f"{self}: heartbeat frame not received for more than {wait_time} seconds"
)
def _print_dangling_tasks(self):
tasks = [t.get_name() for t in self._task_manager.current_tasks()]
if tasks:
logger.warning(f"Dangling tasks detected: {tasks}")
def __str__(self):
return self.name

18
src/pipecat/pipeline/task_observer.py
View File

@@ -12,7 +12,7 @@ from attr import dataclass
from pipecat.frames.frames import Frame
from pipecat.observers.base_observer import BaseObserver
from pipecat.processors.frame_processor import FrameDirection, FrameProcessor
from pipecat.utils.asyncio import cancel_task, create_task
from pipecat.utils.asyncio import TaskManager
from pipecat.utils.utils import obj_count, obj_id
@@ -55,10 +55,12 @@ class TaskObserver(BaseObserver):
"""
def __init__(self, observers: List[BaseObserver] = []):
def __init__(self, *, observers: List[BaseObserver] = [], task_manager: TaskManager):
self._id: int = obj_id()
self._name: str = f"{self.__class__.__name__}#{obj_count(self)}"
self._proxies: List[Proxy] = self._create_proxies(observers)
self._observers = observers
self._task_manager = task_manager
self._proxies: List[Proxy] = []
@property
def id(self) -> int:
@@ -68,10 +70,14 @@ class TaskObserver(BaseObserver):
def name(self) -> str:
return self._name
async def start(self):
"""Starts all proxy observer tasks."""
self._proxies = self._create_proxies(self._observers)
async def stop(self):
"""Stops all proxy observer tasks."""
for proxy in self._proxies:
await cancel_task(proxy.task)
await self._task_manager.cancel_task(proxy.task)
async def on_push_frame(
self,
@@ -90,11 +96,9 @@ class TaskObserver(BaseObserver):
def _create_proxies(self, observers) -> List[Proxy]:
proxies = []
loop = asyncio.get_running_loop()
for observer in observers:
queue = asyncio.Queue()
task = create_task(
loop,
task = self._task_manager.create_task(
self._proxy_task_handler(queue, observer),
f"{self}::{observer.__class__.__name__}::_proxy_task_handler",
)

2
src/pipecat/processors/filters/stt_mute_filter.py
View File

@@ -121,6 +121,8 @@ class STTMuteFilter(FrameProcessor):
return False
async def process_frame(self, frame: Frame, direction: FrameDirection):
await super().process_frame(frame, direction)
"""Processes incoming frames and manages muting state."""
# Handle function call state changes
if isinstance(frame, FunctionCallInProgressFrame):

87
src/pipecat/processors/frame_processor.py
View File

@@ -6,7 +6,6 @@
import asyncio
import inspect
import sys
from enum import Enum
from typing import Awaitable, Callable, Coroutine, Optional
@@ -24,7 +23,7 @@ from pipecat.frames.frames import (
)
from pipecat.metrics.metrics import LLMTokenUsage, MetricsData
from pipecat.processors.metrics.frame_processor_metrics import FrameProcessorMetrics
from pipecat.utils.asyncio import cancel_task, create_task, wait_for_task
from pipecat.utils.asyncio import TaskManager
from pipecat.utils.utils import obj_count, obj_id
@@ -37,24 +36,25 @@ class FrameProcessor:
def __init__(
self,
*,
name: str | None = None,
metrics: FrameProcessorMetrics | None = None,
loop: asyncio.AbstractEventLoop | None = None,
name: Optional[str] = None,
metrics: Optional[FrameProcessorMetrics] = None,
**kwargs,
):
self._id: int = obj_id()
self._name = name or f"{self.__class__.__name__}#{obj_count(self)}"
self._parent: "FrameProcessor" | None = None
self._prev: "FrameProcessor" | None = None
self._next: "FrameProcessor" | None = None
self._loop: asyncio.AbstractEventLoop = loop or asyncio.get_running_loop()
self._parent: Optional["FrameProcessor"] = None
self._prev: Optional["FrameProcessor"] = None
self._next: Optional["FrameProcessor"] = None
self._event_handlers: dict = {}
# Clock
self._clock: BaseClock | None = None
self._clock: Optional[BaseClock] = None
# Properties
# Task Manager
self._task_manager: Optional[TaskManager] = None
# Other properties
self._allow_interruptions = False
self._enable_metrics = False
self._enable_usage_metrics = False
@@ -73,16 +73,16 @@ class FrameProcessor:
self._metrics.set_processor_name(self.name)
# Processors have an input queue. The input queue will be processed
# immediately (default) or it will block if `pause_processing_frames()`
# is called. To resume processing frames we need to call
# immediately (default) or it will block if `pause_processing_frames()` is
# called. To resume processing frames we need to call
# `resume_processing_frames()`.
self.__should_block_frames = False
self.__create_input_task()
self.__input_frame_task: Optional[asyncio.Task] = None
# Every processor in Pipecat should only output frames from a single
# task. This avoid problems like audio overlapping. System frames are
# the exception to this rule. This create this task.
self.__create_push_task()
# task. This avoid problems like audio overlapping. System frames are the
# exception to this rule. This create this task.
self.__push_frame_task: Optional[asyncio.Task] = None
@property
def id(self) -> int:
@@ -151,14 +151,20 @@ class FrameProcessor:
await self.stop_processing_metrics()
def create_task(self, coroutine: Coroutine) -> asyncio.Task:
if not self._task_manager:
raise Exception(f"{self} TaskManager is still not initialized.")
name = f"{self}::{coroutine.cr_code.co_name}"
return create_task(self.get_event_loop(), coroutine, name)
return self._task_manager.create_task(coroutine, name)
async def cancel_task(self, task: asyncio.Task, timeout: Optional[float] = None):
await cancel_task(task, timeout)
if not self._task_manager:
raise Exception(f"{self} TaskManager is still not initialized.")
await self._task_manager.cancel_task(task, timeout)
async def wait_for_task(self, task: asyncio.Task, timeout: Optional[float] = None):
await wait_for_task(task, timeout)
if not self._task_manager:
raise Exception(f"{self} TaskManager is still not initialized.")
await self._task_manager.wait_for_task(task, timeout)
async def cleanup(self):
await self.__cancel_input_task()
@@ -170,17 +176,26 @@ class FrameProcessor:
logger.debug(f"Linking {self} -> {self._next}")
def get_event_loop(self) -> asyncio.AbstractEventLoop:
return self._loop
if not self._task_manager:
raise Exception(f"{self} TaskManager is still not initialized.")
return self._task_manager.get_event_loop()
def set_parent(self, parent: "FrameProcessor"):
self._parent = parent
def get_parent(self) -> "FrameProcessor":
def get_parent(self) -> Optional["FrameProcessor"]:
return self._parent
def get_clock(self) -> BaseClock:
if not self._clock:
raise Exception(f"{self} Clock is still not initialized.")
return self._clock
def get_task_manager(self) -> TaskManager:
if not self._task_manager:
raise Exception(f"{self} TaskManager is still not initialized.")
return self._task_manager
async def queue_frame(
self,
frame: Frame,
@@ -211,11 +226,13 @@ class FrameProcessor:
async def process_frame(self, frame: Frame, direction: FrameDirection):
if isinstance(frame, StartFrame):
self._clock = frame.clock
self._task_manager = frame.task_manager
self._allow_interruptions = frame.allow_interruptions
self._enable_metrics = frame.enable_metrics
self._enable_usage_metrics = frame.enable_usage_metrics
self._report_only_initial_ttfb = frame.report_only_initial_ttfb
self._observer = frame.observer
await self.__start(frame)
elif isinstance(frame, StartInterruptionFrame):
await self._start_interruption()
await self.stop_all_metrics()
@@ -250,6 +267,10 @@ class FrameProcessor:
raise Exception(f"Event handler {event_name} already registered")
self._event_handlers[event_name] = []
async def __start(self, frame: StartFrame):
self.__create_input_task()
self.__create_push_task()
#
# Handle interruptions
#
@@ -299,13 +320,16 @@ class FrameProcessor:
raise
def __create_input_task(self):
self.__should_block_frames = False
self.__input_queue = asyncio.Queue()
self.__input_event = asyncio.Event()
self.__input_frame_task = self.create_task(self.__input_frame_task_handler())
if not self.__input_frame_task:
self.__should_block_frames = False
self.__input_queue = asyncio.Queue()
self.__input_event = asyncio.Event()
self.__input_frame_task = self.create_task(self.__input_frame_task_handler())
async def __cancel_input_task(self):
await self.cancel_task(self.__input_frame_task)
if self.__input_frame_task:
await self.cancel_task(self.__input_frame_task)
self.__input_frame_task = None
async def __input_frame_task_handler(self):
while True:
@@ -328,11 +352,14 @@ class FrameProcessor:
self.__input_queue.task_done()
def __create_push_task(self):
self.__push_queue = asyncio.Queue()
self.__push_frame_task = self.create_task(self.__push_frame_task_handler())
if not self.__push_frame_task:
self.__push_queue = asyncio.Queue()
self.__push_frame_task = self.create_task(self.__push_frame_task_handler())
async def __cancel_push_task(self):
await self.cancel_task(self.__push_frame_task)
if self.__push_frame_task:
await self.cancel_task(self.__push_frame_task)
self.__push_frame_task = None
async def __push_frame_task_handler(self):
while True:

6
src/pipecat/processors/frameworks/rtvi.py
View File

@@ -764,11 +764,11 @@ class RTVIProcessor(FrameProcessor):
# A task to process incoming action frames.
self._action_queue = asyncio.Queue()
self._action_task = self.create_task(self._action_task_handler())
self._action_task: Optional[asyncio.Task] = None
# A task to process incoming transport messages.
self._message_queue = asyncio.Queue()
self._message_task = self.create_task(self._message_task_handler())
self._message_task: Optional[asyncio.Task] = None
self._register_event_handler("on_bot_started")
self._register_event_handler("on_client_ready")
@@ -863,6 +863,8 @@ class RTVIProcessor(FrameProcessor):
await self._pipeline.cleanup()
async def _start(self, frame: StartFrame):
self._action_task = self.create_task(self._action_task_handler())
self._message_task = self.create_task(self._message_task_handler())
await self._call_event_handler("on_bot_started")
async def _stop(self, frame: EndFrame):

2
src/pipecat/processors/logger.py
View File

@@ -31,6 +31,8 @@ class FrameLogger(FrameProcessor):
self._ignored_frame_types = tuple(ignored_frame_types) if ignored_frame_types else None
async def process_frame(self, frame: Frame, direction: FrameDirection):
await super().process_frame(frame, direction)
if self._ignored_frame_types and not isinstance(frame, self._ignored_frame_types):
dir = "<" if direction is FrameDirection.UPSTREAM else ">"
msg = f"{dir} {self._prefix}: {frame}"

9
src/pipecat/services/ai_services.py
View File

@@ -399,7 +399,7 @@ class WordTTSService(TTSService):
super().__init__(**kwargs)
self._initial_word_timestamp = -1
self._words_queue = asyncio.Queue()
self._words_task = self.create_task(self._words_task_handler())
self._words_task = None
def start_word_timestamps(self):
if self._initial_word_timestamp == -1:
@@ -412,6 +412,10 @@ class WordTTSService(TTSService):
for word, timestamp in word_times:
await self._words_queue.put((word, seconds_to_nanoseconds(timestamp)))
async def start(self, frame: StartFrame):
await super().start(frame)
await self._create_words_task()
async def stop(self, frame: EndFrame):
await super().stop(frame)
await self._stop_words_task()
@@ -430,6 +434,9 @@ class WordTTSService(TTSService):
await super()._handle_interruption(frame, direction)
self.reset_word_timestamps()
async def _create_words_task(self):
self._words_task = self.create_task(self._words_task_handler())
async def _stop_words_task(self):
if self._words_task:
await self.cancel_task(self._words_task)

2
src/pipecat/services/gladia.py
View File

@@ -186,10 +186,12 @@ class GladiaSTTService(STTService):
await super().stop(frame)
await self._send_stop_recording()
await self._websocket.close()
await self.wait_for_task(self._receive_task)
async def cancel(self, frame: CancelFrame):
await super().cancel(frame)
await self._websocket.close()
await self.cancel_task(self._receive_task)
async def run_stt(self, audio: bytes) -> AsyncGenerator[Frame, None]:
await self.start_processing_metrics()

20
src/pipecat/tests/utils.py
View File

@@ -17,6 +17,7 @@ from pipecat.frames.frames import (
)
from pipecat.observers.base_observer import BaseObserver
from pipecat.processors.frame_processor import FrameDirection, FrameProcessor
from pipecat.utils.asyncio import TaskManager
@dataclass
@@ -54,9 +55,12 @@ class QueuedFrameProcessor(FrameProcessor):
async def process_frame(self, frame: Frame, direction: FrameDirection):
await super().process_frame(frame, direction)
if self._ignore_start and isinstance(frame, StartFrame):
return
await self._queue.put(frame)
await self.push_frame(frame, direction)
else:
await self._queue.put(frame)
await self.push_frame(frame, direction)
async def run_test(
@@ -67,13 +71,15 @@ async def run_test(
) -> Tuple[Sequence[Frame], Sequence[Frame]]:
received_up = asyncio.Queue()
received_down = asyncio.Queue()
up_processor = QueuedFrameProcessor(received_up)
down_processor = QueuedFrameProcessor(received_down)
source = QueuedFrameProcessor(received_up)
sink = QueuedFrameProcessor(received_down)
up_processor.link(processor)
processor.link(down_processor)
source.link(processor)
processor.link(sink)
await processor.queue_frame(StartFrame(clock=SystemClock()))
task_manager = TaskManager()
task_manager.set_event_loop(asyncio.get_event_loop())
await source.queue_frame(StartFrame(clock=SystemClock(), task_manager=task_manager))
for frame in frames_to_send:
await processor.process_frame(frame, FrameDirection.DOWNSTREAM)

5
src/pipecat/transports/base_output.py
View File

@@ -35,7 +35,6 @@ from pipecat.frames.frames import (
)
from pipecat.processors.frame_processor import FrameDirection, FrameProcessor
from pipecat.transports.base_transport import TransportParams
from pipecat.utils.asyncio import wait_for_task
from pipecat.utils.time import nanoseconds_to_seconds
@@ -88,9 +87,9 @@ class BaseOutputTransport(FrameProcessor):
# for these tasks before cancelling the camera and audio tasks below
# because they might be still rendering.
if self._sink_task:
await wait_for_task(self._sink_task)
await self.wait_for_task(self._sink_task)
if self._sink_clock_task:
await wait_for_task(self._sink_clock_task)
await self.wait_for_task(self._sink_clock_task)
# We can now cancel the camera task.
await self._cancel_camera_task()

2
src/pipecat/transports/base_transport.py
View File

@@ -51,13 +51,11 @@ class BaseTransport(ABC):
name: Optional[str] = None,
input_name: Optional[str] = None,
output_name: Optional[str] = None,
loop: Optional[asyncio.AbstractEventLoop] = None,
):
self._id: int = obj_id()
self._name = name or f"{self.__class__.__name__}#{obj_count(self)}"
self._input_name = input_name
self._output_name = output_name
self._loop = loop or asyncio.get_running_loop()
self._event_handlers: dict = {}
@property

36
src/pipecat/transports/network/fastapi_websocket.py
View File

@@ -29,7 +29,6 @@ from pipecat.serializers.base_serializer import FrameSerializer, FrameSerializer
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 cancel_task
try:
from fastapi import WebSocket
@@ -77,11 +76,11 @@ class FastAPIWebsocketInputTransport(BaseInputTransport):
async def stop(self, frame: EndFrame):
await super().stop(frame)
await cancel_task(self._receive_task)
await self.cancel_task(self._receive_task)
async def cancel(self, frame: CancelFrame):
await super().cancel(frame)
await cancel_task(self._receive_task)
await self.cancel_task(self._receive_task)
def _iter_data(self) -> typing.AsyncIterator[bytes | str]:
if self._params.serializer.type == FrameSerializerType.BINARY:
@@ -90,16 +89,19 @@ class FastAPIWebsocketInputTransport(BaseInputTransport):
return self._websocket.iter_text()
async def _receive_messages(self):
async for message in self._iter_data():
frame = self._params.serializer.deserialize(message)
try:
async for message in self._iter_data():
frame = self._params.serializer.deserialize(message)
if not frame:
continue
if not frame:
continue
if isinstance(frame, InputAudioRawFrame):
await self.push_audio_frame(frame)
else:
await self.push_frame(frame)
if isinstance(frame, InputAudioRawFrame):
await self.push_audio_frame(frame)
else:
await self.push_frame(frame)
except Exception as e:
logger.error(f"{self} exception receiving data (class: {e.__class__.__name__})")
await self._callbacks.on_client_disconnected(self._websocket)
@@ -160,9 +162,12 @@ class FastAPIWebsocketOutputTransport(BaseOutputTransport):
await self._write_audio_sleep()
async def _write_frame(self, frame: Frame):
payload = self._params.serializer.serialize(frame)
if payload and self._websocket.client_state == WebSocketState.CONNECTED:
await self._send_data(payload)
try:
payload = self._params.serializer.serialize(frame)
if payload and self._websocket.client_state == WebSocketState.CONNECTED:
await self._send_data(payload)
except Exception as e:
logger.error(f"{self} exception sending data (class: {e.__class__.__name__})")
def _send_data(self, data: str | bytes):
if self._params.serializer.type == FrameSerializerType.BINARY:
@@ -188,9 +193,8 @@ class FastAPIWebsocketTransport(BaseTransport):
params: FastAPIWebsocketParams,
input_name: str | None = None,
output_name: str | None = None,
loop: asyncio.AbstractEventLoop | None = None,
):
super().__init__(input_name=input_name, output_name=output_name, loop=loop)
super().__init__(input_name=input_name, output_name=output_name)
self._params = params
self._callbacks = FastAPIWebsocketCallbacks(

51
src/pipecat/transports/network/websocket_client.py
View File

@@ -9,7 +9,7 @@ import asyncio
import io
import time
import wave
from typing import Awaitable, Callable
from typing import Awaitable, Callable, Optional
import websockets
from loguru import logger
@@ -30,7 +30,7 @@ from pipecat.serializers.protobuf import ProtobufFrameSerializer
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 cancel_task, create_task
from pipecat.utils.asyncio import TaskManager
class WebsocketClientParams(TransportParams):
@@ -50,25 +50,36 @@ class WebsocketClientSession:
uri: str,
params: WebsocketClientParams,
callbacks: WebsocketClientCallbacks,
loop: asyncio.AbstractEventLoop,
transport_name: str,
):
self._uri = uri
self._params = params
self._callbacks = callbacks
self._loop = loop
self._transport_name = transport_name
self._task_manager: Optional[TaskManager] = None
self._websocket: websockets.WebSocketClientProtocol | None = None
@property
def task_manager(self) -> TaskManager:
if not self._task_manager:
raise Exception(
f"{self._transport_name}::WebsocketClientSession: TaskManager not initialized (pipeline not started?)"
)
return self._task_manager
async def setup(self, frame: StartFrame):
if not self._task_manager:
self._task_manager = frame.task_manager
async def connect(self):
if self._websocket:
return
try:
self._websocket = await websockets.connect(uri=self._uri, open_timeout=10)
self._client_task = create_task(
self._loop,
self._client_task = self.task_manager.create_task(
self._client_task_handler(),
f"{self._transport_name}::WebsocketClientSession::_client_task_handler",
)
@@ -80,22 +91,31 @@ class WebsocketClientSession:
if not self._websocket:
return
await cancel_task(self._client_task)
await self.task_manager.cancel_task(self._client_task)
await self._websocket.close()
self._websocket = None
async def send(self, message: websockets.Data):
if self._websocket:
await self._websocket.send(message)
try:
if self._websocket:
await self._websocket.send(message)
except Exception as e:
logger.error(f"{self} exception sending data (class: {e.__class__.__name__})")
async def _client_task_handler(self):
# Handle incoming messages
async for message in self._websocket:
await self._callbacks.on_message(self._websocket, message)
try:
# Handle incoming messages
async for message in self._websocket:
await self._callbacks.on_message(self._websocket, message)
except Exception as e:
logger.error(f"{self} exception receiving data (class: {e.__class__.__name__})")
await self._callbacks.on_disconnected(self._websocket)
def __str__(self):
return f"{self._transport_name}::WebsocketClientSession"
class WebsocketClientInputTransport(BaseInputTransport):
def __init__(self, session: WebsocketClientSession, params: WebsocketClientParams):
@@ -106,6 +126,7 @@ class WebsocketClientInputTransport(BaseInputTransport):
async def start(self, frame: StartFrame):
await super().start(frame)
await self._session.setup(frame)
await self._session.connect()
async def stop(self, frame: EndFrame):
@@ -138,6 +159,7 @@ class WebsocketClientOutputTransport(BaseOutputTransport):
async def start(self, frame: StartFrame):
await super().start(frame)
await self._session.setup(frame)
await self._session.connect()
async def stop(self, frame: EndFrame):
@@ -198,9 +220,8 @@ class WebsocketClientTransport(BaseTransport):
self,
uri: str,
params: WebsocketClientParams = WebsocketClientParams(),
loop: asyncio.AbstractEventLoop | None = None,
):
super().__init__(loop=loop)
super().__init__()
self._params = params
@@ -210,7 +231,7 @@ class WebsocketClientTransport(BaseTransport):
on_message=self._on_message,
)
self._session = WebsocketClientSession(uri, params, callbacks, self._loop, self.name)
self._session = WebsocketClientSession(uri, params, callbacks, self.name)
self._input: WebsocketClientInputTransport | None = None
self._output: WebsocketClientOutputTransport | None = None

30
src/pipecat/transports/network/websocket_server.py
View File

@@ -100,19 +100,22 @@ class WebsocketServerInputTransport(BaseInputTransport):
# Create a task to monitor the websocket connection
if self._params.session_timeout:
self.get_event_loop().create_task(self._monitor_websocket(websocket))
self.create_task(self._monitor_websocket(websocket))
# Handle incoming messages
async for message in websocket:
frame = self._params.serializer.deserialize(message)
try:
async for message in websocket:
frame = self._params.serializer.deserialize(message)
if not frame:
continue
if not frame:
continue
if isinstance(frame, InputAudioRawFrame):
await self.push_audio_frame(frame)
else:
await self.push_frame(frame)
if isinstance(frame, InputAudioRawFrame):
await self.push_audio_frame(frame)
else:
await self.push_frame(frame)
except Exception as e:
logger.error(f"{self} exception receiving data (class: {e.__class__.__name__})")
# Notify disconnection
await self._callbacks.on_client_disconnected(websocket)
@@ -189,9 +192,12 @@ class WebsocketServerOutputTransport(BaseOutputTransport):
await self._write_audio_sleep()
async def _write_frame(self, frame: Frame):
payload = self._params.serializer.serialize(frame)
if payload and self._websocket:
await self._websocket.send(payload)
try:
payload = self._params.serializer.serialize(frame)
if payload and self._websocket:
await self._websocket.send(payload)
except Exception as e:
logger.error(f"{self} exception sending data (class: {e.__class__.__name__})")
async def _write_audio_sleep(self):
# Simulate a clock.

86
src/pipecat/transports/services/daily.py
View File

@@ -46,7 +46,7 @@ from pipecat.transcriptions.language import Language
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 cancel_task, create_task
from pipecat.utils.asyncio import TaskManager
try:
from daily import CallClient, Daily, EventHandler
@@ -179,7 +179,6 @@ class DailyTransportClient(EventHandler):
bot_name: str,
params: DailyParams,
callbacks: DailyCallbacks,
loop: asyncio.AbstractEventLoop,
transport_name: str,
):
super().__init__()
@@ -193,7 +192,6 @@ class DailyTransportClient(EventHandler):
self._bot_name: str = bot_name
self._params: DailyParams = params
self._callbacks = callbacks
self._loop = loop
self._transport_name = transport_name
self._participant_id: str = ""
@@ -205,6 +203,8 @@ class DailyTransportClient(EventHandler):
self._joined_event = asyncio.Event()
self._leave_counter = 0
self._task_manager: Optional[TaskManager] = None
# We use the executor to cleanup the client. We just do it from one
# place, so only one thread is really needed.
self._executor = ThreadPoolExecutor(max_workers=1)
@@ -220,11 +220,7 @@ class DailyTransportClient(EventHandler):
# future) we will deadlock because completions use event handlers (which
# are holding the GIL).
self._callback_queue = asyncio.Queue()
self._callback_task = create_task(
self._loop,
self._callback_task_handler(),
f"{self._transport_name}::DailyTransportClient::_callback_task_handler",
)
self._callback_task = None
self._camera: VirtualCameraDevice | None = None
if self._params.camera_out_enabled:
@@ -267,9 +263,6 @@ class DailyTransportClient(EventHandler):
def participant_id(self) -> str:
return self._participant_id
def set_callbacks(self, callbacks: DailyCallbacks):
self._callbacks = callbacks
async def send_message(self, frame: TransportMessageFrame | TransportMessageUrgentFrame):
if not self._joined:
return
@@ -278,7 +271,7 @@ class DailyTransportClient(EventHandler):
if isinstance(frame, (DailyTransportMessageFrame, DailyTransportMessageUrgentFrame)):
participant_id = frame.participant_id
future = self._loop.create_future()
future = self._get_event_loop().create_future()
self._client.send_app_message(
frame.message, participant_id, completion=completion_callback(future)
)
@@ -292,7 +285,7 @@ class DailyTransportClient(EventHandler):
num_channels = self._params.audio_in_channels
num_frames = int(sample_rate / 100) * 2 # 20ms of audio
future = self._loop.create_future()
future = self._get_event_loop().create_future()
self._speaker.read_frames(num_frames, completion=completion_callback(future))
audio = await future
@@ -311,7 +304,7 @@ class DailyTransportClient(EventHandler):
if not self._mic:
return None
future = self._loop.create_future()
future = self._get_event_loop().create_future()
self._mic.write_frames(frames, completion=completion_callback(future))
await future
@@ -321,6 +314,14 @@ class DailyTransportClient(EventHandler):
self._camera.write_frame(frame.image)
async def setup(self, frame: StartFrame):
if not self._task_manager:
self._task_manager = frame.task_manager
self._callback_task = self._task_manager.create_task(
self._callback_task_handler(),
f"{self}::callback_task",
)
async def join(self):
# Transport already joined, ignore.
if self._joined:
@@ -370,7 +371,7 @@ class DailyTransportClient(EventHandler):
logger.info(f"Enabling transcription with settings {self._params.transcription_settings}")
future = self._loop.create_future()
future = self._get_event_loop().create_future()
self._client.start_transcription(
settings=self._params.transcription_settings.model_dump(exclude_none=True),
completion=completion_callback(future),
@@ -381,7 +382,7 @@ class DailyTransportClient(EventHandler):
return
async def _join(self):
future = self._loop.create_future()
future = self._get_event_loop().create_future()
self._client.join(
self._room_url,
@@ -466,24 +467,24 @@ class DailyTransportClient(EventHandler):
async def _stop_transcription(self):
if not self._token:
return
future = self._loop.create_future()
future = self._get_event_loop().create_future()
self._client.stop_transcription(completion=completion_callback(future))
error = await future
if error:
logger.error(f"Unable to stop transcription: {error}")
async def _leave(self):
future = self._loop.create_future()
future = self._get_event_loop().create_future()
self._client.leave(completion=completion_callback(future))
return await asyncio.wait_for(future, timeout=10)
async def cleanup(self):
if self._callback_task:
await cancel_task(self._callback_task)
if self._callback_task and self._task_manager:
await self._task_manager.cancel_task(self._callback_task)
self._callback_task = None
# Make sure we don't block the event loop in case `client.release()`
# takes extra time.
await self._loop.run_in_executor(self._executor, self._cleanup)
await self._get_event_loop().run_in_executor(self._executor, self._cleanup)
def _cleanup(self):
if self._client:
@@ -497,39 +498,39 @@ class DailyTransportClient(EventHandler):
return self._client.participant_counts()
async def start_dialout(self, settings):
future = self._loop.create_future()
future = self._get_event_loop().create_future()
self._client.start_dialout(settings, completion=completion_callback(future))
await future
async def stop_dialout(self, participant_id):
future = self._loop.create_future()
future = self._get_event_loop().create_future()
self._client.stop_dialout(participant_id, completion=completion_callback(future))
await future
async def send_dtmf(self, settings):
future = self._loop.create_future()
future = self._get_event_loop().create_future()
self._client.send_dtmf(settings, completion=completion_callback(future))
await future
async def sip_call_transfer(self, settings):
future = self._loop.create_future()
future = self._get_event_loop().create_future()
self._client.sip_call_transfer(settings, completion=completion_callback(future))
await future
async def sip_refer(self, settings):
future = self._loop.create_future()
future = self._get_event_loop().create_future()
self._client.sip_refer(settings, completion=completion_callback(future))
await future
async def start_recording(self, streaming_settings, stream_id, force_new):
future = self._loop.create_future()
future = self._get_event_loop().create_future()
self._client.start_recording(
streaming_settings, stream_id, force_new, completion=completion_callback(future)
)
await future
async def stop_recording(self, stream_id):
future = self._loop.create_future()
future = self._get_event_loop().create_future()
self._client.stop_recording(stream_id, completion=completion_callback(future))
await future
@@ -537,7 +538,7 @@ class DailyTransportClient(EventHandler):
if not self._joined:
return
future = self._loop.create_future()
future = self._get_event_loop().create_future()
self._client.send_prebuilt_chat_message(
message, user_name=user_name, completion=completion_callback(future)
)
@@ -574,14 +575,14 @@ class DailyTransportClient(EventHandler):
)
async def update_transcription(self, participants=None, instance_id=None):
future = self._loop.create_future()
future = self._get_event_loop().create_future()
self._client.update_transcription(
participants, instance_id, completion=completion_callback(future)
)
await future
async def update_subscriptions(self, participant_settings=None, profile_settings=None):
future = self._loop.create_future()
future = self._get_event_loop().create_future()
self._client.update_subscriptions(
participant_settings=participant_settings,
profile_settings=profile_settings,
@@ -681,7 +682,7 @@ class DailyTransportClient(EventHandler):
def _call_async_callback(self, callback, *args):
future = asyncio.run_coroutine_threadsafe(
self._callback_queue.put((callback, *args)), self._loop
self._callback_queue.put((callback, *args)), self._get_event_loop()
)
future.result()
@@ -692,6 +693,14 @@ class DailyTransportClient(EventHandler):
(callback, *args) = await self._callback_queue.get()
await callback(*args)
def _get_event_loop(self) -> asyncio.AbstractEventLoop:
if not self._task_manager:
raise Exception(f"{self}: missing task manager (pipeline not started?)")
return self._task_manager.get_event_loop()
def __str__(self):
return f"{self._transport_name}::DailyTransportClient"
class DailyInputTransport(BaseInputTransport):
def __init__(self, client: DailyTransportClient, params: DailyParams, **kwargs):
@@ -715,6 +724,8 @@ class DailyInputTransport(BaseInputTransport):
async def start(self, frame: StartFrame):
# Parent start.
await super().start(frame)
# Setup client.
await self._client.setup(frame)
# Join the room.
await self._client.join()
# Create audio task. It reads audio frames from Daily and push them
@@ -837,6 +848,8 @@ class DailyOutputTransport(BaseOutputTransport):
async def start(self, frame: StartFrame):
# Parent start.
await super().start(frame)
# Setup client.
await self._client.setup(frame)
# Join the room.
await self._client.join()
@@ -875,9 +888,8 @@ class DailyTransport(BaseTransport):
params: DailyParams = DailyParams(),
input_name: str | None = None,
output_name: str | None = None,
loop: asyncio.AbstractEventLoop | None = None,
):
super().__init__(input_name=input_name, output_name=output_name, loop=loop)
super().__init__(input_name=input_name, output_name=output_name)
callbacks = DailyCallbacks(
on_joined=self._on_joined,
@@ -905,9 +917,7 @@ class DailyTransport(BaseTransport):
)
self._params = params
self._client = DailyTransportClient(
room_url, token, bot_name, params, callbacks, self._loop, self.name
)
self._client = DailyTransportClient(room_url, token, bot_name, params, callbacks, self.name)
self._input: DailyInputTransport | None = None
self._output: DailyOutputTransport | None = None
@@ -1042,7 +1052,7 @@ class DailyTransport(BaseTransport):
await self._output.push_error(error_frame)
else:
logger.error("Both input and output are None while trying to push error")
raise RuntimeError("No valid input or output channel to push error")
raise Exception("No valid input or output channel to push error")
async def _on_app_message(self, message: Any, sender: str):
if self._input:

124
src/pipecat/transports/services/livekit.py
View File

@@ -6,7 +6,7 @@
import asyncio
from dataclasses import dataclass
from typing import Any, Awaitable, Callable, List
from typing import Any, Awaitable, Callable, List, Optional
from loguru import logger
from pydantic import BaseModel
@@ -17,7 +17,6 @@ from pipecat.frames.frames import (
AudioRawFrame,
CancelFrame,
EndFrame,
Frame,
InputAudioRawFrame,
OutputAudioRawFrame,
StartFrame,
@@ -28,7 +27,7 @@ from pipecat.processors.frame_processor import FrameDirection
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 create_task
from pipecat.utils.asyncio import TaskManager
try:
from livekit import rtc
@@ -72,7 +71,6 @@ class LiveKitTransportClient:
room_name: str,
params: LiveKitParams,
callbacks: LiveKitCallbacks,
loop: asyncio.AbstractEventLoop,
transport_name: str,
):
self._url = url
@@ -80,9 +78,8 @@ class LiveKitTransportClient:
self._room_name = room_name
self._params = params
self._callbacks = callbacks
self._loop = loop
self._transport_name = transport_name
self._room = rtc.Room(loop=loop)
self._room: rtc.Room | None = None
self._participant_id: str = ""
self._connected = False
self._disconnect_counter = 0
@@ -91,20 +88,32 @@ class LiveKitTransportClient:
self._audio_tracks = {}
self._audio_queue = asyncio.Queue()
self._other_participant_has_joined = False
# Set up room event handlers
self._room.on("participant_connected")(self._on_participant_connected_wrapper)
self._room.on("participant_disconnected")(self._on_participant_disconnected_wrapper)
self._room.on("track_subscribed")(self._on_track_subscribed_wrapper)
self._room.on("track_unsubscribed")(self._on_track_unsubscribed_wrapper)
self._room.on("data_received")(self._on_data_received_wrapper)
self._room.on("connected")(self._on_connected_wrapper)
self._room.on("disconnected")(self._on_disconnected_wrapper)
self._task_manager: Optional[TaskManager] = None
@property
def participant_id(self) -> str:
return self._participant_id
@property
def room(self) -> rtc.Room:
if not self._room:
raise Exception(f"{self}: missing room object (pipeline not started?)")
return self._room
async def setup(self, frame: StartFrame):
if not self._task_manager:
self._task_manager = frame.task_manager
self._room = rtc.Room(loop=self._task_manager.get_event_loop())
# Set up room event handlers
self.room.on("participant_connected")(self._on_participant_connected_wrapper)
self.room.on("participant_disconnected")(self._on_participant_disconnected_wrapper)
self.room.on("track_subscribed")(self._on_track_subscribed_wrapper)
self.room.on("track_unsubscribed")(self._on_track_unsubscribed_wrapper)
self.room.on("data_received")(self._on_data_received_wrapper)
self.room.on("connected")(self._on_connected_wrapper)
self.room.on("disconnected")(self._on_disconnected_wrapper)
@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10))
async def connect(self):
if self._connected:
@@ -115,7 +124,7 @@ class LiveKitTransportClient:
logger.info(f"Connecting to {self._room_name}")
try:
await self._room.connect(
await self.room.connect(
self._url,
self._token,
options=rtc.RoomOptions(auto_subscribe=True),
@@ -124,7 +133,7 @@ class LiveKitTransportClient:
# Increment disconnect counter if we successfully connected.
self._disconnect_counter += 1
self._participant_id = self._room.local_participant.sid
self._participant_id = self.room.local_participant.sid
logger.info(f"Connected to {self._room_name}")
# Set up audio source and track
@@ -136,7 +145,7 @@ class LiveKitTransportClient:
)
options = rtc.TrackPublishOptions()
options.source = rtc.TrackSource.SOURCE_MICROPHONE
await self._room.local_participant.publish_track(self._audio_track, options)
await self.room.local_participant.publish_track(self._audio_track, options)
await self._callbacks.on_connected()
@@ -157,7 +166,7 @@ class LiveKitTransportClient:
return
logger.info(f"Disconnecting from {self._room_name}")
await self._room.disconnect()
await self.room.disconnect()
self._connected = False
logger.info(f"Disconnected from {self._room_name}")
await self._callbacks.on_disconnected()
@@ -168,11 +177,11 @@ class LiveKitTransportClient:
try:
if participant_id:
await self._room.local_participant.publish_data(
await self.room.local_participant.publish_data(
data, reliable=True, destination_identities=[participant_id]
)
else:
await self._room.local_participant.publish_data(data, reliable=True)
await self.room.local_participant.publish_data(data, reliable=True)
except Exception as e:
logger.error(f"Error sending data: {e}")
@@ -186,10 +195,10 @@ class LiveKitTransportClient:
logger.error(f"Error publishing audio: {e}")
def get_participants(self) -> List[str]:
return [p.sid for p in self._room.remote_participants.values()]
return [p.sid for p in self.room.remote_participants.values()]
async def get_participant_metadata(self, participant_id: str) -> dict:
participant = self._room.remote_participants.get(participant_id)
participant = self.room.remote_participants.get(participant_id)
if participant:
return {
"id": participant.sid,
@@ -200,17 +209,17 @@ class LiveKitTransportClient:
return {}
async def set_participant_metadata(self, metadata: str):
await self._room.local_participant.set_metadata(metadata)
await self.room.local_participant.set_metadata(metadata)
async def mute_participant(self, participant_id: str):
participant = self._room.remote_participants.get(participant_id)
participant = self.room.remote_participants.get(participant_id)
if participant:
for track in participant.tracks.values():
if track.kind == "audio":
await track.set_enabled(False)
async def unmute_participant(self, participant_id: str):
participant = self._room.remote_participants.get(participant_id)
participant = self.room.remote_participants.get(participant_id)
if participant:
for track in participant.tracks.values():
if track.kind == "audio":
@@ -218,17 +227,15 @@ class LiveKitTransportClient:
# Wrapper methods for event handlers
def _on_participant_connected_wrapper(self, participant: rtc.RemoteParticipant):
create_task(
self._loop,
self._task_manager.create_task(
self._async_on_participant_connected(participant),
f"{self._transport_name}::LiveKitTransportClient::_async_on_participant_connected",
f"{self}::_async_on_participant_connected",
)
def _on_participant_disconnected_wrapper(self, participant: rtc.RemoteParticipant):
create_task(
self._loop,
self._task_manager.create_task(
self._async_on_participant_disconnected(participant),
f"{self._transport_name}::LiveKitTransportClient::_async_on_participant_disconnected",
f"{self}::_async_on_participant_disconnected",
)
def _on_track_subscribed_wrapper(
@@ -237,10 +244,9 @@ class LiveKitTransportClient:
publication: rtc.RemoteTrackPublication,
participant: rtc.RemoteParticipant,
):
create_task(
self._loop,
self._task_manager.create_task(
self._async_on_track_subscribed(track, publication, participant),
f"{self._transport_name}::LiveKitTransportClient::_async_on_track_subscribed",
f"{self}::_async_on_track_subscribed",
)
def _on_track_unsubscribed_wrapper(
@@ -249,31 +255,23 @@ class LiveKitTransportClient:
publication: rtc.RemoteTrackPublication,
participant: rtc.RemoteParticipant,
):
create_task(
self._loop,
self._task_manager.create_task(
self._async_on_track_unsubscribed(track, publication, participant),
f"{self._transport_name}::LiveKitTransportClient::_async_on_track_unsubscribed",
f"{self}::_async_on_track_unsubscribed",
)
def _on_data_received_wrapper(self, data: rtc.DataPacket):
create_task(
self._loop,
self._task_manager.create_task(
self._async_on_data_received(data),
f"{self._transport_name}::LiveKitTransportClient::_async_on_data_received",
f"{self}::_async_on_data_received",
)
def _on_connected_wrapper(self):
create_task(
self._loop,
self._async_on_connected(),
f"{self._transport_name}::LiveKitTransportClient::_async_on_connected",
)
self._task_manager.create_task(self._async_on_connected(), f"{self}::_async_on_connected")
def _on_disconnected_wrapper(self):
create_task(
self._loop,
self._async_on_disconnected(),
f"{self._transport_name}::LiveKitTransportClient::_async_on_disconnected",
self._task_manager.create_task(
self._async_on_disconnected(), f"{self}::_async_on_disconnected"
)
# Async methods for event handling
@@ -300,10 +298,9 @@ class LiveKitTransportClient:
logger.info(f"Audio track subscribed: {track.sid} from participant {participant.sid}")
self._audio_tracks[participant.sid] = track
audio_stream = rtc.AudioStream(track)
create_task(
self._loop,
self._task_manager.create_task(
self._process_audio_stream(audio_stream, participant.sid),
f"{self._transport_name}::LiveKitTransportClient::_process_audio_stream",
f"{self}::_process_audio_stream",
)
async def _async_on_track_unsubscribed(
@@ -342,6 +339,9 @@ class LiveKitTransportClient:
frame, participant_id = await self._audio_queue.get()
return frame, participant_id
def __str__(self):
return f"{self._transport_name}::LiveKitTransportClient"
class LiveKitInputTransport(BaseInputTransport):
def __init__(self, client: LiveKitTransportClient, params: LiveKitParams, **kwargs):
@@ -352,6 +352,7 @@ class LiveKitInputTransport(BaseInputTransport):
async def start(self, frame: StartFrame):
await super().start(frame)
await self._client.setup(frame)
await self._client.connect()
if self._params.audio_in_enabled or self._params.vad_enabled:
self._audio_in_task = self.create_task(self._audio_in_task_handler())
@@ -395,13 +396,10 @@ class LiveKitInputTransport(BaseInputTransport):
self, audio_frame_event: rtc.AudioFrameEvent
) -> AudioRawFrame:
audio_frame = audio_frame_event.frame
audio_data = audio_frame.data
original_sample_rate = audio_frame.sample_rate
if original_sample_rate != self._params.audio_in_sample_rate:
audio_data = resample_audio(
audio_data, original_sample_rate, self._params.audio_in_sample_rate
)
audio_data = resample_audio(
audio_frame.data.tobytes(), audio_frame.sample_rate, self._params.audio_in_sample_rate
)
return AudioRawFrame(
audio=audio_data,
@@ -417,6 +415,7 @@ class LiveKitOutputTransport(BaseOutputTransport):
async def start(self, frame: StartFrame):
await super().start(frame)
await self._client.setup(frame)
await self._client.connect()
logger.info("LiveKitOutputTransport started")
@@ -461,9 +460,8 @@ class LiveKitTransport(BaseTransport):
params: LiveKitParams = LiveKitParams(),
input_name: str | None = None,
output_name: str | None = None,
loop: asyncio.AbstractEventLoop | None = None,
):
super().__init__(input_name=input_name, output_name=output_name, loop=loop)
super().__init__(input_name=input_name, output_name=output_name)
callbacks = LiveKitCallbacks(
on_connected=self._on_connected,
@@ -478,7 +476,7 @@ class LiveKitTransport(BaseTransport):
self._params = params
self._client = LiveKitTransportClient(
url, token, room_name, self._params, callbacks, self._loop, self.name
url, token, room_name, self._params, callbacks, self.name
)
self._input: LiveKitInputTransport | None = None
self._output: LiveKitOutputTransport | None = None
@@ -544,7 +542,7 @@ class LiveKitTransport(BaseTransport):
async def _on_audio_track_subscribed(self, participant_id: str):
await self._call_event_handler("on_audio_track_subscribed", participant_id)
participant = self._client._room.remote_participants.get(participant_id)
participant = self._client.room.remote_participants.get(participant_id)
if participant:
for publication in participant.audio_tracks.values():
self._client._on_track_subscribed_wrapper(

189
src/pipecat/utils/asyncio.py
View File

@@ -9,106 +9,121 @@ from typing import Coroutine, Optional, Set
from loguru import logger
_TASKS: Set[asyncio.Task] = set()
class TaskManager:
def __init__(self) -> None:
self._tasks: Set[asyncio.Task] = set()
self._loop: Optional[asyncio.AbstractEventLoop] = None
def create_task(loop: asyncio.AbstractEventLoop, coroutine: Coroutine, name: str) -> asyncio.Task:
"""
Creates and schedules a new asyncio Task that runs the given coroutine.
def set_event_loop(self, loop: asyncio.AbstractEventLoop):
self._loop = loop
The task is added to a global set of created tasks.
def get_event_loop(self) -> asyncio.AbstractEventLoop:
if not self._loop:
raise Exception("TaskManager missing event loop, use TaskManager.set_event_loop().")
return self._loop
Args:
loop (asyncio.AbstractEventLoop): The event loop to use for creating the task.
coroutine (Coroutine): The coroutine to be executed within the task.
name (str): The name to assign to the task for identification.
def create_task(self, coroutine: Coroutine, name: str) -> asyncio.Task:
"""
Creates and schedules a new asyncio Task that runs the given coroutine.
Returns:
asyncio.Task: The created task object.
"""
The task is added to a global set of created tasks.
async def run_coroutine():
Args:
loop (asyncio.AbstractEventLoop): The event loop to use for creating the task.
coroutine (Coroutine): The coroutine to be executed within the task.
name (str): The name to assign to the task for identification.
Returns:
asyncio.Task: The created task object.
"""
async def run_coroutine():
try:
await coroutine
except asyncio.CancelledError:
logger.trace(f"{name}: task cancelled")
# Re-raise the exception to ensure the task is cancelled.
raise
except Exception as e:
logger.exception(f"{name}: unexpected exception: {e}")
if not self._loop:
raise Exception("TaskManager missing event loop, use TaskManager.set_event_loop().")
task = self._loop.create_task(run_coroutine())
task.set_name(name)
self._add_task(task)
logger.trace(f"{name}: task created")
return task
async def wait_for_task(self, task: asyncio.Task, timeout: Optional[float] = None):
"""Wait for an asyncio.Task to complete with optional timeout handling.
This function awaits the specified asyncio.Task and handles scenarios for
timeouts, cancellations, and other exceptions. It also ensures that the task
is removed from the set of registered tasks upon completion or failure.
Args:
task (asyncio.Task): The asyncio Task to wait for.
timeout (Optional[float], optional): The maximum number of seconds
to wait for the task to complete. If None, waits indefinitely.
Defaults to None.
"""
name = task.get_name()
try:
await coroutine
if timeout:
await asyncio.wait_for(task, timeout=timeout)
else:
await task
except asyncio.TimeoutError:
logger.warning(f"{name}: timed out waiting for task to finish")
except asyncio.CancelledError:
logger.trace(f"{name}: task cancelled")
# Re-raise the exception to ensure the task is cancelled.
raise
logger.trace(f"{name}: unexpected task cancellation (maybe Ctrl-C?)")
except Exception as e:
logger.exception(f"{name}: unexpected exception: {e}")
logger.exception(f"{name}: unexpected exception while stopping task: {e}")
finally:
self._remove_task(task)
task = loop.create_task(run_coroutine())
task.set_name(name)
_TASKS.add(task)
logger.trace(f"{name}: task created")
return task
async def cancel_task(self, task: asyncio.Task, timeout: Optional[float] = None):
"""Cancels the given asyncio Task and awaits its completion with an
optional timeout.
This function removes the task from the set of registered tasks upon
completion or failure.
async def wait_for_task(task: asyncio.Task, timeout: Optional[float] = None):
"""Wait for an asyncio.Task to complete with optional timeout handling.
Args:
task (asyncio.Task): The task to be cancelled.
timeout (Optional[float]): The optional timeout in seconds to wait for the task to cancel.
This function awaits the specified asyncio.Task and handles scenarios for
timeouts, cancellations, and other exceptions. It also ensures that the task
is removed from the set of registered tasks upon completion or failure.
Args:
task (asyncio.Task): The asyncio Task to wait for.
timeout (Optional[float], optional): The maximum number of seconds
to wait for the task to complete. If None, waits indefinitely.
Defaults to None.
"""
name = task.get_name()
try:
if timeout:
await asyncio.wait_for(task, timeout=timeout)
else:
await task
except asyncio.TimeoutError:
logger.warning(f"{name}: timed out waiting for task to finish")
except asyncio.CancelledError:
logger.trace(f"{name}: unexpected task cancellation (maybe Ctrl-C?)")
except Exception as e:
logger.exception(f"{name}: unexpected exception while stopping task: {e}")
finally:
"""
name = task.get_name()
task.cancel()
try:
_TASKS.remove(task)
if timeout:
await asyncio.wait_for(task, timeout=timeout)
else:
await task
except asyncio.TimeoutError:
logger.warning(f"{name}: timed out waiting for task to cancel")
except asyncio.CancelledError:
# Here are sure the task is cancelled properly.
pass
except Exception as e:
logger.exception(f"{name}: unexpected exception while cancelling task: {e}")
finally:
self._remove_task(task)
def current_tasks(self) -> Set[asyncio.Task]:
"""Returns the list of currently created/registered tasks."""
return self._tasks
def _add_task(self, task: asyncio.Task):
self._tasks.add(task)
def _remove_task(self, task: asyncio.Task):
name = task.get_name()
try:
self._tasks.remove(task)
except KeyError as e:
logger.trace(f"{name}: unable to remove task (already removed?): {e}")
async def cancel_task(task: asyncio.Task, timeout: Optional[float] = None):
"""Cancels the given asyncio Task and awaits its completion with an
optional timeout.
This function removes the task from the set of registered tasks upon
completion or failure.
Args:
task (asyncio.Task): The task to be cancelled.
timeout (Optional[float]): The optional timeout in seconds to wait for the task to cancel.
"""
name = task.get_name()
task.cancel()
try:
if timeout:
await asyncio.wait_for(task, timeout=timeout)
else:
await task
except asyncio.TimeoutError:
logger.warning(f"{name}: timed out waiting for task to cancel")
except asyncio.CancelledError:
# Here are sure the task is cancelled properly.
pass
except Exception as e:
logger.exception(f"{name}: unexpected exception while cancelling task: {e}")
finally:
try:
_TASKS.remove(task)
except KeyError as e:
logger.trace(f"{name}: unable to remove task (already removed?): {e}")
def current_tasks() -> Set[asyncio.Task]:
"""Returns the list of currently created/registered tasks."""
return _TASKS

9
src/pipecat/utils/utils.py
View File

@@ -6,9 +6,12 @@
import collections
import itertools
import threading
_COUNTS = collections.defaultdict(itertools.count)
_COUNTS_LOCK = threading.Lock()
_ID = itertools.count()
_ID_LOCK = threading.Lock()
def obj_id() -> int:
@@ -21,7 +24,8 @@ def obj_id() -> int:
>>> obj_id()
2
"""
return next(_ID)
with _ID_LOCK:
return next(_ID)
def obj_count(obj) -> int:
@@ -35,4 +39,5 @@ def obj_count(obj) -> int:
>>> obj_count(new_type())
0
"""
return next(_COUNTS[obj.__class__.__name__])
with _COUNTS_LOCK:
return next(_COUNTS[obj.__class__.__name__])

2
tests/test_pipeline.py
View File

@@ -57,6 +57,7 @@ class TestPipelineTask(unittest.IsolatedAsyncioTestCase):
async def test_task_single(self):
pipeline = Pipeline([IdentityFilter()])
task = PipelineTask(pipeline)
task.set_event_loop(asyncio.get_event_loop())
await task.queue_frame(TextFrame(text="Hello!"))
await task.queue_frames([TextFrame(text="Bye!"), EndFrame()])
@@ -81,6 +82,7 @@ class TestPipelineTask(unittest.IsolatedAsyncioTestCase):
enable_heartbeats=True, heartbeats_period_secs=0.2, observers=[heartbeats_observer]
),
)
task.set_event_loop(asyncio.get_event_loop())
expected_heartbeats = 1.0 / 0.2