From 090bc81ec59e2d00277cca0a548a36570a4ca941 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aleix=20Conchillo=20Flaqu=C3=A9?= Date: Mon, 20 Jan 2025 17:41:21 -0800 Subject: [PATCH 01/15] tests: add some initial run_test() utilities --- .github/workflows/tests.yaml | 2 +- README.md | 8 +-- pyproject.toml | 3 + tests/__init__.py | 0 tests/test_aggregators.py | 6 ++ tests/test_ai_services.py | 6 ++ tests/test_daily_transport_service.py | 6 ++ tests/test_filters.py | 42 ++++++++++++ tests/test_openai_tts.py | 6 ++ tests/test_pipeline.py | 6 ++ tests/test_protobuf_serializer.py | 6 ++ tests/test_websocket_transport.py | 6 ++ tests/utils.py | 96 +++++++++++++++++++++++++++ 13 files changed, 185 insertions(+), 8 deletions(-) create mode 100644 tests/__init__.py create mode 100644 tests/test_filters.py create mode 100644 tests/utils.py diff --git a/.github/workflows/tests.yaml b/.github/workflows/tests.yaml index b806efad4..628f7369b 100644 --- a/.github/workflows/tests.yaml +++ b/.github/workflows/tests.yaml @@ -49,4 +49,4 @@ jobs: - name: Test with pytest run: | source .venv/bin/activate - pytest --ignore-glob="*to_be_updated*" --ignore-glob=*pipeline_source* src tests + pytest diff --git a/README.md b/README.md index 52c6f831b..4f4e08771 100644 --- a/README.md +++ b/README.md @@ -53,12 +53,6 @@ To keep things lightweight, only the core framework is included by default. If y pip install "pipecat-ai[option,...]" ``` -Or you can install all of them with: - -```shell -pip install "pipecat-ai[all]" -``` - Available options include: | Category | Services | Install Command Example | @@ -195,7 +189,7 @@ pip install "path_to_this_repo[option,...]" From the root directory, run: ```shell -pytest --doctest-modules --ignore-glob="*to_be_updated*" --ignore-glob=*pipeline_source* src tests +pytest ``` ## Setting up your editor diff --git a/pyproject.toml b/pyproject.toml index 38bf6d902..fff189c6a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -85,7 +85,10 @@ openrouter = [ "openai~=1.59.6" ] where = ["src"] [tool.pytest.ini_options] +addopts = "--verbose --disable-warnings" +testpaths = ["tests"] pythonpath = ["src"] +asyncio_default_fixture_loop_scope = "function" [tool.setuptools_scm] local_scheme = "no-local-version" diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/test_aggregators.py b/tests/test_aggregators.py index dcf27ad6e..463c2ffed 100644 --- a/tests/test_aggregators.py +++ b/tests/test_aggregators.py @@ -1,3 +1,9 @@ +# +# Copyright (c) 2024-2025 Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + import asyncio import doctest import functools diff --git a/tests/test_ai_services.py b/tests/test_ai_services.py index 13aa20467..4b3a28408 100644 --- a/tests/test_ai_services.py +++ b/tests/test_ai_services.py @@ -1,3 +1,9 @@ +# +# Copyright (c) 2024-2025 Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + import unittest from typing import AsyncGenerator diff --git a/tests/test_daily_transport_service.py b/tests/test_daily_transport_service.py index 8c2788c9e..aabbd733d 100644 --- a/tests/test_daily_transport_service.py +++ b/tests/test_daily_transport_service.py @@ -1,3 +1,9 @@ +# +# Copyright (c) 2024-2025 Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + import unittest diff --git a/tests/test_filters.py b/tests/test_filters.py new file mode 100644 index 000000000..706ca1cca --- /dev/null +++ b/tests/test_filters.py @@ -0,0 +1,42 @@ +# +# Copyright (c) 2024-2025 Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + +import unittest + +from pipecat.frames.frames import ( + TranscriptionFrame, + UserStartedSpeakingFrame, + UserStoppedSpeakingFrame, +) +from pipecat.processors.filters.identity_filter import IdentityFilter +from pipecat.processors.filters.wake_check_filter import WakeCheckFilter +from tests.utils import run_test + + +class TestIdentifyFilter(unittest.IsolatedAsyncioTestCase): + async def test_identity(self): + filter = IdentityFilter() + frames_to_send = [UserStartedSpeakingFrame(), UserStoppedSpeakingFrame()] + expected_returned_frames = [UserStartedSpeakingFrame, UserStoppedSpeakingFrame] + await run_test(filter, frames_to_send, expected_returned_frames) + + +class TestWakeCheckFilter(unittest.IsolatedAsyncioTestCase): + async def test_no_wake_word(self): + filter = WakeCheckFilter(wake_phrases=["Hey, Pipecat"]) + frames_to_send = [TranscriptionFrame(user_id="test", text="Phrase 1", timestamp="")] + expected_returned_frames = [] + await run_test(filter, frames_to_send, expected_returned_frames) + + async def test_wake_word(self): + filter = WakeCheckFilter(wake_phrases=["Hey, Pipecat"]) + frames_to_send = [ + TranscriptionFrame(user_id="test", text="Hey, Pipecat", timestamp=""), + TranscriptionFrame(user_id="test", text="Phrase 1", timestamp=""), + ] + expected_returned_frames = [TranscriptionFrame, TranscriptionFrame] + (received_down, _) = await run_test(filter, frames_to_send, expected_returned_frames) + assert received_down[-1].text == "Phrase 1" diff --git a/tests/test_openai_tts.py b/tests/test_openai_tts.py index 1dc3929a6..171542232 100644 --- a/tests/test_openai_tts.py +++ b/tests/test_openai_tts.py @@ -1,3 +1,9 @@ +# +# Copyright (c) 2024-2025 Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + import asyncio import unittest diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index 7d703d5ed..0acdb034d 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -1,3 +1,9 @@ +# +# Copyright (c) 2024-2025 Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + import asyncio import unittest from unittest.mock import Mock diff --git a/tests/test_protobuf_serializer.py b/tests/test_protobuf_serializer.py index 7f9841622..5e289286f 100644 --- a/tests/test_protobuf_serializer.py +++ b/tests/test_protobuf_serializer.py @@ -1,3 +1,9 @@ +# +# Copyright (c) 2024-2025 Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + import unittest from pipecat.frames.frames import AudioRawFrame, TextFrame, TranscriptionFrame diff --git a/tests/test_websocket_transport.py b/tests/test_websocket_transport.py index b24caa5b9..c918acd8a 100644 --- a/tests/test_websocket_transport.py +++ b/tests/test_websocket_transport.py @@ -1,3 +1,9 @@ +# +# Copyright (c) 2024-2025 Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + # import asyncio # import unittest # from unittest.mock import AsyncMock, patch, Mock diff --git a/tests/utils.py b/tests/utils.py new file mode 100644 index 000000000..95da9f4aa --- /dev/null +++ b/tests/utils.py @@ -0,0 +1,96 @@ +# +# Copyright (c) 2024-2025 Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + +import asyncio +from dataclasses import dataclass +from typing import Sequence, Tuple + +from pipecat.clocks.system_clock import SystemClock +from pipecat.frames.frames import ( + ControlFrame, + Frame, + StartFrame, +) +from pipecat.processors.frame_processor import FrameDirection, FrameProcessor + + +@dataclass +class EndTestFrame(ControlFrame): + pass + + +class QueuedFrameProcessor(FrameProcessor): + def __init__(self, queue: asyncio.Queue, ignore_start: bool = True): + super().__init__() + self._queue = queue + self._ignore_start = ignore_start + + 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) + + +async def run_test( + processor: FrameProcessor, + frames_to_send: Sequence[Frame], + expected_down_frames: Sequence[type], + expected_up_frames: Sequence[type] = [], +) -> Tuple[Sequence[Frame], Sequence[Frame]]: + received_up = asyncio.Queue() + received_down = asyncio.Queue() + up_processor = QueuedFrameProcessor(received_up) + down_processor = QueuedFrameProcessor(received_down) + + up_processor.link(processor) + processor.link(down_processor) + + await processor.queue_frame(StartFrame(clock=SystemClock())) + + for frame in frames_to_send: + await processor.process_frame(frame, FrameDirection.DOWNSTREAM) + + await processor.queue_frame(EndTestFrame()) + await processor.queue_frame(EndTestFrame(), FrameDirection.UPSTREAM) + + # + # Down frames + # + received_down_frames: Sequence[Frame] = [] + running = True + while running: + frame = await received_down.get() + running = not isinstance(frame, EndTestFrame) + if running: + received_down_frames.append(frame) + + print("received DOWN frames =", received_down_frames) + + assert len(received_down_frames) == len(expected_down_frames) + + for real, expected in zip(received_down_frames, expected_down_frames): + assert isinstance(real, expected) + + # + # Up frames + # + received_up_frames: Sequence[Frame] = [] + running = True + while running: + frame = await received_up.get() + running = not isinstance(frame, EndTestFrame) + if running: + received_up_frames.append(frame) + + print("received UP frames =", received_up_frames) + + assert len(received_up_frames) == len(expected_up_frames) + + for real, expected in zip(received_up_frames, expected_up_frames): + assert isinstance(real, expected) + + return (received_down_frames, received_up_frames) From d4d9c3b7aee9ab9d4a7ad608ad9a5d5b2c64595f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aleix=20Conchillo=20Flaqu=C3=A9?= Date: Mon, 20 Jan 2025 18:16:02 -0800 Subject: [PATCH 02/15] tests: fix test_aggregators.py --- tests/test_aggregators.py | 120 ++++++++++---------------------------- 1 file changed, 31 insertions(+), 89 deletions(-) diff --git a/tests/test_aggregators.py b/tests/test_aggregators.py index 463c2ffed..cc03fda52 100644 --- a/tests/test_aggregators.py +++ b/tests/test_aggregators.py @@ -4,125 +4,67 @@ # SPDX-License-Identifier: BSD 2-Clause License # -import asyncio -import doctest -import functools import unittest from pipecat.frames.frames import ( - AudioRawFrame, - EndFrame, - Frame, ImageRawFrame, LLMFullResponseEndFrame, LLMFullResponseStartFrame, + OutputAudioRawFrame, + OutputImageRawFrame, TextFrame, ) -from pipecat.pipeline.parallel_pipeline import ParallelPipeline -from pipecat.pipeline.pipeline import Pipeline from pipecat.processors.aggregators.gated import GatedAggregator from pipecat.processors.aggregators.sentence import SentenceAggregator -from pipecat.processors.text_transformer import StatelessTextTransformer +from tests.utils import run_test -class TestDailyFrameAggregators(unittest.IsolatedAsyncioTestCase): - @unittest.skip("FIXME: This test is failing") +class TestSentenceAggregator(unittest.IsolatedAsyncioTestCase): async def test_sentence_aggregator(self): - sentence = "Hello, world. How are you? I am fine" - expected_sentences = ["Hello, world.", " How are you?", " I am fine "] aggregator = SentenceAggregator() + + sentence = "Hello, world. How are you? I am fine!" + + frames_to_send = [] for word in sentence.split(" "): - async for sentence in aggregator.process_frame(TextFrame(word + " ")): - self.assertIsInstance(sentence, TextFrame) - if isinstance(sentence, TextFrame): - self.assertEqual(sentence.text, expected_sentences.pop(0)) + frames_to_send.append(TextFrame(text=word + " ")) - async for sentence in aggregator.process_frame(EndFrame()): - if len(expected_sentences): - self.assertIsInstance(sentence, TextFrame) - if isinstance(sentence, TextFrame): - self.assertEqual(sentence.text, expected_sentences.pop(0)) - else: - self.assertIsInstance(sentence, EndFrame) + expected_returned_frames = [TextFrame, TextFrame, TextFrame] - self.assertEqual(expected_sentences, []) + (received_down, _) = await run_test(aggregator, frames_to_send, expected_returned_frames) + assert received_down[-3].text == "Hello, world. " + assert received_down[-2].text == "How are you? " + assert received_down[-1].text == "I am fine! " - @unittest.skip("FIXME: This test is failing") - async def test_gated_accumulator(self): + +class TestGatedAggregator(unittest.IsolatedAsyncioTestCase): + async def test_gated_aggregator(self): gated_aggregator = GatedAggregator( gate_open_fn=lambda frame: isinstance(frame, ImageRawFrame), gate_close_fn=lambda frame: isinstance(frame, LLMFullResponseStartFrame), start_open=False, ) - frames = [ + frames_to_send = [ LLMFullResponseStartFrame(), TextFrame("Hello, "), TextFrame("world."), - AudioRawFrame(b"hello"), - ImageRawFrame(b"image", (0, 0)), - AudioRawFrame(b"world"), + OutputAudioRawFrame(audio=b"hello", sample_rate=16000, num_channels=1), + OutputImageRawFrame(image=b"image", size=(0, 0), format="RGB"), + OutputAudioRawFrame(audio=b"world", sample_rate=16000, num_channels=1), LLMFullResponseEndFrame(), ] - expected_output_frames = [ - ImageRawFrame(b"image", (0, 0)), - LLMFullResponseStartFrame(), - TextFrame("Hello, "), - TextFrame("world."), - AudioRawFrame(b"hello"), - AudioRawFrame(b"world"), - LLMFullResponseEndFrame(), + expected_returned_frames = [ + OutputImageRawFrame, + LLMFullResponseStartFrame, + TextFrame, + TextFrame, + OutputAudioRawFrame, + OutputAudioRawFrame, + LLMFullResponseEndFrame, ] - for frame in frames: - async for out_frame in gated_aggregator.process_frame(frame): - self.assertEqual(out_frame, expected_output_frames.pop(0)) - self.assertEqual(expected_output_frames, []) - @unittest.skip("FIXME: This test is failing") - async def test_parallel_pipeline(self): - async def slow_add(sleep_time: float, name: str, x: str): - await asyncio.sleep(sleep_time) - return ":".join([x, name]) - - pipe1_annotation = StatelessTextTransformer(functools.partial(slow_add, 0.1, "pipe1")) - pipe2_annotation = StatelessTextTransformer(functools.partial(slow_add, 0.2, "pipe2")) - sentence_aggregator = SentenceAggregator() - add_dots = StatelessTextTransformer(lambda x: x + ".") - - source = asyncio.Queue() - sink = asyncio.Queue() - pipeline = Pipeline( - [ - ParallelPipeline([[pipe1_annotation], [sentence_aggregator, pipe2_annotation]]), - add_dots, - ], - source, - sink, + (received_down, _) = await run_test( + gated_aggregator, frames_to_send, expected_returned_frames ) - - frames = [TextFrame("Hello, "), TextFrame("world."), EndFrame()] - - expected_output_frames: list[Frame] = [ - TextFrame(text="Hello, :pipe1."), - TextFrame(text="world.:pipe1."), - TextFrame(text="Hello, world.:pipe2."), - EndFrame(), - ] - - for frame in frames: - await source.put(frame) - - await pipeline.run_pipeline() - - while not sink.empty(): - frame = await sink.get() - self.assertEqual(frame, expected_output_frames.pop(0)) - - -def load_tests(loader, tests, ignore): - """Run doctests on the aggregators module.""" - from pipecat.processors import aggregators - - tests.addTests(doctest.DocTestSuite(aggregators)) - return tests From b7dd9748cf90aa049117b12200ed4993796b0990 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aleix=20Conchillo=20Flaqu=C3=A9?= Date: Mon, 20 Jan 2025 18:39:41 -0800 Subject: [PATCH 03/15] serializers: fix special fix initialization --- src/pipecat/serializers/protobuf.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/src/pipecat/serializers/protobuf.py b/src/pipecat/serializers/protobuf.py index f35a1e17a..13d7ded34 100644 --- a/src/pipecat/serializers/protobuf.py +++ b/src/pipecat/serializers/protobuf.py @@ -93,11 +93,11 @@ class ProtobufFrameSerializer(FrameSerializer): id = getattr(args, "id", None) name = getattr(args, "name", None) pts = getattr(args, "pts", None) - if not id and "id" in args_dict: + if "id" in args_dict: del args_dict["id"] - if not name and "name" in args_dict: + if "name" in args_dict: del args_dict["name"] - if not pts and "pts" in args_dict: + if "pts" in args_dict: del args_dict["pts"] # Create the instance @@ -105,10 +105,10 @@ class ProtobufFrameSerializer(FrameSerializer): # Set special fields if id: - setattr(instance, "id", getattr(args, "id", None)) + setattr(instance, "id", id) if name: - setattr(instance, "name", getattr(args, "name", None)) + setattr(instance, "name", name) if pts: - setattr(instance, "pts", getattr(args, "pts", None)) + setattr(instance, "pts", pts) return instance From 2f23693bf3facdcf85729bf335ba965e62286c06 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aleix=20Conchillo=20Flaqu=C3=A9?= Date: Mon, 20 Jan 2025 18:39:59 -0800 Subject: [PATCH 04/15] tests: fix test_protobuf_serializer.py --- tests/test_protobuf_serializer.py | 21 +++++++++++---------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/tests/test_protobuf_serializer.py b/tests/test_protobuf_serializer.py index 5e289286f..88fd7867e 100644 --- a/tests/test_protobuf_serializer.py +++ b/tests/test_protobuf_serializer.py @@ -6,7 +6,11 @@ import unittest -from pipecat.frames.frames import AudioRawFrame, TextFrame, TranscriptionFrame +from pipecat.frames.frames import ( + OutputAudioRawFrame, + TextFrame, + TranscriptionFrame, +) from pipecat.serializers.protobuf import ProtobufFrameSerializer @@ -14,22 +18,19 @@ class TestProtobufFrameSerializer(unittest.IsolatedAsyncioTestCase): def setUp(self): self.serializer = ProtobufFrameSerializer() - @unittest.skip("FIXME: This test is failing") async def test_roundtrip(self): text_frame = TextFrame(text="hello world") frame = self.serializer.deserialize(self.serializer.serialize(text_frame)) - self.assertEqual(frame, TextFrame(text="hello world")) + self.assertEqual(text_frame, frame) transcription_frame = TranscriptionFrame( - text="Hello there!", participantId="123", timestamp="2021-01-01" + text="Hello there!", user_id="123", timestamp="2021-01-01" ) frame = self.serializer.deserialize(self.serializer.serialize(transcription_frame)) self.assertEqual(frame, transcription_frame) - audio_frame = AudioRawFrame(data=b"1234567890") + audio_frame = OutputAudioRawFrame(audio=b"1234567890", sample_rate=16000, num_channels=1) frame = self.serializer.deserialize(self.serializer.serialize(audio_frame)) - self.assertEqual(frame, audio_frame) - - -if __name__ == "__main__": - unittest.main() + self.assertEqual(frame.audio, audio_frame.audio) + self.assertEqual(frame.sample_rate, audio_frame.sample_rate) + self.assertEqual(frame.num_channels, audio_frame.num_channels) From 2e0fb198bf351158c8ed916d8e0407cc0c6e3c6c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aleix=20Conchillo=20Flaqu=C3=A9?= Date: Tue, 21 Jan 2025 09:42:15 -0800 Subject: [PATCH 05/15] frame_processor: allow pushing more frames after EndFrame This can be useful for testing purposes. In real practice, there shouldn't be any frames after an EndFrame is pushed. --- src/pipecat/processors/frame_processor.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/src/pipecat/processors/frame_processor.py b/src/pipecat/processors/frame_processor.py index ae830c52c..228143b2d 100644 --- a/src/pipecat/processors/frame_processor.py +++ b/src/pipecat/processors/frame_processor.py @@ -293,8 +293,7 @@ class FrameProcessor: await self.__input_frame_task async def __input_frame_task_handler(self): - running = True - while running: + while True: try: if self.__should_block_frames: logger.trace(f"{self}: frame processing paused") @@ -311,8 +310,6 @@ class FrameProcessor: if callback: await callback(self, frame, direction) - running = not isinstance(frame, EndFrame) - self.__input_queue.task_done() except asyncio.CancelledError: logger.trace(f"{self}: cancelled input task") @@ -330,12 +327,10 @@ class FrameProcessor: await self.__push_frame_task async def __push_frame_task_handler(self): - running = True - while running: + while True: try: (frame, direction) = await self.__push_queue.get() await self.__internal_push_frame(frame, direction) - running = not isinstance(frame, EndFrame) self.__push_queue.task_done() except asyncio.CancelledError: logger.trace(f"{self}: cancelled push task") From af02f8f1cdc931caa53db3b8722eac83d6bc16ca Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aleix=20Conchillo=20Flaqu=C3=A9?= Date: Tue, 21 Jan 2025 09:43:33 -0800 Subject: [PATCH 06/15] filters(frame_filter): allow more than one frame --- src/pipecat/processors/filters/frame_filter.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pipecat/processors/filters/frame_filter.py b/src/pipecat/processors/filters/frame_filter.py index a2fb90505..43e2dab95 100644 --- a/src/pipecat/processors/filters/frame_filter.py +++ b/src/pipecat/processors/filters/frame_filter.py @@ -11,7 +11,7 @@ from pipecat.processors.frame_processor import FrameDirection, FrameProcessor class FrameFilter(FrameProcessor): - def __init__(self, types: Tuple[Type[Frame]]): + def __init__(self, types: Tuple[Type[Frame], ...]): super().__init__() self._types = types From 3c970a3ceecd8cc2287441c41bd5e6a35525ccef Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aleix=20Conchillo=20Flaqu=C3=A9?= Date: Tue, 21 Jan 2025 09:43:57 -0800 Subject: [PATCH 07/15] tests: add more filter tests --- tests/test_filters.py | 54 ++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 53 insertions(+), 1 deletion(-) diff --git a/tests/test_filters.py b/tests/test_filters.py index 706ca1cca..e831f4071 100644 --- a/tests/test_filters.py +++ b/tests/test_filters.py @@ -4,16 +4,22 @@ # SPDX-License-Identifier: BSD 2-Clause License # +import asyncio import unittest from pipecat.frames.frames import ( + EndFrame, + Frame, + TextFrame, TranscriptionFrame, UserStartedSpeakingFrame, UserStoppedSpeakingFrame, ) +from pipecat.processors.filters.frame_filter import FrameFilter +from pipecat.processors.filters.function_filter import FunctionFilter from pipecat.processors.filters.identity_filter import IdentityFilter from pipecat.processors.filters.wake_check_filter import WakeCheckFilter -from tests.utils import run_test +from tests.utils import EndTestFrame, run_test class TestIdentifyFilter(unittest.IsolatedAsyncioTestCase): @@ -24,6 +30,52 @@ class TestIdentifyFilter(unittest.IsolatedAsyncioTestCase): await run_test(filter, frames_to_send, expected_returned_frames) +class TestFrameFilter(unittest.IsolatedAsyncioTestCase): + async def test_text_frame(self): + filter = FrameFilter(types=(TextFrame, EndTestFrame)) + frames_to_send = [TextFrame(text="Hello Pipecat!")] + expected_returned_frames = [TextFrame] + await run_test(filter, frames_to_send, expected_returned_frames) + + async def test_end_frame(self): + filter = FrameFilter(types=(EndFrame, EndTestFrame)) + frames_to_send = [EndFrame()] + expected_returned_frames = [EndFrame] + await run_test(filter, frames_to_send, expected_returned_frames) + + async def test_system_frame(self): + filter = FrameFilter(types=(EndTestFrame,)) + frames_to_send = [UserStartedSpeakingFrame()] + expected_returned_frames = [UserStartedSpeakingFrame] + await run_test(filter, frames_to_send, expected_returned_frames) + + +class TestFunctionFilter(unittest.IsolatedAsyncioTestCase): + async def test_passthrough(self): + async def passthrough(frame: Frame): + return True + + filter = FunctionFilter(filter=passthrough) + frames_to_send = [TextFrame(text="Hello Pipecat!")] + expected_returned_frames = [TextFrame] + await run_test(filter, frames_to_send, expected_returned_frames) + + async def test_no_passthrough(self): + async def no_passthrough(frame: Frame): + return False + + filter = FunctionFilter(filter=no_passthrough) + frames_to_send = [TextFrame(text="Hello Pipecat!")] + expected_returned_frames = [TextFrame] + try: + await asyncio.wait_for( + run_test(filter, frames_to_send, expected_returned_frames), timeout=0.5 + ) + assert False + except asyncio.TimeoutError: + pass + + class TestWakeCheckFilter(unittest.IsolatedAsyncioTestCase): async def test_no_wake_word(self): filter = WakeCheckFilter(wake_phrases=["Hey, Pipecat"]) From 177cb2ca8b72a1a3f8cdf55f5e8d189978c0ab2f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aleix=20Conchillo=20Flaqu=C3=A9?= Date: Tue, 21 Jan 2025 09:57:54 -0800 Subject: [PATCH 08/15] tests: initial pipeline and parallelpipeline tests --- tests/test_pipeline.py | 132 ++++++++++------------------------------- 1 file changed, 30 insertions(+), 102 deletions(-) diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index 0acdb034d..dddd1c0de 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -4,119 +4,47 @@ # SPDX-License-Identifier: BSD 2-Clause License # -import asyncio import unittest -from unittest.mock import Mock -from pipecat.frames.frames import EndFrame, TextFrame +from pipecat.frames.frames import TextFrame +from pipecat.pipeline.parallel_pipeline import ParallelPipeline from pipecat.pipeline.pipeline import Pipeline -from pipecat.processors.aggregators.sentence import SentenceAggregator -from pipecat.processors.frame_processor import FrameProcessor -from pipecat.processors.text_transformer import StatelessTextTransformer +from pipecat.processors.filters.identity_filter import IdentityFilter +from tests.utils import run_test -class TestDailyPipeline(unittest.IsolatedAsyncioTestCase): - @unittest.skip("FIXME: This test is failing") - async def test_pipeline_simple(self): - aggregator = SentenceAggregator() +class TestPipeline(unittest.IsolatedAsyncioTestCase): + async def test_pipeline_single(self): + pipeline = Pipeline([IdentityFilter()]) - outgoing_queue = asyncio.Queue() - incoming_queue = asyncio.Queue() - pipeline = Pipeline([aggregator], incoming_queue, outgoing_queue) + frames_to_send = [TextFrame(text="Hello from Pipecat!")] + expected_returned_frames = [TextFrame] + await run_test(pipeline, frames_to_send, expected_returned_frames) - await incoming_queue.put(TextFrame("Hello, ")) - await incoming_queue.put(TextFrame("world.")) - await incoming_queue.put(EndFrame()) + async def test_pipeline_multiple(self): + identity1 = IdentityFilter() + identity2 = IdentityFilter() + identity3 = IdentityFilter() - await pipeline.run_pipeline() + pipeline = Pipeline([identity1, identity2, identity3]) - self.assertEqual(await outgoing_queue.get(), TextFrame("Hello, world.")) - self.assertIsInstance(await outgoing_queue.get(), EndFrame) - - @unittest.skip("FIXME: This test is failing") - async def test_pipeline_multiple_stages(self): - sentence_aggregator = SentenceAggregator() - to_upper = StatelessTextTransformer(lambda x: x.upper()) - add_space = StatelessTextTransformer(lambda x: x + " ") - - outgoing_queue = asyncio.Queue() - incoming_queue = asyncio.Queue() - pipeline = Pipeline( - [add_space, sentence_aggregator, to_upper], incoming_queue, outgoing_queue - ) - - sentence = "Hello, world. It's me, a pipeline." - for c in sentence: - await incoming_queue.put(TextFrame(c)) - await incoming_queue.put(EndFrame()) - - await pipeline.run_pipeline() - - self.assertEqual(await outgoing_queue.get(), TextFrame("H E L L O , W O R L D .")) - self.assertEqual( - await outgoing_queue.get(), - TextFrame(" I T ' S M E , A P I P E L I N E ."), - ) - # leftover little bit because of the spacing - self.assertEqual( - await outgoing_queue.get(), - TextFrame(" "), - ) - self.assertIsInstance(await outgoing_queue.get(), EndFrame) + frames_to_send = [TextFrame(text="Hello from Pipecat!")] + expected_returned_frames = [TextFrame] + await run_test(pipeline, frames_to_send, expected_returned_frames) -class TestLogFrame(unittest.TestCase): - class MockProcessor(FrameProcessor): - def __init__(self, name): - self.name = name +class TestParallelPipeline(unittest.IsolatedAsyncioTestCase): + async def test_parallel_single(self): + pipeline = ParallelPipeline([IdentityFilter()]) - def __str__(self): - return self.name + frames_to_send = [TextFrame(text="Hello from Pipecat!")] + expected_returned_frames = [TextFrame] + await run_test(pipeline, frames_to_send, expected_returned_frames) - def setUp(self): - self.processor1 = self.MockProcessor("processor1") - self.processor2 = self.MockProcessor("processor2") - self.pipeline = Pipeline(processors=[self.processor1, self.processor2]) - self.pipeline._name = "MyClass" - self.pipeline._logger = Mock() + async def test_parallel_multiple(self): + """Should only passthrough one instance of TextFrame.""" + pipeline = ParallelPipeline([IdentityFilter()], [IdentityFilter()]) - @unittest.skip("FIXME: This test is failing") - def test_log_frame_from_source(self): - frame = Mock(__class__=Mock(__name__="MyFrame")) - self.pipeline._log_frame(frame, depth=1) - self.pipeline._logger.debug.assert_called_once_with( - "MyClass source -> MyFrame -> processor1" - ) - - @unittest.skip("FIXME: This test is failing") - def test_log_frame_to_sink(self): - frame = Mock(__class__=Mock(__name__="MyFrame")) - self.pipeline._log_frame(frame, depth=3) - self.pipeline._logger.debug.assert_called_once_with( - "MyClass processor2 -> MyFrame -> sink" - ) - - @unittest.skip("FIXME: This test is failing") - def test_log_frame_repeated_log(self): - frame = Mock(__class__=Mock(__name__="MyFrame")) - self.pipeline._log_frame(frame, depth=2) - self.pipeline._logger.debug.assert_called_once_with( - "MyClass processor1 -> MyFrame -> processor2" - ) - self.pipeline._log_frame(frame, depth=2) - self.pipeline._logger.debug.assert_called_with("MyClass ... repeated") - - @unittest.skip("FIXME: This test is failing") - def test_log_frame_reset_repeated_log(self): - frame1 = Mock(__class__=Mock(__name__="MyFrame1")) - frame2 = Mock(__class__=Mock(__name__="MyFrame2")) - self.pipeline._log_frame(frame1, depth=2) - self.pipeline._logger.debug.assert_called_once_with( - "MyClass processor1 -> MyFrame1 -> processor2" - ) - self.pipeline._log_frame(frame1, depth=2) - self.pipeline._logger.debug.assert_called_with("MyClass ... repeated") - self.pipeline._log_frame(frame2, depth=2) - self.pipeline._logger.debug.assert_called_with( - "MyClass processor1 -> MyFrame2 -> processor2" - ) + frames_to_send = [TextFrame(text="Hello from Pipecat!")] + expected_returned_frames = [TextFrame] + await run_test(pipeline, frames_to_send, expected_returned_frames) From a27fe4bde2ede48187bf7faa3b4e67555720857b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aleix=20Conchillo=20Flaqu=C3=A9?= Date: Tue, 21 Jan 2025 10:06:14 -0800 Subject: [PATCH 09/15] tests: move test_ai_services to test_utils_string --- tests/test_openai_tts.py | 43 ------------------- ...st_ai_services.py => test_utils_string.py} | 27 +----------- 2 files changed, 2 insertions(+), 68 deletions(-) delete mode 100644 tests/test_openai_tts.py rename tests/{test_ai_services.py => test_utils_string.py} (64%) diff --git a/tests/test_openai_tts.py b/tests/test_openai_tts.py deleted file mode 100644 index 171542232..000000000 --- a/tests/test_openai_tts.py +++ /dev/null @@ -1,43 +0,0 @@ -# -# Copyright (c) 2024-2025 Daily -# -# SPDX-License-Identifier: BSD 2-Clause License -# - -import asyncio -import unittest - -import openai -import pyaudio -from dotenv import load_dotenv - -from pipecat.frames.frames import AudioRawFrame, ErrorFrame -from pipecat.services.openai import OpenAITTSService - -load_dotenv() - - -class TestWhisperOpenAIService(unittest.IsolatedAsyncioTestCase): - @unittest.skip("FIXME: This test is failing") - async def test_whisper_tts(self): - pa = pyaudio.PyAudio() - stream = pa.open(format=pyaudio.paInt16, channels=1, rate=24_000, output=True) - - tts = OpenAITTSService(voice="nova") - - async for frame in tts.run_tts("Hello, there. Nice to meet you, seems to work well"): - self.assertIsInstance(frame, AudioRawFrame) - stream.write(frame.audio) - - await asyncio.sleep(0.5) - stream.stop_stream() - pa.terminate() - - tts = OpenAITTSService(voice="invalid_voice") - with self.assertRaises(openai.BadRequestError): - async for frame in tts.run_tts("wont work"): - self.assertIsInstance(frame, ErrorFrame) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/test_ai_services.py b/tests/test_utils_string.py similarity index 64% rename from tests/test_ai_services.py rename to tests/test_utils_string.py index 4b3a28408..9587544a1 100644 --- a/tests/test_ai_services.py +++ b/tests/test_utils_string.py @@ -5,30 +5,11 @@ # import unittest -from typing import AsyncGenerator -from pipecat.frames.frames import EndFrame, Frame, TextFrame -from pipecat.services.ai_services import AIService, match_endofsentence +from pipecat.utils.string import match_endofsentence -class SimpleAIService(AIService): - async def process_frame(self, frame: Frame) -> AsyncGenerator[Frame, None]: - yield frame - - -class TestBaseAIService(unittest.IsolatedAsyncioTestCase): - async def test_simple_processing(self): - service = SimpleAIService() - - input_frames = [TextFrame("hello"), EndFrame()] - - output_frames = [] - for input_frame in input_frames: - async for output_frame in service.process_frame(input_frame): - output_frames.append(output_frame) - - self.assertEqual(input_frames, output_frames) - +class TestUtilsString(unittest.IsolatedAsyncioTestCase): async def test_endofsentence(self): assert match_endofsentence("This is a sentence.") assert match_endofsentence("This is a sentence! ") @@ -57,7 +38,3 @@ class TestBaseAIService(unittest.IsolatedAsyncioTestCase): for i in chinese_sentences: assert match_endofsentence(i) assert not match_endofsentence("你好，") - - -if __name__ == "__main__": - unittest.main() From 0d6c68013365e47c9327c24dcb8a1e2d19dad46a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aleix=20Conchillo=20Flaqu=C3=A9?= Date: Tue, 21 Jan 2025 10:11:50 -0800 Subject: [PATCH 10/15] README: add unit tests badge --- .github/workflows/tests.yaml | 2 +- README.md | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/tests.yaml b/.github/workflows/tests.yaml index 628f7369b..bd6961dd5 100644 --- a/.github/workflows/tests.yaml +++ b/.github/workflows/tests.yaml @@ -1,4 +1,4 @@ -name: test +name: tests on: workflow_dispatch: diff --git a/README.md b/README.md index 4f4e08771..f3f1401a1 100644 --- a/README.md +++ b/README.md @@ -2,7 +2,7 @@  pipecat -[![PyPI](https://img.shields.io/pypi/v/pipecat-ai)](https://pypi.org/project/pipecat-ai) [![Docs](https://img.shields.io/badge/Documentation-blue)](https://docs.pipecat.ai) [![Discord](https://img.shields.io/discord/1239284677165056021)](https://discord.gg/pipecat) +[![PyPI](https://img.shields.io/pypi/v/pipecat-ai)](https://pypi.org/project/pipecat-ai) ![Tests](https://github.com/pipecat-ai/pipecat/actions/workflows/tests.yaml/badge.svg) [![Docs](https://img.shields.io/badge/Documentation-blue)](https://docs.pipecat.ai) [![Discord](https://img.shields.io/discord/1239284677165056021)](https://discord.gg/pipecat) Pipecat is an open source Python framework for building voice and multimodal conversational agents. It handles the complex orchestration of AI services, network transport, audio processing, and multimodal interactions, letting you focus on creating engaging experiences. From 76884877dd599b86f355400e882faec36809f734 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aleix=20Conchillo=20Flaqu=C3=A9?= Date: Tue, 21 Jan 2025 10:23:19 -0800 Subject: [PATCH 11/15] tests: add pytest-asyncio dependency --- dev-requirements.txt | 1 + pyproject.toml | 2 +- tests/test_langchain.py | 4 ---- 3 files changed, 2 insertions(+), 5 deletions(-) diff --git a/dev-requirements.txt b/dev-requirements.txt index 2a9c68581..e3f52f9cd 100644 --- a/dev-requirements.txt +++ b/dev-requirements.txt @@ -4,6 +4,7 @@ pip-tools~=7.4.1 pre-commit~=4.0.1 pyright~=1.1.392 pytest~=8.3.4 +pytest-asyncio~=0.25.2 ruff~=0.9.1 setuptools~=75.8.0 setuptools_scm~=8.1.0 diff --git a/pyproject.toml b/pyproject.toml index fff189c6a..1fd2033a3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -85,7 +85,7 @@ openrouter = [ "openai~=1.59.6" ] where = ["src"] [tool.pytest.ini_options] -addopts = "--verbose --disable-warnings" +addopts = "--verbose" testpaths = ["tests"] pythonpath = ["src"] asyncio_default_fixture_loop_scope = "function" diff --git a/tests/test_langchain.py b/tests/test_langchain.py index 22dc9d38d..b09eee86e 100644 --- a/tests/test_langchain.py +++ b/tests/test_langchain.py @@ -93,7 +93,3 @@ class TestLangchain(unittest.IsolatedAsyncioTestCase): # This next one would fail with: # AssertionError: ' H e l l o d e a r h u m a n' != 'Hello dear human' # self.assertEqual(tma_out.messages[-1]["content"], self.expected_response) - - -if __name__ == "__main__": - unittest.main() From dd21b424d6808ee83cb889d7f1e7b955226cfea0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aleix=20Conchillo=20Flaqu=C3=A9?= Date: Tue, 21 Jan 2025 10:27:34 -0800 Subject: [PATCH 12/15] pyproject: ignore 'audioop' deprecation warning --- pyproject.toml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 1fd2033a3..57d1f2703 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -89,6 +89,9 @@ addopts = "--verbose" testpaths = ["tests"] pythonpath = ["src"] asyncio_default_fixture_loop_scope = "function" +filterwarnings = [ + "ignore:'audioop' is deprecated:DeprecationWarning", +] [tool.setuptools_scm] local_scheme = "no-local-version" From bd6f82cf945c38c40f31ed9467d2ac09350b6a73 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aleix=20Conchillo=20Flaqu=C3=A9?= Date: Tue, 21 Jan 2025 11:43:21 -0800 Subject: [PATCH 13/15] task: allow specifying heartbeat period --- CHANGELOG.md | 5 +++++ src/pipecat/pipeline/task.py | 5 +++-- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e18113c0e..e0d48cbaa 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added + +- It is now possible to specify the period of the `PipelineTask` heartbeat + frames with `heartbeats_period_secs`. + ### Fixed - Fixed a type error when using `voice_settings` in `ElevenLabsHttpTTSService`. diff --git a/src/pipecat/pipeline/task.py b/src/pipecat/pipeline/task.py index 9263352e4..b3cc1438e 100644 --- a/src/pipecat/pipeline/task.py +++ b/src/pipecat/pipeline/task.py @@ -45,6 +45,7 @@ class PipelineParams(BaseModel): send_initial_empty_metrics: bool = True report_only_initial_ttfb: bool = False observers: List[BaseObserver] = [] + heartbeats_period_secs: float = HEARTBEAT_SECONDS class Source(FrameProcessor): @@ -315,7 +316,7 @@ class PipelineTask: async def _heartbeat_push_handler(self): """ - This tasks pushes a heartbeat frame every HEARTBEAT_SECONDS. + This tasks pushes a heartbeat frame every heartbeat period. """ while True: try: @@ -323,7 +324,7 @@ class PipelineTask: # task will just stop waiting for the pipeline to finish not # allowing more frames to be pushed. await self._source.queue_frame(HeartbeatFrame(timestamp=self._clock.get_time())) - await asyncio.sleep(HEARTBEAT_SECONDS) + await asyncio.sleep(self._params.heartbeats_period_secs) except asyncio.CancelledError: break From ab4221a4db6278b9ddf23c6326a5e30f7656d832 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aleix=20Conchillo=20Flaqu=C3=A9?= Date: Tue, 21 Jan 2025 11:43:54 -0800 Subject: [PATCH 14/15] task: added BaseTask --- src/pipecat/pipeline/base_task.py | 56 +++++++++++++++++++++++++++++++ src/pipecat/pipeline/task.py | 5 +-- 2 files changed, 59 insertions(+), 2 deletions(-) create mode 100644 src/pipecat/pipeline/base_task.py diff --git a/src/pipecat/pipeline/base_task.py b/src/pipecat/pipeline/base_task.py new file mode 100644 index 000000000..51971e8ce --- /dev/null +++ b/src/pipecat/pipeline/base_task.py @@ -0,0 +1,56 @@ +# +# Copyright (c) 2024–2025, Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + +from abc import ABC, abstractmethod +from typing import AsyncIterable, Iterable + +from pipecat.frames.frames import Frame + + +class BaseTask(ABC): + @abstractmethod + def has_finished(self) -> bool: + """Indicates whether the tasks has finished. That is, all processors + have stopped. + + """ + pass + + @abstractmethod + async def stop_when_done(self): + """This is a helper function that sends an EndFrame to the pipeline in + order to stop the task after everything in it has been processed. + + """ + pass + + @abstractmethod + async def cancel(self): + """ + Stops the running pipeline immediately. + """ + pass + + @abstractmethod + async def run(self): + """ + Starts running the given pipeline. + """ + pass + + @abstractmethod + async def queue_frame(self, frame: Frame): + """ + 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. + """ + pass diff --git a/src/pipecat/pipeline/task.py b/src/pipecat/pipeline/task.py index b3cc1438e..82918c239 100644 --- a/src/pipecat/pipeline/task.py +++ b/src/pipecat/pipeline/task.py @@ -27,6 +27,7 @@ from pipecat.frames.frames import ( from pipecat.metrics.metrics import ProcessingMetricsData, TTFBMetricsData from pipecat.observers.base_observer import BaseObserver 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.utils import obj_count, obj_id @@ -86,7 +87,7 @@ class Sink(FrameProcessor): await self._down_queue.put(frame) -class PipelineTask: +class PipelineTask(BaseTask): def __init__( self, pipeline: BasePipeline, @@ -122,7 +123,7 @@ class PipelineTask: self._observer = TaskObserver(params.observers) - def has_finished(self): + def has_finished(self) -> bool: """Indicates whether the tasks has finished. That is, all processors have stopped. From 401d3ff2670047417f7d24d68ff2a2d25e0303c0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aleix=20Conchillo=20Flaqu=C3=A9?= Date: Tue, 21 Jan 2025 11:44:13 -0800 Subject: [PATCH 15/15] tests: added PipelineTask tests --- tests/test_pipeline.py | 46 ++++++++++++++++++++++++++++++++++++++++-- tests/utils.py | 26 +++++++++++++++++++++++- 2 files changed, 69 insertions(+), 3 deletions(-) diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index dddd1c0de..fefdfa4b6 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -4,13 +4,16 @@ # SPDX-License-Identifier: BSD 2-Clause License # +import asyncio import unittest -from pipecat.frames.frames import TextFrame +from pipecat.frames.frames import EndFrame, HeartbeatFrame, TextFrame from pipecat.pipeline.parallel_pipeline import ParallelPipeline from pipecat.pipeline.pipeline import Pipeline +from pipecat.pipeline.task import PipelineParams, PipelineTask from pipecat.processors.filters.identity_filter import IdentityFilter -from tests.utils import run_test +from pipecat.processors.frame_processor import FrameProcessor +from tests.utils import HeartbeatsObserver, run_test class TestPipeline(unittest.IsolatedAsyncioTestCase): @@ -48,3 +51,42 @@ class TestParallelPipeline(unittest.IsolatedAsyncioTestCase): frames_to_send = [TextFrame(text="Hello from Pipecat!")] expected_returned_frames = [TextFrame] await run_test(pipeline, frames_to_send, expected_returned_frames) + + +class TestPipelineTask(unittest.IsolatedAsyncioTestCase): + async def test_task_single(self): + pipeline = Pipeline([IdentityFilter()]) + task = PipelineTask(pipeline) + + await task.queue_frame(TextFrame(text="Hello!")) + await task.queue_frames([TextFrame(text="Bye!"), EndFrame()]) + await task.run() + assert task.has_finished() + + async def test_task_heartbeats(self): + heartbeats_counter = 0 + + async def heartbeat_received(processor: FrameProcessor, heartbeat: HeartbeatFrame): + nonlocal heartbeats_counter + heartbeats_counter += 1 + + identity = IdentityFilter() + pipeline = Pipeline([identity]) + heartbeats_observer = HeartbeatsObserver( + target=identity, heartbeat_callback=heartbeat_received + ) + task = PipelineTask( + pipeline, + params=PipelineParams( + enable_heartbeats=True, heartbeats_period_secs=0.2, observers=[heartbeats_observer] + ), + ) + + expected_heartbeats = 1.0 / 0.2 + + await task.queue_frame(TextFrame(text="Hello!")) + try: + await asyncio.wait_for(task.run(), timeout=1.0) + except asyncio.TimeoutError: + pass + assert heartbeats_counter == expected_heartbeats diff --git a/tests/utils.py b/tests/utils.py index 95da9f4aa..3b6a0d009 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -6,14 +6,16 @@ import asyncio from dataclasses import dataclass -from typing import Sequence, Tuple +from typing import Awaitable, Callable, Sequence, Tuple from pipecat.clocks.system_clock import SystemClock from pipecat.frames.frames import ( ControlFrame, Frame, + HeartbeatFrame, StartFrame, ) +from pipecat.observers.base_observer import BaseObserver from pipecat.processors.frame_processor import FrameDirection, FrameProcessor @@ -22,6 +24,28 @@ class EndTestFrame(ControlFrame): pass +class HeartbeatsObserver(BaseObserver): + def __init__( + self, + *, + target: FrameProcessor, + heartbeat_callback: Callable[[FrameProcessor, HeartbeatFrame], Awaitable[None]], + ): + self._target = target + self._callback = heartbeat_callback + + async def on_push_frame( + self, + src: FrameProcessor, + dst: FrameProcessor, + frame: Frame, + direction: FrameDirection, + timestamp: int, + ): + if src == self._target and isinstance(frame, HeartbeatFrame): + await self._callback(self._target, frame) + + class QueuedFrameProcessor(FrameProcessor): def __init__(self, queue: asyncio.Queue, ignore_start: bool = True): super().__init__()