import os
import json
import time
import datetime
import yaml
from swarms.structs.conversation import (
Conversation,
generate_conversation_id,
)
def run_all_tests():
"""Run all tests for the Conversation class"""
test_results = []
def run_test(test_func):
try:
test_func()
test_results.append(f"✅ {test_func.__name__} passed")
except Exception as e:
test_results.append(
f"❌ {test_func.__name__} failed: {str(e)}"
)
def test_basic_initialization():
"""Test basic initialization of Conversation"""
conv = Conversation()
assert conv.id is not None
assert conv.conversation_history is not None
assert isinstance(conv.conversation_history, list)
# Test with custom ID
custom_id = generate_conversation_id()
conv_with_id = Conversation(id=custom_id)
assert conv_with_id.id == custom_id
# Test with custom name
conv_with_name = Conversation(name="Test Conversation")
assert conv_with_name.name == "Test Conversation"
def test_initialization_with_settings():
"""Test initialization with various settings"""
conv = Conversation(
system_prompt="Test system prompt",
time_enabled=True,
autosave=True,
token_count=True,
provider="in-memory",
context_length=4096,
rules="Test rules",
custom_rules_prompt="Custom rules",
user="TestUser:",
save_as_yaml=True,
save_as_json_bool=True,
)
# Test all settings
assert conv.system_prompt == "Test system prompt"
assert conv.time_enabled is True
assert conv.autosave is True
assert conv.token_count is True
assert conv.provider == "in-memory"
assert conv.context_length == 4096
assert conv.rules == "Test rules"
assert conv.custom_rules_prompt == "Custom rules"
assert conv.user == "TestUser:"
assert conv.save_as_yaml is True
assert conv.save_as_json_bool is True
def test_message_manipulation():
"""Test adding, deleting, and updating messages"""
conv = Conversation()
# Test adding messages with different content types
conv.add("user", "Hello") # String content
conv.add("assistant", {"response": "Hi"}) # Dict content
conv.add("system", ["Hello", "Hi"]) # List content
assert len(conv.conversation_history) == 3
assert isinstance(
conv.conversation_history[1]["content"], dict
)
assert isinstance(
conv.conversation_history[2]["content"], list
)
# Test adding multiple messages
conv.add_multiple(
["user", "assistant", "system"],
["Hi", "Hello there", "System message"],
)
assert len(conv.conversation_history) == 6
# Test updating message with different content type
conv.update(0, "user", {"updated": "content"})
assert isinstance(
conv.conversation_history[0]["content"], dict
)
# Test deleting multiple messages
conv.delete(0)
conv.delete(0)
assert len(conv.conversation_history) == 4
def test_message_retrieval():
"""Test message retrieval methods"""
conv = Conversation()
# Add messages in specific order for testing
conv.add("user", "Test message")
conv.add("assistant", "Test response")
conv.add("system", "System message")
# Test query - note: messages might have system prompt prepended
message = conv.query(0)
assert "Test message" in message["content"]
# Test search with multiple results
results = conv.search("Test")
assert (
len(results) >= 2
) # At least two messages should contain "Test"
assert any(
"Test message" in str(msg["content"]) for msg in results
)
assert any(
"Test response" in str(msg["content"]) for msg in results
)
# Test get_last_message_as_string
last_message = conv.get_last_message_as_string()
assert "System message" in last_message
# Test return_messages_as_list
messages_list = conv.return_messages_as_list()
assert (
len(messages_list) >= 3
) # At least our 3 added messages
assert any("Test message" in msg for msg in messages_list)
# Test return_messages_as_dictionary
messages_dict = conv.return_messages_as_dictionary()
assert (
len(messages_dict) >= 3
) # At least our 3 added messages
assert all(isinstance(m, dict) for m in messages_dict)
assert all(
{"role", "content"} <= set(m.keys())
for m in messages_dict
)
# Test get_final_message and content
assert "System message" in conv.get_final_message()
assert "System message" in conv.get_final_message_content()
# Test return_all_except_first
remaining_messages = conv.return_all_except_first()
assert (
len(remaining_messages) >= 2
) # At least 2 messages after removing first
# Test return_all_except_first_string
remaining_string = conv.return_all_except_first_string()
assert isinstance(remaining_string, str)
def test_saving_loading():
"""Test saving and loading conversation"""
# Test with save_enabled
conv = Conversation(
save_enabled=True,
conversations_dir="./test_conversations",
)
conv.add("user", "Test save message")
# Test save_as_json
test_file = os.path.join(
"./test_conversations", "test_conversation.json"
)
conv.save_as_json(test_file)
assert os.path.exists(test_file)
# Test load_from_json
new_conv = Conversation()
new_conv.load_from_json(test_file)
assert len(new_conv.conversation_history) == 1
assert (
new_conv.conversation_history[0]["content"]
== "Test save message"
)
# Test class method load_conversation
loaded_conv = Conversation.load_conversation(
name=conv.id, conversations_dir="./test_conversations"
)
assert loaded_conv.id == conv.id
# Cleanup
os.remove(test_file)
os.rmdir("./test_conversations")
def test_output_formats():
"""Test different output formats"""
conv = Conversation()
conv.add("user", "Test message")
conv.add("assistant", {"response": "Test"})
# Test JSON output
json_output = conv.to_json()
assert isinstance(json_output, str)
parsed_json = json.loads(json_output)
assert len(parsed_json) == 2
# Test dict output
dict_output = conv.to_dict()
assert isinstance(dict_output, list)
assert len(dict_output) == 2
# Test YAML output
yaml_output = conv.to_yaml()
assert isinstance(yaml_output, str)
parsed_yaml = yaml.safe_load(yaml_output)
assert len(parsed_yaml) == 2
# Test return_json
json_str = conv.return_json()
assert isinstance(json_str, str)
assert len(json.loads(json_str)) == 2
def test_memory_management():
"""Test memory management functions"""
conv = Conversation()
# Test clear
conv.add("user", "Test message")
conv.clear()
assert len(conv.conversation_history) == 0
# Test clear_memory
conv.add("user", "Test message")
conv.clear_memory()
assert len(conv.conversation_history) == 0
# Test batch operations
messages = [
{"role": "user", "content": "Message 1"},
{"role": "assistant", "content": "Response 1"},
]
conv.batch_add(messages)
assert len(conv.conversation_history) == 2
# Test truncate_memory_with_tokenizer
if conv.tokenizer: # Only if tokenizer is available
conv.truncate_memory_with_tokenizer()
assert len(conv.conversation_history) > 0
def test_conversation_metadata():
"""Test conversation metadata and listing"""
test_dir = "./test_conversations_metadata"
os.makedirs(test_dir, exist_ok=True)
try:
# Create a conversation with metadata
conv = Conversation(
name="Test Conv",
system_prompt="System",
rules="Rules",
custom_rules_prompt="Custom",
conversations_dir=test_dir,
save_enabled=True,
autosave=True,
)
# Add a message to trigger save
conv.add("user", "Test message")
# Give a small delay for autosave
time.sleep(0.1)
# List conversations and verify
conversations = Conversation.list_conversations(test_dir)
assert len(conversations) >= 1
found_conv = next(
(
c
for c in conversations
if c["name"] == "Test Conv"
),
None,
)
assert found_conv is not None
assert found_conv["id"] == conv.id
finally:
# Cleanup
import shutil
if os.path.exists(test_dir):
shutil.rmtree(test_dir)
def test_time_enabled_messages():
"""Test time-enabled messages"""
conv = Conversation(time_enabled=True)
conv.add("user", "Time test")
# Verify timestamp in message
message = conv.conversation_history[0]
assert "timestamp" in message
assert isinstance(message["timestamp"], str)
# Verify time in content when time_enabled is True
assert "Time:" in message["content"]
def test_provider_specific():
"""Test provider-specific functionality"""
# Test in-memory provider
conv_memory = Conversation(provider="in-memory")
conv_memory.add("user", "Test")
assert len(conv_memory.conversation_history) == 1
# Test mem0 provider if available
try:
conv_mem0 = Conversation(provider="mem0")
conv_mem0.add("user", "Test")
# Add appropriate assertions based on mem0 behavior
except:
pass # Skip if mem0 is not available
def test_tool_output():
"""Test tool output handling"""
conv = Conversation()
tool_output = {
"tool_name": "test_tool",
"output": "test result",
}
conv.add_tool_output_to_agent("tool", tool_output)
assert len(conv.conversation_history) == 1
assert conv.conversation_history[0]["role"] == "tool"
assert conv.conversation_history[0]["content"] == tool_output
def test_autosave_functionality():
"""Test autosave functionality and related features"""
test_dir = "./test_conversations_autosave"
os.makedirs(test_dir, exist_ok=True)
try:
# Test with autosave and save_enabled True
conv = Conversation(
autosave=True,
save_enabled=True,
conversations_dir=test_dir,
name="autosave_test",
)
# Add a message and verify it was auto-saved
conv.add("user", "Test autosave message")
save_path = os.path.join(test_dir, f"{conv.id}.json")
# Give a small delay for autosave to complete
time.sleep(0.1)
assert os.path.exists(
save_path
), f"Save file not found at {save_path}"
# Load the saved conversation and verify content
loaded_conv = Conversation.load_conversation(
name=conv.id, conversations_dir=test_dir
)
found_message = False
for msg in loaded_conv.conversation_history:
if "Test autosave message" in str(msg["content"]):
found_message = True
break
assert (
found_message
), "Message not found in loaded conversation"
# Clean up first conversation files
if os.path.exists(save_path):
os.remove(save_path)
# Test with save_enabled=False
conv_no_save = Conversation(
autosave=False, # Changed to False to prevent autosave
save_enabled=False,
conversations_dir=test_dir,
)
conv_no_save.add("user", "This shouldn't be saved")
save_path_no_save = os.path.join(
test_dir, f"{conv_no_save.id}.json"
)
time.sleep(0.1) # Give time for potential save
assert not os.path.exists(
save_path_no_save
), "File should not exist when save_enabled is False"
finally:
# Cleanup
import shutil
if os.path.exists(test_dir):
shutil.rmtree(test_dir)
def test_advanced_message_handling():
"""Test advanced message handling features"""
conv = Conversation()
# Test adding messages with metadata
metadata = {"timestamp": "2024-01-01", "session_id": "123"}
conv.add("user", "Test with metadata", metadata=metadata)
# Test batch operations with different content types
messages = [
{"role": "user", "content": "Message 1"},
{
"role": "assistant",
"content": {"response": "Complex response"},
},
{"role": "system", "content": ["Multiple", "Items"]},
]
conv.batch_add(messages)
assert (
len(conv.conversation_history) == 4
) # Including the first message
# Test message format consistency
for msg in conv.conversation_history:
assert "role" in msg
assert "content" in msg
if "timestamp" in msg:
assert isinstance(msg["timestamp"], str)
def test_conversation_metadata_handling():
"""Test handling of conversation metadata and attributes"""
test_dir = "./test_conversations_metadata_handling"
os.makedirs(test_dir, exist_ok=True)
try:
# Test initialization with all optional parameters
conv = Conversation(
name="Test Conv",
system_prompt="System Prompt",
time_enabled=True,
context_length=2048,
rules="Test Rules",
custom_rules_prompt="Custom Rules",
user="CustomUser:",
provider="in-memory",
conversations_dir=test_dir,
save_enabled=True,
)
# Verify all attributes are set correctly
assert conv.name == "Test Conv"
assert conv.system_prompt == "System Prompt"
assert conv.time_enabled is True
assert conv.context_length == 2048
assert conv.rules == "Test Rules"
assert conv.custom_rules_prompt == "Custom Rules"
assert conv.user == "CustomUser:"
assert conv.provider == "in-memory"
# Test saving and loading preserves metadata
conv.save_as_json()
# Load using load_conversation
loaded_conv = Conversation.load_conversation(
name=conv.id, conversations_dir=test_dir
)
# Verify metadata was preserved
assert loaded_conv.name == "Test Conv"
assert loaded_conv.system_prompt == "System Prompt"
assert loaded_conv.rules == "Test Rules"
finally:
# Cleanup
import shutil
shutil.rmtree(test_dir)
def test_time_enabled_features():
"""Test time-enabled message features"""
conv = Conversation(time_enabled=True)
# Add message and verify timestamp
conv.add("user", "Time test message")
message = conv.conversation_history[0]
# Verify timestamp format
assert "timestamp" in message
try:
datetime.datetime.fromisoformat(message["timestamp"])
except ValueError:
assert False, "Invalid timestamp format"
# Verify time in content
assert "Time:" in message["content"]
assert (
datetime.datetime.now().strftime("%Y-%m-%d")
in message["content"]
)
def test_provider_specific_features():
"""Test provider-specific features and behaviors"""
# Test in-memory provider
conv_memory = Conversation(provider="in-memory")
conv_memory.add("user", "Test in-memory")
assert len(conv_memory.conversation_history) == 1
assert (
"Test in-memory"
in conv_memory.get_last_message_as_string()
)
# Test mem0 provider if available
try:
from mem0 import AsyncMemory
# Skip actual mem0 testing since it requires async
pass
except ImportError:
pass
# Test invalid provider
invalid_provider = "invalid_provider"
try:
Conversation(provider=invalid_provider)
# If we get here, the provider was accepted when it shouldn't have been
raise AssertionError(
f"Should have raised ValueError for provider '{invalid_provider}'"
)
except ValueError:
# This is the expected behavior
pass
# Run all tests
tests = [
test_basic_initialization,
test_initialization_with_settings,
test_message_manipulation,
test_message_retrieval,
test_saving_loading,
test_output_formats,
test_memory_management,
test_conversation_metadata,
test_time_enabled_messages,
test_provider_specific,
test_tool_output,
test_autosave_functionality,
test_advanced_message_handling,
test_conversation_metadata_handling,
test_time_enabled_features,
test_provider_specific_features,
]
for test in tests:
run_test(test)
# Print results
print("\nTest Results:")
for result in test_results:
print(result)
if __name__ == "__main__":
run_all_tests()