209 lines
8.7 KiB
Python
209 lines
8.7 KiB
Python
#
|
|
# Copyright (c) 2024-2025 Daily
|
|
#
|
|
# SPDX-License-Identifier: BSD 2-Clause License
|
|
#
|
|
|
|
import unittest
|
|
|
|
from pipecat.processors.aggregators.llm_context import LLMContext, LLMSpecificMessage
|
|
|
|
|
|
class TestGetMessagesForPersistentStorage(unittest.TestCase):
|
|
"""Test suite for LLMContext.get_messages_for_persistent_storage method."""
|
|
|
|
def test_no_system_instruction_returns_messages_as_is(self):
|
|
"""Test that without system instruction, messages are returned unchanged."""
|
|
messages = [
|
|
{"role": "user", "content": "Hello"},
|
|
{"role": "assistant", "content": "Hi there!"},
|
|
]
|
|
context = LLMContext(messages=messages)
|
|
|
|
result = context.get_messages_for_persistent_storage()
|
|
|
|
self.assertEqual(result, messages)
|
|
self.assertEqual(len(result), 2)
|
|
|
|
def test_empty_messages_with_system_instruction_adds_system_message(self):
|
|
"""Test that system instruction is added when messages list is empty."""
|
|
context = LLMContext()
|
|
system_instruction = "You are a helpful assistant."
|
|
|
|
result = context.get_messages_for_persistent_storage(system_instruction)
|
|
|
|
self.assertEqual(len(result), 1)
|
|
self.assertEqual(result[0]["role"], "system")
|
|
self.assertEqual(result[0]["content"], system_instruction)
|
|
|
|
def test_non_system_first_message_prepends_system_instruction(self):
|
|
"""Test that system instruction is prepended when first message is not system."""
|
|
messages = [
|
|
{"role": "user", "content": "Hello"},
|
|
{"role": "assistant", "content": "Hi there!"},
|
|
]
|
|
context = LLMContext(messages=messages)
|
|
system_instruction = "You are a helpful assistant."
|
|
|
|
result = context.get_messages_for_persistent_storage(system_instruction)
|
|
|
|
self.assertEqual(len(result), 3)
|
|
self.assertEqual(result[0]["role"], "system")
|
|
self.assertEqual(result[0]["content"], system_instruction)
|
|
self.assertEqual(result[1], messages[0])
|
|
self.assertEqual(result[2], messages[1])
|
|
|
|
def test_existing_system_message_not_duplicated(self):
|
|
"""Test that system instruction is not added when first message is already system."""
|
|
messages = [
|
|
{"role": "system", "content": "Existing system message"},
|
|
{"role": "user", "content": "Hello"},
|
|
{"role": "assistant", "content": "Hi there!"},
|
|
]
|
|
context = LLMContext(messages=messages)
|
|
system_instruction = "You are a helpful assistant."
|
|
|
|
result = context.get_messages_for_persistent_storage(system_instruction)
|
|
|
|
self.assertEqual(len(result), 3)
|
|
self.assertEqual(result, messages)
|
|
self.assertEqual(result[0]["role"], "system")
|
|
self.assertEqual(result[0]["content"], "Existing system message")
|
|
|
|
def test_empty_system_instruction_does_not_add_message(self):
|
|
"""Test that empty system instruction does not add a system message."""
|
|
messages = [{"role": "user", "content": "Hello"}]
|
|
context = LLMContext(messages=messages)
|
|
|
|
result = context.get_messages_for_persistent_storage("")
|
|
|
|
self.assertEqual(result, messages)
|
|
self.assertEqual(len(result), 1)
|
|
|
|
def test_none_system_instruction_does_not_add_message(self):
|
|
"""Test that None system instruction does not add a system message."""
|
|
messages = [{"role": "user", "content": "Hello"}]
|
|
context = LLMContext(messages=messages)
|
|
|
|
result = context.get_messages_for_persistent_storage(None)
|
|
|
|
self.assertEqual(result, messages)
|
|
self.assertEqual(len(result), 1)
|
|
|
|
def test_whitespace_only_system_instruction_adds_message(self):
|
|
"""Test that whitespace-only system instruction still adds a system message."""
|
|
messages = [{"role": "user", "content": "Hello"}]
|
|
context = LLMContext(messages=messages)
|
|
system_instruction = " "
|
|
|
|
result = context.get_messages_for_persistent_storage(system_instruction)
|
|
|
|
self.assertEqual(len(result), 2)
|
|
self.assertEqual(result[0]["role"], "system")
|
|
self.assertEqual(result[0]["content"], system_instruction)
|
|
|
|
def test_with_llm_specific_messages(self):
|
|
"""Test that method works correctly with LLMSpecificMessage objects."""
|
|
llm_specific = LLMSpecificMessage(
|
|
llm="test-llm", message={"role": "user", "content": "Specific"}
|
|
)
|
|
messages = [{"role": "user", "content": "Standard message"}, llm_specific]
|
|
context = LLMContext(messages=messages)
|
|
system_instruction = "You are a helpful assistant."
|
|
|
|
result = context.get_messages_for_persistent_storage(system_instruction)
|
|
|
|
self.assertEqual(len(result), 3)
|
|
self.assertEqual(result[0]["role"], "system")
|
|
self.assertEqual(result[0]["content"], system_instruction)
|
|
self.assertEqual(result[1], messages[0])
|
|
self.assertEqual(result[2], llm_specific)
|
|
|
|
def test_system_message_detection_case_sensitivity(self):
|
|
"""Test that system message detection is case sensitive."""
|
|
messages = [
|
|
{"role": "System", "content": "Mixed case system"}, # Capital S
|
|
{"role": "user", "content": "Hello"},
|
|
]
|
|
context = LLMContext(messages=messages)
|
|
system_instruction = "You are a helpful assistant."
|
|
|
|
result = context.get_messages_for_persistent_storage(system_instruction)
|
|
|
|
# Should prepend because "System" != "system"
|
|
self.assertEqual(len(result), 3)
|
|
self.assertEqual(result[0]["role"], "system")
|
|
self.assertEqual(result[0]["content"], system_instruction)
|
|
self.assertEqual(result[1], messages[0])
|
|
|
|
def test_message_without_role_key_does_not_crash(self):
|
|
"""Test that messages without 'role' key are handled gracefully."""
|
|
messages = [{"content": "Message without role"}, {"role": "user", "content": "Hello"}]
|
|
context = LLMContext(messages=messages)
|
|
system_instruction = "You are a helpful assistant."
|
|
|
|
result = context.get_messages_for_persistent_storage(system_instruction)
|
|
|
|
# Should prepend system instruction since first message doesn't have role="system"
|
|
self.assertEqual(len(result), 3)
|
|
self.assertEqual(result[0]["role"], "system")
|
|
self.assertEqual(result[0]["content"], system_instruction)
|
|
|
|
def test_original_messages_not_modified(self):
|
|
"""Test that the original messages list is not modified."""
|
|
original_messages = [{"role": "user", "content": "Hello"}]
|
|
context = LLMContext(messages=original_messages)
|
|
system_instruction = "You are a helpful assistant."
|
|
|
|
result = context.get_messages_for_persistent_storage(system_instruction)
|
|
|
|
# Original messages should remain unchanged
|
|
self.assertEqual(len(original_messages), 1)
|
|
self.assertEqual(original_messages[0]["role"], "user")
|
|
|
|
# Result should have system message prepended
|
|
self.assertEqual(len(result), 2)
|
|
self.assertEqual(result[0]["role"], "system")
|
|
self.assertEqual(result[1], original_messages[0])
|
|
|
|
def test_complex_message_structure_preserved(self):
|
|
"""Test that complex message structures are preserved."""
|
|
complex_message = {
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "text", "text": "Complex message"},
|
|
{"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,..."}},
|
|
],
|
|
}
|
|
messages = [complex_message]
|
|
context = LLMContext(messages=messages)
|
|
system_instruction = "You are a helpful assistant."
|
|
|
|
result = context.get_messages_for_persistent_storage(system_instruction)
|
|
|
|
self.assertEqual(len(result), 2)
|
|
self.assertEqual(result[0]["role"], "system")
|
|
self.assertEqual(result[1], complex_message)
|
|
self.assertEqual(result[1]["content"], complex_message["content"])
|
|
|
|
def test_deep_copy_prevents_nested_mutation(self):
|
|
"""Test that deep copy prevents mutation of nested message content."""
|
|
nested_content = {"nested": {"data": "original"}}
|
|
complex_message = {"role": "user", "content": nested_content}
|
|
messages = [complex_message]
|
|
context = LLMContext(messages=messages)
|
|
system_instruction = "You are a helpful assistant."
|
|
|
|
result = context.get_messages_for_persistent_storage(system_instruction)
|
|
|
|
# Modify the nested content in the result
|
|
result[1]["content"]["nested"]["data"] = "modified"
|
|
|
|
# Original message should remain unchanged
|
|
self.assertEqual(complex_message["content"]["nested"]["data"], "original")
|
|
self.assertEqual(context.get_messages()[0]["content"]["nested"]["data"], "original")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|