diff --git a/src/dailyai/pipeline/merge_pipeline.py b/src/dailyai/pipeline/merge_pipeline.py index 736903e9d..e4b865f4a 100644 --- a/src/dailyai/pipeline/merge_pipeline.py +++ b/src/dailyai/pipeline/merge_pipeline.py @@ -12,9 +12,10 @@ class SequentialMergePipeline(Pipeline): self.pipelines = pipelines async def run_pipeline(self): - for pipeline in self.pipelines: + for idx, pipeline in enumerate(self.pipelines): while True: frame = await pipeline.sink.get() + print(idx, frame) if isinstance( frame, EndFrame) or isinstance( frame, EndPipeFrame): diff --git a/src/dailyai/transports/abstract_transport.py b/src/dailyai/transports/abstract_transport.py index 6b9c9549c..f5ecc5880 100644 --- a/src/dailyai/transports/abstract_transport.py +++ b/src/dailyai/transports/abstract_transport.py @@ -1,12 +1,41 @@ from abc import abstractmethod +import asyncio +import logging +import time +from dailyai.pipeline.frame_processor import FrameProcessor from dailyai.pipeline.pipeline import Pipeline class AbstractTransport: def __init__(self, **kwargs): - pass + self.send_queue = asyncio.Queue() + self.receive_queue = asyncio.Queue() + self.completed_queue = asyncio.Queue() + + duration_minutes = kwargs.get("duration_minutes") or 10 + self._expiration = time.time() + duration_minutes * 60 + + self._mic_enabled = kwargs.get("mic_enabled") or False + self._mic_sample_rate = kwargs.get("mic_sample_rate") or 16000 + self._camera_enabled = kwargs.get("camera_enabled") or False + self._camera_width = kwargs.get("camera_width") or 1024 + self._camera_height = kwargs.get("camera_height") or 768 + self._speaker_enabled = kwargs.get("speaker_enabled") or False + self._speaker_sample_rate = kwargs.get("speaker_sample_rate") or 16000 + self._fps = kwargs.get("fps") or 8 + + self._logger: logging.Logger = logging.getLogger("dailyai.transport") @abstractmethod async def run(self, pipeline: Pipeline, override_pipeline_source_queue=True): pass + + @abstractmethod + async def run_interruptible_pipeline( + self, + pipeline: Pipeline, + pre_processor: FrameProcessor | None = None, + post_processor: FrameProcessor | None = None, + ): + pass diff --git a/src/dailyai/transports/threaded_transport.py b/src/dailyai/transports/threaded_transport.py index 50bd56cda..d529731cd 100644 --- a/src/dailyai/transports/threaded_transport.py +++ b/src/dailyai/transports/threaded_transport.py @@ -27,6 +27,7 @@ from dailyai.pipeline.frames import ( ) from dailyai.pipeline.pipeline import Pipeline from dailyai.services.ai_services import TTSService +from dailyai.transports.abstract_transport import AbstractTransport torch.set_num_threads(1) @@ -72,20 +73,13 @@ class VADState(Enum): STOPPING = 4 -class ThreadedTransport: +class ThreadedTransport(AbstractTransport): def __init__( self, **kwargs, ) -> None: - self._mic_enabled = kwargs.get("mic_enabled") or False - self._mic_sample_rate = kwargs.get("mic_sample_rate") or 16000 - self._camera_enabled = kwargs.get("camera_enabled") or False - self._camera_width = kwargs.get("camera_width") or 1024 - self._camera_height = kwargs.get("camera_height") or 768 - self._speaker_enabled = kwargs.get("speaker_enabled") or False - self._speaker_sample_rate = kwargs.get("speaker_sample_rate") or 16000 - self._fps = kwargs.get("fps") or 8 + super().__init__(**kwargs) self._vad_start_s = kwargs.get("vad_start_s") or 0.2 self._vad_stop_s = kwargs.get("vad_stop_s") or 0.8 self._context = kwargs.get("context") or [] @@ -105,14 +99,6 @@ class ThreadedTransport: self._vad_state = VADState.QUIET self._user_is_speaking = False - duration_minutes = kwargs.get("duration_minutes") or 10 - self._expiration = time.time() + duration_minutes * 60 - - self.send_queue = asyncio.Queue() - self.receive_queue = asyncio.Queue() - - self.completed_queue = asyncio.Queue() - self._threadsafe_send_queue = queue.Queue() self._images = None @@ -125,8 +111,6 @@ class ThreadedTransport: self._stop_threads = threading.Event() self._is_interrupted = threading.Event() - self._logger: logging.Logger = logging.getLogger() - async def run(self, pipeline: Pipeline | None = None, override_pipeline_source_queue=True): self._prerun() diff --git a/src/dailyai/transports/websocket_transport.py b/src/dailyai/transports/websocket_transport.py index 84c185a3d..0784eb89f 100644 --- a/src/dailyai/transports/websocket_transport.py +++ b/src/dailyai/transports/websocket_transport.py @@ -7,6 +7,7 @@ from dailyai.pipeline.frame_processor import FrameProcessor from dailyai.pipeline.frames import AudioFrame, ControlFrame, EndFrame, Frame, TTSEndFrame, TTSStartFrame, TextFrame from dailyai.pipeline.pipeline import Pipeline from dailyai.serializers.protobuf_serializer import ProtobufFrameSerializer +from dailyai.transports.abstract_transport import AbstractTransport from dailyai.transports.threaded_transport import ThreadedTransport @@ -45,7 +46,7 @@ class WebSocketFrameProcessor(FrameProcessor): yield frame -class WebsocketTransport(ThreadedTransport): +class WebsocketTransport(AbstractTransport): def __init__(self, **kwargs): super().__init__(**kwargs) self._sample_width = kwargs.get("sample_width", 2)