diff --git a/tests/test_context_aggregators_universal.py b/tests/test_context_aggregators_universal.py index 94acf383a..6aa38d362 100644 --- a/tests/test_context_aggregators_universal.py +++ b/tests/test_context_aggregators_universal.py @@ -13,10 +13,17 @@ from pipecat.frames.frames import ( FunctionCallResultFrame, FunctionCallsStartedFrame, InterruptionFrame, + LLMContextAssistantTimestampFrame, LLMContextFrame, + LLMFullResponseEndFrame, + LLMFullResponseStartFrame, LLMMessagesAppendFrame, LLMMessagesUpdateFrame, LLMRunFrame, + LLMTextFrame, + LLMThoughtEndFrame, + LLMThoughtStartFrame, + LLMThoughtTextFrame, TranscriptionFrame, UserStartedSpeakingFrame, UserStoppedSpeakingFrame, @@ -26,6 +33,9 @@ from pipecat.frames.frames import ( from pipecat.pipeline.pipeline import Pipeline from pipecat.processors.aggregators.llm_context import LLMContext from pipecat.processors.aggregators.llm_response_universal import ( + AssistantThoughtMessage, + AssistantTurnStoppedMessage, + LLMAssistantAggregator, LLMUserAggregator, LLMUserAggregatorParams, ) @@ -333,3 +343,172 @@ class TestLLMUserAggregator(unittest.IsolatedAsyncioTestCase): # The user mute strategies should have muted the user. self.assertFalse(user_turn) + + +class TestLLMAssistantAggregator(unittest.IsolatedAsyncioTestCase): + async def test_empty(self): + context = LLMContext() + + aggregator = LLMAssistantAggregator(context) + + should_start = None + should_stop = None + stop_message = None + + @aggregator.event_handler("on_assistant_turn_started") + async def on_assistant_turn_started(aggregator): + nonlocal should_start + should_start = True + + @aggregator.event_handler("on_assistant_turn_stopped") + async def on_assistant_turn_stopped(aggregator, message: AssistantTurnStoppedMessage): + nonlocal should_stop, stop_message + should_stop = True + stop_message = message + + frames_to_send = [LLMFullResponseStartFrame(), LLMFullResponseEndFrame()] + await run_test(aggregator, frames_to_send=frames_to_send) + self.assertTrue(should_start) + self.assertIsNone(should_stop) + self.assertIsNone(stop_message) + + async def test_simple(self): + context = LLMContext() + + aggregator = LLMAssistantAggregator(context) + + should_start = None + should_stop = None + stop_message = None + + @aggregator.event_handler("on_assistant_turn_started") + async def on_assistant_turn_started(aggregator): + nonlocal should_start + should_start = True + + @aggregator.event_handler("on_assistant_turn_stopped") + async def on_assistant_turn_stopped(aggregator, message: AssistantTurnStoppedMessage): + nonlocal should_stop, stop_message + should_stop = True + stop_message = message + + frames_to_send = [ + LLMFullResponseStartFrame(), + LLMTextFrame("Hello from Pipecat!"), + LLMFullResponseEndFrame(), + ] + expected_down_frames = [LLMContextFrame, LLMContextAssistantTimestampFrame] + await run_test( + aggregator, + frames_to_send=frames_to_send, + expected_down_frames=expected_down_frames, + ) + self.assertTrue(should_start) + self.assertTrue(should_stop) + self.assertEqual(stop_message.content, "Hello from Pipecat!") + + async def test_multiple(self): + context = LLMContext() + + aggregator = LLMAssistantAggregator(context) + + should_start = None + should_stop = None + stop_message = None + + @aggregator.event_handler("on_assistant_turn_started") + async def on_assistant_turn_started(aggregator): + nonlocal should_start + should_start = True + + @aggregator.event_handler("on_assistant_turn_stopped") + async def on_assistant_turn_stopped(aggregator, message: AssistantTurnStoppedMessage): + nonlocal should_stop, stop_message + should_stop = True + stop_message = message + + frames_to_send = [ + LLMFullResponseStartFrame(), + LLMTextFrame("Hello "), + LLMTextFrame("from "), + LLMTextFrame("Pipecat!"), + LLMFullResponseEndFrame(), + ] + expected_down_frames = [LLMContextFrame, LLMContextAssistantTimestampFrame] + await run_test( + aggregator, + frames_to_send=frames_to_send, + expected_down_frames=expected_down_frames, + ) + self.assertTrue(should_start) + self.assertTrue(should_stop) + self.assertEqual(stop_message.content, "Hello from Pipecat!") + + async def test_interruption(self): + context = LLMContext() + + aggregator = LLMAssistantAggregator(context) + + should_start = 0 + should_stop = 0 + stop_messages = [] + + @aggregator.event_handler("on_assistant_turn_started") + async def on_assistant_turn_started(aggregator): + nonlocal should_start + should_start += 1 + + @aggregator.event_handler("on_assistant_turn_stopped") + async def on_assistant_turn_stopped(aggregator, message: AssistantTurnStoppedMessage): + nonlocal should_stop, stop_messages + should_stop += 1 + stop_messages.append(message) + + frames_to_send = [ + LLMFullResponseStartFrame(), + LLMTextFrame("Hello "), + SleepFrame(), + InterruptionFrame(), + LLMFullResponseStartFrame(), + LLMTextFrame("Hello "), + LLMTextFrame("there!"), + LLMFullResponseEndFrame(), + ] + expected_down_frames = [ + LLMContextFrame, + LLMContextAssistantTimestampFrame, + InterruptionFrame, + LLMContextFrame, + LLMContextAssistantTimestampFrame, + ] + await run_test( + aggregator, + frames_to_send=frames_to_send, + expected_down_frames=expected_down_frames, + ) + self.assertEqual(should_start, 2) + self.assertEqual(should_stop, 2) + self.assertEqual(stop_messages[0].content, "Hello") + self.assertEqual(stop_messages[1].content, "Hello there!") + + async def test_thought(self): + context = LLMContext() + + aggregator = LLMAssistantAggregator(context) + + thought_message = None + + @aggregator.event_handler("on_assistant_thought") + async def on_assistant_thought(aggregator, message: AssistantThoughtMessage): + nonlocal thought_message + thought_message = message + + frames_to_send = [ + LLMFullResponseStartFrame(), + LLMThoughtStartFrame(), + LLMThoughtTextFrame(text="I'm thinking!"), + LLMThoughtEndFrame(), + LLMFullResponseEndFrame(), + ] + await run_test(aggregator, frames_to_send=frames_to_send) + self.assertEqual(thought_message.content, "I'm thinking!")