You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
354 lines
12 KiB
354 lines
12 KiB
import pytest
|
|
|
|
from swarms.structs.agent import Agent
|
|
from swarms.structs.multi_agent_router import MultiAgentRouter
|
|
|
|
|
|
# Test fixtures
|
|
def real_agents():
|
|
"""Create real agents for testing"""
|
|
return [
|
|
Agent(
|
|
agent_name="ResearchAgent",
|
|
agent_description="Specializes in researching topics and providing detailed, factual information",
|
|
system_prompt="You are a research specialist. Provide detailed, well-researched information about any topic, citing sources when possible.",
|
|
model_name="gpt-4o-mini",
|
|
max_loops=1,
|
|
verbose=False,
|
|
print_on=False,
|
|
),
|
|
Agent(
|
|
agent_name="CodeExpertAgent",
|
|
agent_description="Expert in writing, reviewing, and explaining code across multiple programming languages",
|
|
system_prompt="You are a coding expert. Write, review, and explain code with a focus on best practices and clean code principles.",
|
|
model_name="gpt-4o-mini",
|
|
max_loops=1,
|
|
verbose=False,
|
|
print_on=False,
|
|
),
|
|
Agent(
|
|
agent_name="WritingAgent",
|
|
agent_description="Skilled in creative and technical writing, content creation, and editing",
|
|
system_prompt="You are a writing specialist. Create, edit, and improve written content while maintaining appropriate tone and style.",
|
|
model_name="gpt-4o-mini",
|
|
max_loops=1,
|
|
verbose=False,
|
|
print_on=False,
|
|
),
|
|
Agent(
|
|
agent_name="MathAgent",
|
|
agent_description="Expert in mathematical calculations and problem solving",
|
|
system_prompt="You are a math expert. Solve mathematical problems and explain solutions clearly.",
|
|
model_name="gpt-4o-mini",
|
|
max_loops=1,
|
|
verbose=False,
|
|
print_on=False,
|
|
),
|
|
]
|
|
|
|
|
|
# ============================================================================
|
|
# INITIALIZATION TESTS
|
|
# ============================================================================
|
|
|
|
|
|
def test_multi_agent_router_initialization_default():
|
|
"""Test MultiAgentRouter initialization with default parameters"""
|
|
router = MultiAgentRouter(agents=real_agents())
|
|
|
|
assert router.name == "swarm-router"
|
|
assert (
|
|
router.description
|
|
== "Routes tasks to specialized agents based on their capabilities"
|
|
)
|
|
assert router.model == "gpt-4o-mini"
|
|
assert router.temperature == 0.1
|
|
assert router.output_type == "dict"
|
|
assert router.print_on is True
|
|
assert router.skip_null_tasks is True
|
|
assert len(router.agents) == 4
|
|
assert all(
|
|
agent_name in router.agents
|
|
for agent_name in [
|
|
"ResearchAgent",
|
|
"CodeExpertAgent",
|
|
"WritingAgent",
|
|
"MathAgent",
|
|
]
|
|
)
|
|
assert isinstance(
|
|
router.conversation, object
|
|
) # Conversation object
|
|
assert hasattr(router.function_caller, "run")
|
|
|
|
|
|
def test_multi_agent_router_initialization_custom_params():
|
|
"""Test MultiAgentRouter initialization with custom parameters"""
|
|
custom_name = "custom-router"
|
|
custom_description = "Custom description"
|
|
custom_model = "gpt-4"
|
|
custom_temperature = 0.5
|
|
custom_output_type = "json"
|
|
|
|
router = MultiAgentRouter(
|
|
name=custom_name,
|
|
description=custom_description,
|
|
agents=real_agents(),
|
|
model=custom_model,
|
|
temperature=custom_temperature,
|
|
output_type=custom_output_type,
|
|
print_on=False,
|
|
skip_null_tasks=False,
|
|
system_prompt="Custom system prompt",
|
|
)
|
|
|
|
assert router.name == custom_name
|
|
assert router.description == custom_description
|
|
assert router.model == custom_model
|
|
assert router.temperature == custom_temperature
|
|
assert router.output_type == custom_output_type
|
|
assert router.print_on is False
|
|
assert router.skip_null_tasks is False
|
|
assert router.system_prompt == "Custom system prompt"
|
|
|
|
|
|
def test_multi_agent_router_repr():
|
|
"""Test MultiAgentRouter string representation"""
|
|
router = MultiAgentRouter(agents=real_agents())
|
|
|
|
expected_repr = f"MultiAgentRouter(name={router.name}, agents={list(router.agents.keys())})"
|
|
assert repr(router) == expected_repr
|
|
|
|
|
|
# ============================================================================
|
|
# SINGLE HANDOFF TESTS
|
|
# ============================================================================
|
|
|
|
|
|
def test_handle_single_handoff_valid():
|
|
"""Test handling single handoff with valid agent"""
|
|
router = MultiAgentRouter(agents=real_agents())
|
|
|
|
result = router.route_task("Write a fibonacci function")
|
|
|
|
# Check that conversation was updated
|
|
assert len(router.conversation.conversation_history) > 0
|
|
# Check that we got a valid response
|
|
assert result is not None
|
|
assert isinstance(result, (list, dict))
|
|
|
|
|
|
# ============================================================================
|
|
# MULTIPLE HANDOFF TESTS
|
|
# ============================================================================
|
|
|
|
|
|
def test_handle_multiple_handoffs_valid():
|
|
"""Test handling multiple handoffs with valid agents"""
|
|
router = MultiAgentRouter(agents=real_agents())
|
|
|
|
result = router.route_task("Research and implement fibonacci")
|
|
|
|
# Check that conversation was updated
|
|
history = router.conversation.conversation_history
|
|
|
|
assert len(history) > 0
|
|
assert result is not None
|
|
assert isinstance(result, (list, dict))
|
|
|
|
|
|
def test_handle_multiple_handoffs_with_null_tasks():
|
|
"""Test handling multiple handoffs with some null tasks"""
|
|
router = MultiAgentRouter(
|
|
agents=real_agents(), skip_null_tasks=True
|
|
)
|
|
|
|
result = router.route_task("Mixed task")
|
|
|
|
# Should still return a valid result
|
|
history = router.conversation.conversation_history
|
|
assert len(history) > 0
|
|
assert result is not None
|
|
assert isinstance(result, (list, dict))
|
|
|
|
|
|
# ============================================================================
|
|
# ROUTE TASK TESTS
|
|
# ============================================================================
|
|
|
|
|
|
def test_route_task_single_agent():
|
|
"""Test route_task with single agent routing"""
|
|
router = MultiAgentRouter(agents=real_agents())
|
|
|
|
result = router.route_task("Write a fibonacci function")
|
|
|
|
# Check result structure - should be a list of conversation messages
|
|
assert result is not None
|
|
assert isinstance(result, (list, dict))
|
|
assert len(result) > 0 if isinstance(result, list) else True
|
|
|
|
|
|
def test_route_task_multiple_agents():
|
|
"""Test route_task with multiple agent routing"""
|
|
router = MultiAgentRouter(agents=real_agents())
|
|
|
|
result = router.route_task("Research and implement fibonacci")
|
|
|
|
# Check result structure
|
|
assert result is not None
|
|
assert isinstance(result, (list, dict))
|
|
|
|
|
|
def test_route_task_print_on_true():
|
|
"""Test route_task with print_on=True"""
|
|
router = MultiAgentRouter(agents=real_agents(), print_on=True)
|
|
|
|
# Should not raise any exceptions when printing
|
|
result = router.route_task("Test task")
|
|
assert result is not None
|
|
assert isinstance(result, (list, dict))
|
|
|
|
|
|
def test_route_task_print_on_false():
|
|
"""Test route_task with print_on=False"""
|
|
router = MultiAgentRouter(agents=real_agents(), print_on=False)
|
|
|
|
# Should not raise any exceptions when not printing
|
|
result = router.route_task("Test task")
|
|
assert result is not None
|
|
assert isinstance(result, (list, dict))
|
|
|
|
|
|
# ============================================================================
|
|
# ALIAS METHOD TESTS
|
|
# ============================================================================
|
|
|
|
|
|
def test_run_alias():
|
|
"""Test that run() method is an alias for route_task()"""
|
|
router = MultiAgentRouter(agents=real_agents())
|
|
|
|
result1 = router.run(
|
|
"Call your favorite agent to write a fibonacci function"
|
|
)
|
|
result2 = router.route_task(
|
|
"Call your favorite agent to write a fibonacci function"
|
|
)
|
|
|
|
# Results should be valid
|
|
assert result1 is not None
|
|
assert result2 is not None
|
|
assert isinstance(result1, (list, dict))
|
|
assert isinstance(result2, (list, dict))
|
|
|
|
|
|
def test_call_alias():
|
|
"""Test that __call__() method is an alias for route_task()"""
|
|
router = MultiAgentRouter(agents=real_agents())
|
|
|
|
result1 = router(
|
|
"Call your favorite agent to write a fibonacci function"
|
|
)
|
|
result2 = router(
|
|
"Call your favorite agent to write a fibonacci function"
|
|
)
|
|
|
|
# Results should be valid
|
|
assert result1 is not None
|
|
assert result2 is not None
|
|
assert isinstance(result1, (list, dict))
|
|
assert isinstance(result2, (list, dict))
|
|
|
|
|
|
# ============================================================================
|
|
# BATCH PROCESSING TESTS
|
|
# ============================================================================
|
|
|
|
|
|
def test_batch_run():
|
|
"""Test batch_run method"""
|
|
router = MultiAgentRouter(agents=real_agents())
|
|
|
|
tasks = [
|
|
"Call your favorite agent to write a fibonacci function",
|
|
"Call your favorite agent to write a fibonacci function",
|
|
"Call your favorite agent to write a fibonacci function",
|
|
]
|
|
results = router.batch_run(tasks)
|
|
|
|
assert len(results) == 3
|
|
assert all(result is not None for result in results)
|
|
assert all(isinstance(result, (list, dict)) for result in results)
|
|
|
|
|
|
def test_concurrent_batch_run():
|
|
"""Test concurrent_batch_run method"""
|
|
router = MultiAgentRouter(agents=real_agents())
|
|
|
|
tasks = [
|
|
"Call your favorite agent to write a fibonacci function",
|
|
"Call your favorite agent to write a fibonacci function",
|
|
"Call your favorite agent to write a fibonacci function",
|
|
]
|
|
results = router.concurrent_batch_run(tasks)
|
|
|
|
assert len(results) == 3
|
|
assert all(result is not None for result in results)
|
|
assert all(isinstance(result, (list, dict)) for result in results)
|
|
|
|
|
|
# ============================================================================
|
|
# OUTPUT TYPE TESTS
|
|
# ============================================================================
|
|
|
|
|
|
@pytest.mark.parametrize("output_type", ["dict", "json", "string"])
|
|
def test_different_output_types(output_type):
|
|
"""Test different output types"""
|
|
router = MultiAgentRouter(
|
|
agents=real_agents(), output_type=output_type
|
|
)
|
|
|
|
result = router.route_task("Test task")
|
|
|
|
assert result is not None
|
|
# Output format depends on the formatter, but should not raise errors
|
|
assert isinstance(result, (list, dict, str))
|
|
|
|
|
|
# ============================================================================
|
|
# PERFORMANCE AND LOAD TESTS
|
|
# ============================================================================
|
|
|
|
|
|
def test_large_batch_processing():
|
|
"""Test processing a large batch of tasks"""
|
|
router = MultiAgentRouter(agents=real_agents())
|
|
|
|
# Create a smaller number of tasks for testing (reduced from 100 to 5 for performance)
|
|
tasks = [f"Task number {i}" for i in range(5)]
|
|
results = router.batch_run(tasks)
|
|
|
|
assert len(results) == 5
|
|
|
|
|
|
def test_concurrent_large_batch_processing():
|
|
"""Test concurrent processing of a large batch of tasks"""
|
|
router = MultiAgentRouter(agents=real_agents())
|
|
|
|
# Create a small number of tasks for testing
|
|
tasks = [
|
|
f"Route task to your favorite agent to write a fibonacci function {i}"
|
|
for i in range(3)
|
|
]
|
|
results = router.concurrent_batch_run(tasks)
|
|
|
|
assert len(results) == 3
|
|
assert all(result is not None for result in results)
|
|
assert all(isinstance(result, (list, dict)) for result in results)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main([__file__])
|