improved unit tests / add a run_test function to test processors
This commit is contained in:
Aleix Conchillo Flaqué
370
tests/processors/aggregators/test_llm_response.py
Normal file
370
tests/processors/aggregators/test_llm_response.py
Normal file
@@ -0,0 +1,370 @@
|
||||
#
|
||||
# Copyright (c) 2024, Daily
|
||||
#
|
||||
# SPDX-License-Identifier: BSD 2-Clause License
|
||||
#
|
||||
|
||||
import unittest
|
||||
|
||||
from pipecat.frames.frames import (
|
||||
BotInterruptionFrame,
|
||||
InterimTranscriptionFrame,
|
||||
LLMFullResponseEndFrame,
|
||||
LLMFullResponseStartFrame,
|
||||
StartInterruptionFrame,
|
||||
StopInterruptionFrame,
|
||||
TextFrame,
|
||||
TranscriptionFrame,
|
||||
UserStartedSpeakingFrame,
|
||||
UserStoppedSpeakingFrame,
|
||||
)
|
||||
from pipecat.processors.aggregators.llm_response import (
|
||||
LLMAssistantContextAggregator,
|
||||
LLMFullResponseAggregator,
|
||||
LLMUserContextAggregator,
|
||||
)
|
||||
from pipecat.processors.aggregators.openai_llm_context import (
|
||||
OpenAILLMContext,
|
||||
OpenAILLMContextFrame,
|
||||
)
|
||||
from tests.utils import run_test
|
||||
|
||||
|
||||
class TestLLMUserContextAggregator(unittest.IsolatedAsyncioTestCase):
|
||||
# S E ->
|
||||
async def test_s_e(self):
|
||||
"""S E case"""
|
||||
context_aggregator = LLMUserContextAggregator(
|
||||
OpenAILLMContext(messages=[{"role": "", "content": ""}])
|
||||
)
|
||||
frames_to_send = [
|
||||
StartInterruptionFrame(),
|
||||
UserStartedSpeakingFrame(),
|
||||
StopInterruptionFrame(),
|
||||
UserStoppedSpeakingFrame(),
|
||||
]
|
||||
expected_returned_frames = [
|
||||
StartInterruptionFrame,
|
||||
UserStartedSpeakingFrame,
|
||||
StopInterruptionFrame,
|
||||
UserStoppedSpeakingFrame,
|
||||
]
|
||||
await run_test(context_aggregator, frames_to_send, expected_returned_frames)
|
||||
|
||||
# S T E -> T
|
||||
async def test_s_t_e(self):
|
||||
"""S T E case"""
|
||||
context_aggregator = LLMUserContextAggregator(
|
||||
OpenAILLMContext(messages=[{"role": "", "content": ""}])
|
||||
)
|
||||
frames_to_send = [
|
||||
StartInterruptionFrame(),
|
||||
UserStartedSpeakingFrame(),
|
||||
TranscriptionFrame("Hello", "", ""),
|
||||
StopInterruptionFrame(),
|
||||
UserStoppedSpeakingFrame(),
|
||||
]
|
||||
expected_returned_frames = [
|
||||
StartInterruptionFrame,
|
||||
UserStartedSpeakingFrame,
|
||||
StopInterruptionFrame,
|
||||
UserStoppedSpeakingFrame,
|
||||
OpenAILLMContextFrame,
|
||||
]
|
||||
await run_test(context_aggregator, frames_to_send, expected_returned_frames)
|
||||
|
||||
# S I T E -> T
|
||||
async def test_s_i_t_e(self):
|
||||
"""S I T E case"""
|
||||
context_aggregator = LLMUserContextAggregator(
|
||||
OpenAILLMContext(messages=[{"role": "", "content": ""}])
|
||||
)
|
||||
frames_to_send = [
|
||||
StartInterruptionFrame(),
|
||||
UserStartedSpeakingFrame(),
|
||||
InterimTranscriptionFrame("This", "", ""),
|
||||
TranscriptionFrame("This is a test", "", ""),
|
||||
StopInterruptionFrame(),
|
||||
UserStoppedSpeakingFrame(),
|
||||
]
|
||||
expected_returned_frames = [
|
||||
StartInterruptionFrame,
|
||||
UserStartedSpeakingFrame,
|
||||
StopInterruptionFrame,
|
||||
UserStoppedSpeakingFrame,
|
||||
OpenAILLMContextFrame,
|
||||
]
|
||||
await run_test(context_aggregator, frames_to_send, expected_returned_frames)
|
||||
|
||||
# S I E T -> T
|
||||
async def test_s_i_e_t(self):
|
||||
"""S I E T case"""
|
||||
context_aggregator = LLMUserContextAggregator(
|
||||
OpenAILLMContext(messages=[{"role": "", "content": ""}])
|
||||
)
|
||||
frames_to_send = [
|
||||
StartInterruptionFrame(),
|
||||
UserStartedSpeakingFrame(),
|
||||
InterimTranscriptionFrame("This", "", ""),
|
||||
StopInterruptionFrame(),
|
||||
UserStoppedSpeakingFrame(),
|
||||
TranscriptionFrame("This is a test", "", ""),
|
||||
]
|
||||
expected_returned_frames = [
|
||||
StartInterruptionFrame,
|
||||
UserStartedSpeakingFrame,
|
||||
StopInterruptionFrame,
|
||||
UserStoppedSpeakingFrame,
|
||||
OpenAILLMContextFrame,
|
||||
]
|
||||
await run_test(context_aggregator, frames_to_send, expected_returned_frames)
|
||||
|
||||
# S I E I T -> T
|
||||
async def test_s_i_e_i_t(self):
|
||||
"""S I E I T case"""
|
||||
context_aggregator = LLMUserContextAggregator(
|
||||
OpenAILLMContext(messages=[{"role": "", "content": ""}])
|
||||
)
|
||||
frames_to_send = [
|
||||
StartInterruptionFrame(),
|
||||
UserStartedSpeakingFrame(),
|
||||
InterimTranscriptionFrame("This", "", ""),
|
||||
StopInterruptionFrame(),
|
||||
UserStoppedSpeakingFrame(),
|
||||
InterimTranscriptionFrame("This is", "", ""),
|
||||
TranscriptionFrame("This is a test", "", ""),
|
||||
]
|
||||
expected_returned_frames = [
|
||||
StartInterruptionFrame,
|
||||
UserStartedSpeakingFrame,
|
||||
StopInterruptionFrame,
|
||||
UserStoppedSpeakingFrame,
|
||||
OpenAILLMContextFrame,
|
||||
]
|
||||
await run_test(context_aggregator, frames_to_send, expected_returned_frames)
|
||||
|
||||
# S E T -> T
|
||||
async def test_s_e_t(self):
|
||||
"""S E case"""
|
||||
context_aggregator = LLMUserContextAggregator(
|
||||
OpenAILLMContext(messages=[{"role": "", "content": ""}])
|
||||
)
|
||||
frames_to_send = [
|
||||
StartInterruptionFrame(),
|
||||
UserStartedSpeakingFrame(),
|
||||
StopInterruptionFrame(),
|
||||
UserStoppedSpeakingFrame(),
|
||||
TranscriptionFrame("This is a test", "", ""),
|
||||
]
|
||||
expected_returned_frames = [
|
||||
StartInterruptionFrame,
|
||||
UserStartedSpeakingFrame,
|
||||
StopInterruptionFrame,
|
||||
UserStoppedSpeakingFrame,
|
||||
OpenAILLMContextFrame,
|
||||
]
|
||||
await run_test(context_aggregator, frames_to_send, expected_returned_frames)
|
||||
|
||||
# S E I T -> T
|
||||
async def test_s_e_i_t(self):
|
||||
"""S E I T case"""
|
||||
context_aggregator = LLMUserContextAggregator(
|
||||
OpenAILLMContext(messages=[{"role": "", "content": ""}])
|
||||
)
|
||||
frames_to_send = [
|
||||
StartInterruptionFrame(),
|
||||
UserStartedSpeakingFrame(),
|
||||
StopInterruptionFrame(),
|
||||
UserStoppedSpeakingFrame(),
|
||||
InterimTranscriptionFrame("This", "", ""),
|
||||
TranscriptionFrame("This is a test", "", ""),
|
||||
]
|
||||
expected_returned_frames = [
|
||||
StartInterruptionFrame,
|
||||
UserStartedSpeakingFrame,
|
||||
StopInterruptionFrame,
|
||||
UserStoppedSpeakingFrame,
|
||||
OpenAILLMContextFrame,
|
||||
]
|
||||
await run_test(context_aggregator, frames_to_send, expected_returned_frames)
|
||||
|
||||
# S T1 I E S T2 E -> "T1 T2"
|
||||
async def test_s_t1_i_e_s_t2_e(self):
|
||||
"""S T1 I E S T2 E case"""
|
||||
context_aggregator = LLMUserContextAggregator(
|
||||
OpenAILLMContext(messages=[{"role": "", "content": ""}])
|
||||
)
|
||||
frames_to_send = [
|
||||
StartInterruptionFrame(),
|
||||
UserStartedSpeakingFrame(),
|
||||
TranscriptionFrame("T1", "", ""),
|
||||
InterimTranscriptionFrame("", "", ""),
|
||||
StopInterruptionFrame(),
|
||||
UserStoppedSpeakingFrame(),
|
||||
StartInterruptionFrame(),
|
||||
UserStartedSpeakingFrame(),
|
||||
TranscriptionFrame("T2", "", ""),
|
||||
StopInterruptionFrame(),
|
||||
UserStoppedSpeakingFrame(),
|
||||
]
|
||||
expected_returned_frames = [
|
||||
StartInterruptionFrame,
|
||||
UserStartedSpeakingFrame,
|
||||
StopInterruptionFrame,
|
||||
UserStoppedSpeakingFrame,
|
||||
StartInterruptionFrame,
|
||||
UserStartedSpeakingFrame,
|
||||
StopInterruptionFrame,
|
||||
UserStoppedSpeakingFrame,
|
||||
OpenAILLMContextFrame,
|
||||
]
|
||||
(received_down, _) = await run_test(
|
||||
context_aggregator, frames_to_send, expected_returned_frames
|
||||
)
|
||||
assert received_down[-1].context.messages[-1]["content"] == "T1 T2"
|
||||
|
||||
# S I E T1 I T2 -> T1 Interruption T2
|
||||
async def test_s_i_e_t1_i_t2(self):
|
||||
"""S I E T1 I T2 case"""
|
||||
context_aggregator = LLMUserContextAggregator(
|
||||
OpenAILLMContext(messages=[{"role": "", "content": ""}])
|
||||
)
|
||||
frames_to_send = [
|
||||
StartInterruptionFrame(),
|
||||
UserStartedSpeakingFrame(),
|
||||
InterimTranscriptionFrame("", "", ""),
|
||||
StopInterruptionFrame(),
|
||||
UserStoppedSpeakingFrame(),
|
||||
TranscriptionFrame("T1", "", ""),
|
||||
InterimTranscriptionFrame("", "", ""),
|
||||
TranscriptionFrame("T2", "", ""),
|
||||
]
|
||||
expected_down_frames = [
|
||||
StartInterruptionFrame,
|
||||
UserStartedSpeakingFrame,
|
||||
StopInterruptionFrame,
|
||||
UserStoppedSpeakingFrame,
|
||||
OpenAILLMContextFrame,
|
||||
OpenAILLMContextFrame,
|
||||
]
|
||||
expected_up_frames = [
|
||||
BotInterruptionFrame,
|
||||
]
|
||||
(received_down, _) = await run_test(
|
||||
context_aggregator, frames_to_send, expected_down_frames, expected_up_frames
|
||||
)
|
||||
assert received_down[-1].context.messages[-2]["content"] == "T1"
|
||||
assert received_down[-1].context.messages[-1]["content"] == "T2"
|
||||
|
||||
# S T1 E T2 -> T1 Interruption T2
|
||||
async def test_s_t1_e_t2(self):
|
||||
"""S T1 E T2 case"""
|
||||
context_aggregator = LLMUserContextAggregator(
|
||||
OpenAILLMContext(messages=[{"role": "", "content": ""}])
|
||||
)
|
||||
frames_to_send = [
|
||||
StartInterruptionFrame(),
|
||||
UserStartedSpeakingFrame(),
|
||||
TranscriptionFrame("T1", "", ""),
|
||||
StopInterruptionFrame(),
|
||||
UserStoppedSpeakingFrame(),
|
||||
TranscriptionFrame("T2", "", ""),
|
||||
]
|
||||
expected_down_frames = [
|
||||
StartInterruptionFrame,
|
||||
UserStartedSpeakingFrame,
|
||||
StopInterruptionFrame,
|
||||
UserStoppedSpeakingFrame,
|
||||
OpenAILLMContextFrame,
|
||||
OpenAILLMContextFrame,
|
||||
]
|
||||
expected_up_frames = [
|
||||
BotInterruptionFrame,
|
||||
]
|
||||
(received_down, _) = await run_test(
|
||||
context_aggregator, frames_to_send, expected_down_frames, expected_up_frames
|
||||
)
|
||||
assert received_down[-1].context.messages[-2]["content"] == "T1"
|
||||
assert received_down[-1].context.messages[-1]["content"] == "T2"
|
||||
|
||||
# S E T1 T2 -> T1 Interruption T2
|
||||
async def test_s_e_t1_t2(self):
|
||||
"""S E T1 T2 case"""
|
||||
context_aggregator = LLMUserContextAggregator(
|
||||
OpenAILLMContext(messages=[{"role": "", "content": ""}])
|
||||
)
|
||||
frames_to_send = [
|
||||
StartInterruptionFrame(),
|
||||
UserStartedSpeakingFrame(),
|
||||
StopInterruptionFrame(),
|
||||
UserStoppedSpeakingFrame(),
|
||||
TranscriptionFrame("T1", "", ""),
|
||||
TranscriptionFrame("T2", "", ""),
|
||||
]
|
||||
expected_down_frames = [
|
||||
StartInterruptionFrame,
|
||||
UserStartedSpeakingFrame,
|
||||
StopInterruptionFrame,
|
||||
UserStoppedSpeakingFrame,
|
||||
OpenAILLMContextFrame,
|
||||
OpenAILLMContextFrame,
|
||||
]
|
||||
expected_up_frames = [
|
||||
BotInterruptionFrame,
|
||||
]
|
||||
(received_down, _) = await run_test(
|
||||
context_aggregator, frames_to_send, expected_down_frames, expected_up_frames
|
||||
)
|
||||
assert received_down[-1].context.messages[-2]["content"] == "T1"
|
||||
assert received_down[-1].context.messages[-1]["content"] == "T2"
|
||||
|
||||
|
||||
class TestLLMAssistantContextAggregator(unittest.IsolatedAsyncioTestCase):
|
||||
# S T E -> T
|
||||
async def test_s_t_e(self):
|
||||
"""S T E case"""
|
||||
context_aggregator = LLMAssistantContextAggregator(
|
||||
OpenAILLMContext(messages=[{"role": "", "content": ""}])
|
||||
)
|
||||
frames_to_send = [
|
||||
LLMFullResponseStartFrame(),
|
||||
TextFrame("Hello this is Pipecat speaking!"),
|
||||
TextFrame("How are you?"),
|
||||
LLMFullResponseEndFrame(),
|
||||
]
|
||||
expected_returned_frames = [
|
||||
LLMFullResponseStartFrame,
|
||||
OpenAILLMContextFrame,
|
||||
LLMFullResponseEndFrame,
|
||||
]
|
||||
(received_down, _) = await run_test(
|
||||
context_aggregator, frames_to_send, expected_returned_frames
|
||||
)
|
||||
assert (
|
||||
received_down[-2].context.messages[-1]["content"]
|
||||
== "Hello this is Pipecat speaking! How are you?"
|
||||
)
|
||||
|
||||
|
||||
class TestLLMFullResponseAggregator(unittest.IsolatedAsyncioTestCase):
|
||||
# S T E -> T
|
||||
async def test_s_t_e(self):
|
||||
"""S T E case"""
|
||||
response_aggregator = LLMFullResponseAggregator()
|
||||
frames_to_send = [
|
||||
LLMFullResponseStartFrame(),
|
||||
TextFrame("Hello "),
|
||||
TextFrame("this "),
|
||||
TextFrame("is "),
|
||||
TextFrame("Pipecat!"),
|
||||
LLMFullResponseEndFrame(),
|
||||
]
|
||||
expected_returned_frames = [
|
||||
LLMFullResponseStartFrame,
|
||||
TextFrame,
|
||||
LLMFullResponseEndFrame,
|
||||
]
|
||||
(received_down, _) = await run_test(
|
||||
response_aggregator, frames_to_send, expected_returned_frames
|
||||
)
|
||||
assert received_down[-2].text == "Hello this is Pipecat!"
|
||||
Reference in New Issue
Block a user