diff --git a/tests/structs/test_agent_rearrange.py b/tests/structs/test_agent_rearrange.py index 7a9b256c..ffb53278 100644 --- a/tests/structs/test_agent_rearrange.py +++ b/tests/structs/test_agent_rearrange.py @@ -1,128 +1,525 @@ import pytest -from unittest.mock import MagicMock -from swarms import AgentRearrange +from swarms import Agent, AgentRearrange -class MockAgent: - def __init__(self, name): - self.name = name +# ============================================================================ +# Helper Functions +# ============================================================================ - def run(self, task, img=None, *args, **kwargs): - return f"{self.name} processed {task}" - -@pytest.fixture -def mock_agents(): +def create_sample_agents(): + """Create sample agents for testing.""" return [ - MockAgent(name="Agent1"), - MockAgent(name="Agent2"), - MockAgent(name="Agent3"), + Agent( + agent_name="ResearchAgent", + agent_description="Specializes in researching topics", + system_prompt="You are a research specialist. Provide concise answers.", + model_name="gpt-4o-mini", + max_loops=1, + verbose=True, + streaming_on=True, + ), + Agent( + agent_name="WriterAgent", + agent_description="Expert in writing content", + system_prompt="You are a writing expert. Provide concise answers.", + model_name="gpt-4o-mini", + max_loops=1, + verbose=True, + streaming_on=True, + ), + Agent( + agent_name="ReviewerAgent", + agent_description="Expert in reviewing content", + system_prompt="You are a review expert. Provide concise answers.", + model_name="gpt-4o-mini", + max_loops=1, + verbose=True, + streaming_on=True, + ), ] -@pytest.fixture -def agent_rearrange(mock_agents): - return AgentRearrange( - agents=mock_agents, flow="Agent1 -> Agent2 -> Agent3" - ) +# ============================================================================ +# Initialization Tests +# ============================================================================ -def test_initialization(mock_agents): +def test_initialization(): + """Test AgentRearrange initialization.""" + agents = create_sample_agents() + agent_rearrange = AgentRearrange( - agents=mock_agents, flow="Agent1 -> Agent2 -> Agent3" + agents=agents, + flow="ResearchAgent -> WriterAgent -> ReviewerAgent", + max_loops=1, + verbose=True, ) + assert len(agent_rearrange.agents) == 3 - assert agent_rearrange.flow == "Agent1 -> Agent2 -> Agent3" + assert agent_rearrange.flow == "ResearchAgent -> WriterAgent -> ReviewerAgent" + assert agent_rearrange.max_loops == 1 + assert agent_rearrange.verbose is True + print("✓ test_initialization passed") + + +def test_initialization_with_team_awareness(): + """Test AgentRearrange with team_awareness enabled.""" + agents = create_sample_agents() + + agent_rearrange = AgentRearrange( + agents=agents, + flow="ResearchAgent -> WriterAgent -> ReviewerAgent", + team_awareness=True, + verbose=True, + ) + + assert agent_rearrange.flow == "ResearchAgent -> WriterAgent -> ReviewerAgent" + print("✓ test_initialization_with_team_awareness passed") + + +def test_initialization_with_custom_output_type(): + """Test AgentRearrange with custom output types.""" + agents = create_sample_agents() + + for output_type in ["all", "final", "list", "dict"]: + agent_rearrange = AgentRearrange( + agents=agents, + flow="ResearchAgent -> WriterAgent", + output_type=output_type, + verbose=True, + ) + assert agent_rearrange.output_type == output_type + + print("✓ test_initialization_with_custom_output_type passed") + + +# ============================================================================ +# Agent Management Tests +# ============================================================================ -def test_add_agent(agent_rearrange): - new_agent = MockAgent(name="Agent4") +def test_add_agent(): + """Test adding an agent to AgentRearrange.""" + agents = create_sample_agents()[:2] # Only 2 agents + + agent_rearrange = AgentRearrange( + agents=agents, + flow="ResearchAgent -> WriterAgent", + verbose=True, + ) + + new_agent = Agent( + agent_name="EditorAgent", + model_name="gpt-4o-mini", + max_loops=1, + verbose=True, + streaming_on=True, + ) + agent_rearrange.add_agent(new_agent) - assert "Agent4" in agent_rearrange.agents + assert "EditorAgent" in agent_rearrange.agents + assert len(agent_rearrange.agents) == 3 + print("✓ test_add_agent passed") -def test_remove_agent(agent_rearrange): - agent_rearrange.remove_agent("Agent2") - assert "Agent2" not in agent_rearrange.agents +def test_remove_agent(): + """Test removing an agent from AgentRearrange.""" + agents = create_sample_agents() + + agent_rearrange = AgentRearrange( + agents=agents, + flow="ResearchAgent -> WriterAgent -> ReviewerAgent", + verbose=True, + ) + + agent_rearrange.remove_agent("ReviewerAgent") + assert "ReviewerAgent" not in agent_rearrange.agents + assert len(agent_rearrange.agents) == 2 + print("✓ test_remove_agent passed") -def test_add_agents(agent_rearrange): - new_agents = [MockAgent(name="Agent4"), MockAgent(name="Agent5")] +def test_add_agents(): + """Test adding multiple agents to AgentRearrange.""" + agents = create_sample_agents()[:1] # Start with 1 agent + + agent_rearrange = AgentRearrange( + agents=agents, + flow="ResearchAgent", + verbose=True, + ) + + new_agents = [ + Agent( + agent_name="Agent4", + model_name="gpt-4o-mini", + max_loops=1, + verbose=True, + streaming_on=True, + ), + Agent( + agent_name="Agent5", + model_name="gpt-4o-mini", + max_loops=1, + verbose=True, + streaming_on=True, + ), + ] + agent_rearrange.add_agents(new_agents) assert "Agent4" in agent_rearrange.agents assert "Agent5" in agent_rearrange.agents + assert len(agent_rearrange.agents) == 3 + print("✓ test_add_agents passed") -def test_validate_flow_valid(agent_rearrange): +# ============================================================================ +# Flow Validation Tests +# ============================================================================ + + +def test_validate_flow_valid(): + """Test flow validation with valid flow.""" + agents = create_sample_agents() + + agent_rearrange = AgentRearrange( + agents=agents, + flow="ResearchAgent -> WriterAgent -> ReviewerAgent", + verbose=True, + ) + assert agent_rearrange.validate_flow() is True + print("✓ test_validate_flow_valid passed") + + +def test_validate_flow_invalid(): + """Test flow validation with invalid agent name in flow.""" + agents = create_sample_agents() + + agent_rearrange = AgentRearrange( + agents=agents, + flow="ResearchAgent -> WriterAgent -> ReviewerAgent", + verbose=True, + ) + # Change to invalid flow + agent_rearrange.flow = "ResearchAgent -> NonExistentAgent" -def test_validate_flow_invalid(agent_rearrange): - agent_rearrange.flow = "Agent1 -> Agent4" - with pytest.raises(ValueError): + try: agent_rearrange.validate_flow() + assert False, "Should have raised ValueError" + except ValueError as e: + assert "not registered" in str(e) + print("✓ test_validate_flow_invalid passed") + + +def test_validate_flow_no_arrow(): + """Test flow validation without arrow syntax.""" + agents = create_sample_agents() + + agent_rearrange = AgentRearrange( + agents=agents, + flow="ResearchAgent -> WriterAgent", + verbose=True, + ) + + agent_rearrange.flow = "ResearchAgent WriterAgent" + + try: + agent_rearrange.validate_flow() + assert False, "Should have raised ValueError" + except ValueError as e: + assert "'->" in str(e) + print("✓ test_validate_flow_no_arrow passed") + + +# ============================================================================ +# Flow Pattern Tests +# ============================================================================ + + +def test_set_custom_flow(): + """Test setting custom flow.""" + agents = create_sample_agents() + + agent_rearrange = AgentRearrange( + agents=agents, + flow="ResearchAgent -> WriterAgent", + verbose=True, + ) + new_flow = "WriterAgent -> ResearchAgent -> ReviewerAgent" + agent_rearrange.set_custom_flow(new_flow) + assert agent_rearrange.flow == new_flow + print("✓ test_set_custom_flow passed") -def test_run(agent_rearrange): - result = agent_rearrange.run("Test Task") - assert ( - result - == "Agent1 processed Test Task; Agent2 processed Agent1 processed Test Task; Agent3 processed Agent2 processed Agent1 processed Test Task" + +def test_sequential_flow(): + """Test sequential flow execution.""" + agents = create_sample_agents() + + agent_rearrange = AgentRearrange( + agents=agents, + flow="ResearchAgent -> WriterAgent", + max_loops=1, + verbose=True, ) + result = agent_rearrange.run("What is 2+2?") + assert result is not None + print("✓ test_sequential_flow passed") + -def test_run_with_custom_tasks(agent_rearrange): - custom_tasks = {"Agent2": "Custom Task"} - result = agent_rearrange.run( - "Test Task", custom_tasks=custom_tasks +def test_concurrent_flow(): + """Test concurrent flow execution with comma syntax.""" + agents = create_sample_agents() + + agent_rearrange = AgentRearrange( + agents=agents, + flow="ResearchAgent, WriterAgent -> ReviewerAgent", + max_loops=1, + verbose=True, ) - assert ( - result - == "Agent1 processed Test Task; Agent2 processed Custom Task; Agent3 processed Agent2 processed Custom Task" + + result = agent_rearrange.run("What is 3+3?") + assert result is not None + print("✓ test_concurrent_flow passed") + + +# ============================================================================ +# Sequential Awareness Tests +# ============================================================================ + + +def test_get_sequential_flow_structure(): + """Test getting sequential flow structure.""" + agents = create_sample_agents() + + agent_rearrange = AgentRearrange( + agents=agents, + flow="ResearchAgent -> WriterAgent -> ReviewerAgent", + verbose=True, ) + flow_structure = agent_rearrange.get_sequential_flow_structure() + assert flow_structure is not None + assert isinstance(flow_structure, str) + print("✓ test_get_sequential_flow_structure passed") + -def test_run_with_human_intervention(agent_rearrange): - agent_rearrange.human_in_the_loop = True - agent_rearrange.custom_human_in_the_loop = MagicMock( - return_value="Human processed Task" +def test_get_agent_sequential_awareness(): + """Test getting agent sequential awareness.""" + agents = create_sample_agents() + + agent_rearrange = AgentRearrange( + agents=agents, + flow="ResearchAgent -> WriterAgent -> ReviewerAgent", + verbose=True, ) - agent_rearrange.flow = "Agent1 -> H -> Agent3" - result = agent_rearrange.run("Test Task") - assert ( - result - == "Agent1 processed Test Task; Human processed Task; Agent3 processed Human processed Task" + + awareness = agent_rearrange.get_agent_sequential_awareness("WriterAgent") + assert awareness is not None + assert isinstance(awareness, str) + print("✓ test_get_agent_sequential_awareness passed") + + +# ============================================================================ +# Execution Tests +# ============================================================================ + + +def test_run_basic(): + """Test basic run execution.""" + agents = create_sample_agents() + + agent_rearrange = AgentRearrange( + agents=agents, + flow="ResearchAgent -> WriterAgent", + max_loops=1, + verbose=True, ) + result = agent_rearrange.run("Calculate 5+5") + assert result is not None + print("✓ test_run_basic passed") + -def test_run_sub_swarm(agent_rearrange): - sub_swarm_flow = "Agent1 -> Agent3" - agent_rearrange.add_sub_swarm("SubSwarm1", sub_swarm_flow) - result = agent_rearrange.run_sub_swarm( - "SubSwarm1", "Sub Task", None +def test_run_with_different_output_types(): + """Test run with different output types.""" + agents = create_sample_agents() + + for output_type in ["all", "final"]: + agent_rearrange = AgentRearrange( + agents=agents, + flow="ResearchAgent -> WriterAgent", + output_type=output_type, + max_loops=1, + verbose=True, + ) + + result = agent_rearrange.run("What is the capital of France?") + assert result is not None + + print("✓ test_run_with_different_output_types passed") + + +def test_callable_execution(): + """Test __call__ method.""" + agents = create_sample_agents() + + agent_rearrange = AgentRearrange( + agents=agents, + flow="ResearchAgent -> WriterAgent", + max_loops=1, + verbose=True, ) - assert ( - result - == "Agent1 processed Sub Task; Agent3 processed Agent1 processed Sub Task" + + result = agent_rearrange("What is 10+10?") + assert result is not None + print("✓ test_callable_execution passed") + + +# ============================================================================ +# Batch Processing Tests +# ============================================================================ + + +def test_batch_run(): + """Test batch processing.""" + agents = create_sample_agents() + + agent_rearrange = AgentRearrange( + agents=agents, + flow="ResearchAgent -> WriterAgent", + max_loops=1, + verbose=True, ) + tasks = ["What is 1+1?", "What is 2+2?"] + results = agent_rearrange.batch_run(tasks, batch_size=2) + + assert results is not None + assert len(results) == 2 + print("✓ test_batch_run passed") + + +def test_concurrent_run(): + """Test concurrent execution.""" + agents = create_sample_agents() + + agent_rearrange = AgentRearrange( + agents=agents, + flow="ResearchAgent -> WriterAgent", + max_loops=1, + verbose=True, + ) + + tasks = ["What is 3+3?", "What is 4+4?"] + results = agent_rearrange.concurrent_run(tasks, max_workers=2) + + assert results is not None + assert len(results) == 2 + print("✓ test_concurrent_run passed") -def test_process_agent_or_swarm(agent_rearrange): - result = agent_rearrange.process_agent_or_swarm( - "Agent1", "Process Task", None + +# ============================================================================ +# Serialization Tests +# ============================================================================ + + +def test_to_dict(): + """Test serialization to dictionary.""" + agents = create_sample_agents() + + agent_rearrange = AgentRearrange( + agents=agents, + flow="ResearchAgent -> WriterAgent", + verbose=True, ) - assert result == "Agent1 processed Process Task" + + result_dict = agent_rearrange.to_dict() + assert isinstance(result_dict, dict) + assert "name" in result_dict + assert "flow" in result_dict + assert "agents" in result_dict + print("✓ test_to_dict passed") -def test_track_history(agent_rearrange): - agent_rearrange.track_history("Agent1", "Task Result") - assert agent_rearrange.swarm_history["Agent1"] == ["Task Result"] +# ============================================================================ +# Integration Tests +# ============================================================================ -def test_human_intervention(agent_rearrange): - agent_rearrange.human_in_the_loop = True - agent_rearrange.custom_human_in_the_loop = MagicMock( - return_value="Human processed Task" +def test_complete_workflow(): + """Test complete workflow with all features.""" + agents = create_sample_agents() + + agent_rearrange = AgentRearrange( + name="test-workflow", + description="Test complete workflow", + agents=agents, + flow="ResearchAgent -> WriterAgent -> ReviewerAgent", + max_loops=1, + team_awareness=True, + verbose=True, + output_type="all", ) - result = agent_rearrange.human_intervention("Task") - assert result == "Human processed Task" + + # Single run + result1 = agent_rearrange.run("What is Python?") + assert result1 is not None + + # Get flow structure + flow_structure = agent_rearrange.get_sequential_flow_structure() + assert flow_structure is not None + + # Serialize + result_dict = agent_rearrange.to_dict() + assert isinstance(result_dict, dict) + + print("✓ test_complete_workflow passed") + + +def main(): + """Run all tests.""" + tests = [ + test_initialization, + test_initialization_with_team_awareness, + test_initialization_with_custom_output_type, + test_add_agent, + test_remove_agent, + test_add_agents, + test_validate_flow_valid, + test_validate_flow_invalid, + test_validate_flow_no_arrow, + test_set_custom_flow, + test_sequential_flow, + test_concurrent_flow, + test_get_sequential_flow_structure, + test_get_agent_sequential_awareness, + test_run_basic, + test_run_with_different_output_types, + test_callable_execution, + test_batch_run, + test_concurrent_run, + test_to_dict, + test_complete_workflow, + ] + + print("="*60) + print("Running AgentRearrange Tests") + print("="*60) + + for test in tests: + try: + test() + except Exception as e: + print(f"✗ {test.__name__} failed: {e}") + + print("="*60) + print("All tests completed!") + print("="*60) + + +if __name__ == "__main__": + pytest.main([__file__, "-v"])