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.
492 lines
14 KiB
492 lines
14 KiB
import pytest
|
|
|
|
from swarms.structs.swarm_router import (
|
|
SwarmRouter,
|
|
SwarmRouterConfig,
|
|
SwarmRouterRunError,
|
|
SwarmRouterConfigError,
|
|
Document,
|
|
)
|
|
from swarms.structs.agent import Agent
|
|
|
|
# ============================================================================
|
|
# Helper Functions
|
|
# ============================================================================
|
|
|
|
|
|
def create_sample_agents():
|
|
"""Create sample agents for testing."""
|
|
return [
|
|
Agent(
|
|
agent_name="ResearchAgent",
|
|
agent_description="Specializes in researching topics",
|
|
system_prompt="You are a research specialist.",
|
|
model_name="gpt-4o-mini",
|
|
max_loops=1,
|
|
verbose=False,
|
|
print_on=False,
|
|
),
|
|
Agent(
|
|
agent_name="CodeAgent",
|
|
agent_description="Expert in coding",
|
|
system_prompt="You are a coding expert.",
|
|
model_name="gpt-4o-mini",
|
|
max_loops=1,
|
|
verbose=False,
|
|
print_on=False,
|
|
),
|
|
]
|
|
|
|
|
|
# ============================================================================
|
|
# Initialization Tests
|
|
# ============================================================================
|
|
|
|
|
|
def test_default_initialization():
|
|
"""Test SwarmRouter with default parameters."""
|
|
router = SwarmRouter()
|
|
|
|
assert router.name == "swarm-router"
|
|
assert (
|
|
router.description == "Routes your task to the desired swarm"
|
|
)
|
|
assert router.max_loops == 1
|
|
assert router.agents == []
|
|
assert router.swarm_type == "SequentialWorkflow"
|
|
assert router.autosave is False
|
|
assert router.return_json is False
|
|
assert router.auto_generate_prompts is False
|
|
assert router.shared_memory_system is None
|
|
assert router.rules is None
|
|
assert router.documents == []
|
|
assert router.output_type == "dict-all-except-first"
|
|
assert router.verbose is False
|
|
assert router.telemetry_enabled is False
|
|
|
|
|
|
def test_custom_initialization():
|
|
"""Test SwarmRouter with custom parameters."""
|
|
sample_agents = create_sample_agents()
|
|
|
|
router = SwarmRouter(
|
|
name="test-router",
|
|
description="Test router description",
|
|
max_loops=3,
|
|
agents=sample_agents,
|
|
swarm_type="ConcurrentWorkflow",
|
|
autosave=True,
|
|
return_json=True,
|
|
auto_generate_prompts=True,
|
|
rules="Test rules",
|
|
documents=["doc1.txt", "doc2.txt"],
|
|
output_type="json",
|
|
verbose=False, # Keep quiet for tests
|
|
telemetry_enabled=False,
|
|
)
|
|
|
|
assert router.name == "test-router"
|
|
assert router.description == "Test router description"
|
|
assert router.max_loops == 3
|
|
assert router.agents == sample_agents
|
|
assert router.swarm_type == "ConcurrentWorkflow"
|
|
assert router.autosave is True
|
|
assert router.return_json is True
|
|
assert router.auto_generate_prompts is True
|
|
assert router.rules == "Test rules"
|
|
assert router.documents == ["doc1.txt", "doc2.txt"]
|
|
assert router.output_type == "json"
|
|
assert router.verbose is False
|
|
assert router.telemetry_enabled is False
|
|
|
|
|
|
def test_initialization_with_heavy_swarm_config():
|
|
"""Test SwarmRouter with HeavySwarm specific configuration."""
|
|
sample_agents = create_sample_agents()
|
|
|
|
router = SwarmRouter(
|
|
agents=sample_agents,
|
|
swarm_type="HeavySwarm",
|
|
heavy_swarm_loops_per_agent=2,
|
|
heavy_swarm_question_agent_model_name="gpt-4o-mini",
|
|
heavy_swarm_worker_model_name="gpt-4o-mini",
|
|
heavy_swarm_swarm_show_output=False,
|
|
)
|
|
|
|
assert router.swarm_type == "HeavySwarm"
|
|
assert router.heavy_swarm_loops_per_agent == 2
|
|
assert (
|
|
router.heavy_swarm_question_agent_model_name == "gpt-4o-mini"
|
|
)
|
|
assert router.heavy_swarm_worker_model_name == "gpt-4o-mini"
|
|
assert router.heavy_swarm_swarm_show_output is False
|
|
|
|
|
|
def test_initialization_with_agent_rearrange_config():
|
|
"""Test SwarmRouter with AgentRearrange specific configuration."""
|
|
sample_agents = create_sample_agents()
|
|
|
|
router = SwarmRouter(
|
|
agents=sample_agents,
|
|
swarm_type="AgentRearrange",
|
|
rearrange_flow="ResearchAgent -> CodeAgent",
|
|
)
|
|
|
|
assert router.swarm_type == "AgentRearrange"
|
|
assert router.rearrange_flow == "ResearchAgent -> CodeAgent"
|
|
|
|
|
|
# ============================================================================
|
|
# Configuration Tests
|
|
# ============================================================================
|
|
|
|
|
|
def test_initialization_with_shared_memory():
|
|
"""Test SwarmRouter with shared memory system."""
|
|
sample_agents = create_sample_agents()
|
|
|
|
router = SwarmRouter(
|
|
agents=sample_agents,
|
|
shared_memory_system=None, # Test with None for now
|
|
)
|
|
|
|
assert router.shared_memory_system is None
|
|
|
|
|
|
def test_initialization_with_worker_tools():
|
|
"""Test SwarmRouter with worker tools."""
|
|
sample_agents = create_sample_agents()
|
|
|
|
router = SwarmRouter(
|
|
agents=sample_agents,
|
|
worker_tools=[], # Empty list for now
|
|
)
|
|
|
|
assert router.worker_tools == []
|
|
|
|
|
|
# ============================================================================
|
|
# Document Management Tests
|
|
# ============================================================================
|
|
|
|
|
|
def test_document_creation():
|
|
"""Test Document creation."""
|
|
doc = Document(
|
|
file_path="/path/to/test/document.txt",
|
|
data="This is test content",
|
|
)
|
|
|
|
assert doc.file_path == "/path/to/test/document.txt"
|
|
assert doc.data == "This is test content"
|
|
|
|
|
|
def test_router_with_documents():
|
|
"""Test SwarmRouter with document configuration."""
|
|
sample_agents = create_sample_agents()
|
|
documents = [
|
|
Document(file_path="/path/to/doc1.txt", data="Content1"),
|
|
Document(file_path="/path/to/doc2.txt", data="Content2"),
|
|
]
|
|
|
|
router = SwarmRouter(
|
|
agents=sample_agents,
|
|
documents=documents,
|
|
)
|
|
|
|
assert len(router.documents) == 2
|
|
assert router.documents[0].file_path == "/path/to/doc1.txt"
|
|
assert router.documents[1].file_path == "/path/to/doc2.txt"
|
|
|
|
|
|
# ============================================================================
|
|
# Configuration Class Tests
|
|
# ============================================================================
|
|
|
|
|
|
def test_swarm_router_config_creation():
|
|
"""Test SwarmRouterConfig creation."""
|
|
config = SwarmRouterConfig(
|
|
name="test-config",
|
|
description="Test configuration",
|
|
swarm_type="SequentialWorkflow",
|
|
rearrange_flow=None,
|
|
rules=None,
|
|
multi_agent_collab_prompt=True,
|
|
task="Test task",
|
|
)
|
|
|
|
assert config.name == "test-config"
|
|
assert config.description == "Test configuration"
|
|
assert config.swarm_type == "SequentialWorkflow"
|
|
assert config.task == "Test task"
|
|
|
|
|
|
def test_router_with_config():
|
|
"""Test SwarmRouter initialization matches config structure."""
|
|
sample_agents = create_sample_agents()
|
|
config = SwarmRouterConfig(
|
|
name="config-router",
|
|
description="Router from config",
|
|
swarm_type="SequentialWorkflow",
|
|
rearrange_flow=None,
|
|
rules="Test rules",
|
|
multi_agent_collab_prompt=False,
|
|
task="Test task",
|
|
)
|
|
|
|
# SwarmRouter doesn't accept config directly, but we can verify config is valid
|
|
assert config.name == "config-router"
|
|
assert config.description == "Router from config"
|
|
assert config.swarm_type == "SequentialWorkflow"
|
|
assert config.rules == "Test rules"
|
|
|
|
# Create router with matching parameters
|
|
router = SwarmRouter(
|
|
name=config.name,
|
|
description=config.description,
|
|
agents=sample_agents,
|
|
swarm_type=config.swarm_type,
|
|
rules=config.rules,
|
|
)
|
|
|
|
assert router.name == config.name
|
|
assert router.description == config.description
|
|
assert router.swarm_type == config.swarm_type
|
|
assert router.rules == config.rules
|
|
|
|
|
|
# ============================================================================
|
|
# Basic Execution Tests
|
|
# ============================================================================
|
|
|
|
|
|
def test_run_with_sequential_workflow():
|
|
"""Test running SwarmRouter with SequentialWorkflow."""
|
|
sample_agents = create_sample_agents()
|
|
|
|
router = SwarmRouter(
|
|
agents=sample_agents,
|
|
swarm_type="SequentialWorkflow",
|
|
verbose=False,
|
|
)
|
|
|
|
result = router.run("What is 2+2?")
|
|
assert result is not None
|
|
|
|
|
|
def test_run_with_no_agents():
|
|
"""Test running SwarmRouter with no agents."""
|
|
router = SwarmRouter()
|
|
|
|
with pytest.raises(RuntimeError):
|
|
router.run("Test task")
|
|
|
|
|
|
def test_run_with_empty_task():
|
|
"""Test running SwarmRouter with empty task."""
|
|
sample_agents = create_sample_agents()
|
|
|
|
router = SwarmRouter(agents=sample_agents, verbose=False)
|
|
|
|
# Empty task is allowed, router will pass it to the swarm
|
|
result = router.run("")
|
|
assert result is not None
|
|
|
|
|
|
def test_run_with_none_task():
|
|
"""Test running SwarmRouter with None task."""
|
|
sample_agents = create_sample_agents()
|
|
|
|
router = SwarmRouter(agents=sample_agents, verbose=False)
|
|
|
|
# None task is allowed, router will pass it to the swarm
|
|
result = router.run(None)
|
|
assert result is not None
|
|
|
|
|
|
# ============================================================================
|
|
# Batch Processing Tests
|
|
# ============================================================================
|
|
|
|
|
|
def test_batch_run_with_tasks():
|
|
"""Test batch processing with multiple tasks."""
|
|
sample_agents = create_sample_agents()
|
|
|
|
router = SwarmRouter(
|
|
agents=sample_agents,
|
|
verbose=False,
|
|
)
|
|
|
|
tasks = ["What is 1+1?", "What is 2+2?"]
|
|
results = router.batch_run(tasks)
|
|
|
|
assert len(results) == 2
|
|
assert all(result is not None for result in results)
|
|
|
|
|
|
def test_batch_run_with_empty_tasks():
|
|
"""Test batch processing with empty task list."""
|
|
sample_agents = create_sample_agents()
|
|
|
|
router = SwarmRouter(agents=sample_agents)
|
|
|
|
results = router.batch_run([])
|
|
assert results == []
|
|
|
|
|
|
def test_batch_run_with_no_agents():
|
|
"""Test batch processing with no agents."""
|
|
router = SwarmRouter()
|
|
|
|
with pytest.raises(RuntimeError):
|
|
router.batch_run(["Test task"])
|
|
|
|
|
|
# ============================================================================
|
|
# Call Method Tests
|
|
# ============================================================================
|
|
|
|
|
|
def test_call_method():
|
|
"""Test __call__ method."""
|
|
sample_agents = create_sample_agents()
|
|
|
|
router = SwarmRouter(
|
|
agents=sample_agents,
|
|
verbose=False,
|
|
)
|
|
|
|
result = router("What is the capital of France?")
|
|
assert result is not None
|
|
|
|
|
|
def test_call_with_image():
|
|
"""Test __call__ method with image."""
|
|
sample_agents = create_sample_agents()
|
|
|
|
router = SwarmRouter(
|
|
agents=sample_agents,
|
|
verbose=False,
|
|
)
|
|
|
|
# Test with None image (no actual image processing)
|
|
result = router("Describe this image", img=None)
|
|
assert result is not None
|
|
|
|
|
|
# ============================================================================
|
|
# Output Type Tests
|
|
# ============================================================================
|
|
|
|
|
|
def test_different_output_types():
|
|
"""Test router with different output types."""
|
|
sample_agents = create_sample_agents()
|
|
|
|
for output_type in ["dict", "json", "string", "list"]:
|
|
router = SwarmRouter(
|
|
agents=sample_agents,
|
|
output_type=output_type,
|
|
verbose=False,
|
|
)
|
|
|
|
result = router.run("Simple test task")
|
|
assert result is not None
|
|
|
|
|
|
# ============================================================================
|
|
# Error Handling Tests
|
|
# ============================================================================
|
|
|
|
|
|
def test_swarm_router_run_error():
|
|
"""Test SwarmRouterRunError exception."""
|
|
error = SwarmRouterRunError("Test error message")
|
|
assert str(error) == "Test error message"
|
|
|
|
|
|
def test_swarm_router_config_error():
|
|
"""Test SwarmRouterConfigError exception."""
|
|
error = SwarmRouterConfigError("Config error message")
|
|
assert str(error) == "Config error message"
|
|
|
|
|
|
def test_invalid_swarm_type():
|
|
"""Test router with invalid swarm type."""
|
|
sample_agents = create_sample_agents()
|
|
|
|
# This should not raise an error during initialization
|
|
router = SwarmRouter(
|
|
agents=sample_agents,
|
|
swarm_type="InvalidSwarmType",
|
|
)
|
|
|
|
# But should raise ValueError during execution when creating swarm
|
|
with pytest.raises(
|
|
ValueError, match="Invalid swarm type: InvalidSwarmType"
|
|
):
|
|
router.run("Test task")
|
|
|
|
|
|
# ============================================================================
|
|
# Integration Tests
|
|
# ============================================================================
|
|
|
|
|
|
def test_complete_workflow():
|
|
"""Test complete workflow from initialization to execution."""
|
|
# Create agents
|
|
agents = create_sample_agents()
|
|
|
|
# Create router with configuration
|
|
router = SwarmRouter(
|
|
name="integration-test-router",
|
|
description="Router for integration testing",
|
|
agents=agents,
|
|
swarm_type="SequentialWorkflow",
|
|
max_loops=1,
|
|
verbose=False,
|
|
output_type="string",
|
|
)
|
|
|
|
# Execute single task
|
|
result = router.run("Calculate the sum of 5 and 7")
|
|
assert result is not None
|
|
|
|
# Execute batch tasks
|
|
tasks = [
|
|
"What is 10 + 15?",
|
|
"What is 20 - 8?",
|
|
"What is 6 * 7?",
|
|
]
|
|
batch_results = router.batch_run(tasks)
|
|
|
|
assert len(batch_results) == 3
|
|
assert all(result is not None for result in batch_results)
|
|
|
|
|
|
def test_router_reconfiguration():
|
|
"""Test reconfiguring router after initialization."""
|
|
sample_agents = create_sample_agents()
|
|
|
|
router = SwarmRouter(agents=sample_agents)
|
|
|
|
# Change configuration
|
|
router.max_loops = 3
|
|
router.output_type = "json"
|
|
router.verbose = False
|
|
|
|
assert router.max_loops == 3
|
|
assert router.output_type == "json"
|
|
assert router.verbose is False
|
|
|
|
# Test execution with new configuration
|
|
result = router.run("Test reconfiguration")
|
|
assert result is not None
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main([__file__, "-v"])
|