test_auto_swarms_builder.py

import pytest
from dotenv import load_dotenv

from swarms.structs.agent import Agent
from swarms.structs.auto_swarm_builder import (
    AgentSpec,
    AutoSwarmBuilder,
)
from swarms.structs.ma_utils import set_random_models_for_agents

load_dotenv()


def print_separator():
    """Print a separator line for test output formatting."""
    print("\n" + "=" * 50)


def test_initialization():
    """Test basic initialization of AutoSwarmBuilder"""
    print_separator()
    print("Testing AutoSwarmBuilder Initialization")
    try:
        swarm = AutoSwarmBuilder(
            name="TestSwarm",
            description="A test swarm for validation",
            verbose=True,
            max_loops=2,
        )

        print("✓ Created swarm with configuration:")
        print(f"  - Name: {swarm.name}")
        print(f"  - Description: {swarm.description}")
        print(f"  - Max loops: {swarm.max_loops}")
        print(f"  - Verbose: {swarm.verbose}")
        print("✓ Initialization test passed")
        return swarm
    except Exception as e:
        print(f"✗ Initialization test failed: {str(e)}")
        raise


def test_agent_building():
    """Test building individual agents from specs"""
    print_separator()
    print("Testing Agent Building")
    try:
        swarm = AutoSwarmBuilder()
        specs = {
            "agents": [
                {
                    "agent_name": "TestAgent",
                    "description": "A test agent",
                    "system_prompt": "You are a test agent",
                    "max_loops": 1,
                }
            ]
        }
        agents = swarm.create_agents_from_specs(specs)
        agent = agents[0]

        # Create agent from spec
        agents = swarm.create_agents_from_specs({"agents": [agent_spec]})
        agent = agents[0]

        print("✓ Built agent with configuration:")
        print(f"  - Name: {agent.agent_name}")
        print(f"  - Description: {agent.agent_description}")
        print(f"  - Max loops: {agent.max_loops}")
        print("✓ Agent building test passed")
        return agent
    except Exception as e:
        print(f"✗ Agent building test failed: {str(e)}")
        raise


def test_agent_creation():
    """Test creating multiple agents for a task"""
    print_separator()
    print("Testing Agent Creation from Task")
    try:
        import json

        swarm = AutoSwarmBuilder(
            name="ResearchSwarm",
            description="A swarm for research tasks",
        )
        task = "Research the latest developments in quantum computing"
        # create_agents returns a JSON string
        agent_specs_json = swarm.create_agents(task)
        # Parse JSON string to dict
        agent_specs = json.loads(agent_specs_json)
        # Convert specs to actual Agent objects
        agents = swarm.create_agents_from_specs(agent_specs)

        print("✓ Created agents for research task:")
        for i, agent in enumerate(agents, 1):
            print(f"  Agent {i}:")
            print(f"    - Name: {agent.agent_name}")
            print(f"    - Description: {agent.agent_description}")
        print(f"✓ Created {len(agents)} agents successfully")
        return agents
    except Exception as e:
        print(f"✗ Agent creation test failed: {str(e)}")
        raise


def test_swarm_routing():
    """Test routing tasks through the swarm"""
    print_separator()
    print("Testing Swarm Routing")
    try:
        swarm = AutoSwarmBuilder(
            name="RouterTestSwarm",
            description="Testing routing capabilities",
        )
        agents = (
            test_agent_creation()
        )  # Get agents from previous test
        task = "Analyze the impact of AI on healthcare"

        print("Starting task routing...")
        result = swarm.initialize_swarm_router(agents, task)

        print("✓ Task routed successfully")
        print(
            f"  - Result length: {len(str(result)) if result else 0} characters"
        )
        print("✓ Swarm routing test passed")
        return result
    except Exception as e:
        print(f"✗ Swarm routing test failed: {str(e)}")
        raise


def test_full_swarm_execution():
    """Test complete swarm execution with a real task"""
    print_separator()
    print("Testing Full Swarm Execution")
    try:
        swarm = AutoSwarmBuilder(
            name="FullTestSwarm",
            description="Testing complete swarm functionality",
            max_loops=1,
        )
        task = (
            "Create a summary of recent advances in renewable energy"
        )

        print("Starting full swarm execution...")
        result = swarm.run(task)

        print("✓ Full swarm execution completed:")
        print(f"  - Output generated: {bool(result)}")
        print(
            f"  - Output length: {len(str(result)) if result else 0} characters"
        )
        print("✓ Full swarm execution test passed")
        return result
    except Exception as e:
        print(f"✗ Full swarm execution test failed: {str(e)}")
        raise


def test_error_handling():
    """Test error handling in swarm operations"""
    print_separator()
    print("Testing Error Handling")
    try:
        swarm = AutoSwarmBuilder()

        # Test with invalid agent configuration
        print("Testing invalid agent configuration...")
        try:
            swarm.create_agents_from_specs(
                {"agents": [{"agent_name": ""}]}
            )
            print(
                "✗ Should have raised an error for empty agent configuration"
            )
        except Exception as e:
            print(
                f"✓ Correctly handled invalid agent configuration: {type(e).__name__}"
            )

        # Test with None task
        print("\nTesting None task...")
        try:
            swarm.run(None)
            print("✗ Should have raised an error for None task")
        except Exception as e:
            print(
                f"✓ Correctly handled None task: {type(e).__name__}"
            )

        print("✓ Error handling test passed")
    except Exception as e:
        print(f"✗ Error handling test failed: {str(e)}")
        raise


def run_all_tests():
    """Run complete test suite"""
    print("\n=== Starting AutoSwarmBuilder Test Suite ===\n")

    try:
        # Run all tests in sequence
        test_initialization()
        test_agent_building()
        test_agent_creation()
        test_swarm_routing()
        test_full_swarm_execution()
        test_error_handling()

        print_separator()
        print("🎉 All tests completed successfully!")

    except Exception as e:
        print_separator()
        print(f"❌ Test suite failed: {str(e)}")
        raise


# Bug Fix Tests (from test_auto_swarm_builder_fix.py)
class TestAutoSwarmBuilderFix:
    """Tests for bug #1115 fix in AutoSwarmBuilder."""

    def test_create_agents_from_specs_with_dict(self):
        """Test that create_agents_from_specs handles dict input correctly."""
        builder = AutoSwarmBuilder()

        # Create specs as a dictionary
        specs = {
            "agents": [
                {
                    "agent_name": "test_agent_1",
                    "description": "Test agent 1 description",
                    "system_prompt": "You are a helpful assistant",
                    "model_name": "gpt-4o-mini",
                    "max_loops": 1,
                }
            ]
        }

        agents = builder.create_agents_from_specs(specs)

        # Verify agents were created correctly
        assert len(agents) == 1
        assert isinstance(agents[0], Agent)
        assert agents[0].agent_name == "test_agent_1"

        # Verify description was mapped to agent_description
        assert hasattr(agents[0], "agent_description")
        assert (
            agents[0].agent_description == "Test agent 1 description"
        )

    def test_create_agents_from_specs_with_pydantic(self):
        """Test that create_agents_from_specs handles Pydantic model input correctly.

        This is the main test for bug #1115 - it verifies that AgentSpec
        Pydantic models can be unpacked correctly.
        """
        builder = AutoSwarmBuilder()

        # Create specs as Pydantic AgentSpec objects
        agent_spec = AgentSpec(
            agent_name="test_agent_pydantic",
            description="Pydantic test agent",
            system_prompt="You are a helpful assistant",
            model_name="gpt-4o-mini",
            max_loops=1,
        )

        specs = {"agents": [agent_spec]}

        agents = builder.create_agents_from_specs(specs)

        # Verify agents were created correctly
        assert len(agents) == 1
        assert isinstance(agents[0], Agent)
        assert agents[0].agent_name == "test_agent_pydantic"

        # Verify description was mapped to agent_description
        assert hasattr(agents[0], "agent_description")
        assert agents[0].agent_description == "Pydantic test agent"

    def test_parameter_name_mapping(self):
        """Test that 'description' field maps to 'agent_description' correctly."""
        builder = AutoSwarmBuilder()

        # Test with dict that has 'description'
        specs = {
            "agents": [
                {
                    "agent_name": "mapping_test",
                    "description": "This should map to agent_description",
                    "system_prompt": "You are helpful",
                }
            ]
        }

        agents = builder.create_agents_from_specs(specs)

        assert len(agents) == 1
        agent = agents[0]

        # Verify description was mapped
        assert hasattr(agent, "agent_description")
        assert (
            agent.agent_description
            == "This should map to agent_description"
        )

    def test_create_agents_from_specs_mixed_input(self):
        """Test that create_agents_from_specs handles mixed dict and Pydantic input."""
        builder = AutoSwarmBuilder()

        # Mix of dict and Pydantic objects
        dict_spec = {
            "agent_name": "dict_agent",
            "description": "Dict agent description",
            "system_prompt": "You are helpful",
        }

        pydantic_spec = AgentSpec(
            agent_name="pydantic_agent",
            description="Pydantic agent description",
            system_prompt="You are smart",
        )

        specs = {"agents": [dict_spec, pydantic_spec]}

        agents = builder.create_agents_from_specs(specs)

        # Verify both agents were created
        assert len(agents) == 2
        assert all(isinstance(agent, Agent) for agent in agents)

        # Verify both have correct descriptions
        dict_agent = next(
            a for a in agents if a.agent_name == "dict_agent"
        )
        pydantic_agent = next(
            a for a in agents if a.agent_name == "pydantic_agent"
        )

        assert (
            dict_agent.agent_description == "Dict agent description"
        )
        assert (
            pydantic_agent.agent_description
            == "Pydantic agent description"
        )

    def test_set_random_models_for_agents_with_valid_agents(self):
        """Test set_random_models_for_agents with proper Agent objects."""
        # Create proper Agent objects
        agents = [
            Agent(
                agent_name="agent1",
                system_prompt="You are agent 1",
                max_loops=1,
            ),
            Agent(
                agent_name="agent2",
                system_prompt="You are agent 2",
                max_loops=1,
            ),
        ]

        # Set random models
        model_names = ["gpt-4o-mini", "gpt-4o", "claude-3-5-sonnet"]
        result = set_random_models_for_agents(
            agents=agents, model_names=model_names
        )

        # Verify results
        assert len(result) == 2
        assert all(isinstance(agent, Agent) for agent in result)
        assert all(hasattr(agent, "model_name") for agent in result)
        assert all(
            agent.model_name in model_names for agent in result
        )

    def test_set_random_models_for_agents_with_single_agent(self):
        """Test set_random_models_for_agents with a single agent."""
        agent = Agent(
            agent_name="single_agent",
            system_prompt="You are helpful",
            max_loops=1,
        )

        model_names = ["gpt-4o-mini", "gpt-4o"]
        result = set_random_models_for_agents(
            agents=agent, model_names=model_names
        )

        assert isinstance(result, Agent)
        assert hasattr(result, "model_name")
        assert result.model_name in model_names

    def test_set_random_models_for_agents_with_none(self):
        """Test set_random_models_for_agents with None returns random model name."""
        model_names = ["gpt-4o-mini", "gpt-4o", "claude-3-5-sonnet"]
        result = set_random_models_for_agents(
            agents=None, model_names=model_names
        )

        assert isinstance(result, str)
        assert result in model_names

    @pytest.mark.skip(
        reason="This test requires API key and makes LLM calls"
    )
    def test_auto_swarm_builder_return_agents_objects_integration(
        self,
    ):
        """Integration test for AutoSwarmBuilder with execution_type='return-agents-objects'.

        This test requires OPENAI_API_KEY and makes actual LLM calls.
        Run manually with: pytest -k test_auto_swarm_builder_return_agents_objects_integration -v
        """
        builder = AutoSwarmBuilder(
            execution_type="return-agents-objects",
            model_name="gpt-4o-mini",
            max_loops=1,
            verbose=False,
        )

        agents = builder.run(
            "Create a team of 2 data analysis agents with specific roles"
        )

        # Verify agents were created
        assert isinstance(agents, list)
        assert len(agents) >= 1
        assert all(isinstance(agent, Agent) for agent in agents)
        assert all(hasattr(agent, "agent_name") for agent in agents)
        assert all(
            hasattr(agent, "agent_description") for agent in agents
        )

    def test_agent_spec_to_agent_all_fields(self):
        """Test that all AgentSpec fields are properly passed to Agent."""
        builder = AutoSwarmBuilder()

        agent_spec = AgentSpec(
            agent_name="full_test_agent",
            description="Full test description",
            system_prompt="You are a comprehensive test agent",
            model_name="gpt-4o-mini",
            auto_generate_prompt=False,
            max_tokens=4096,
            temperature=0.7,
            role="worker",
            max_loops=3,
            goal="Test all parameters",
        )

        agents = builder.create_agents_from_specs(
            {"agents": [agent_spec]}
        )

        assert len(agents) == 1
        agent = agents[0]

        # Verify all fields were set
        assert agent.agent_name == "full_test_agent"
        assert agent.agent_description == "Full test description"
        # Agent may modify system_prompt by adding additional instructions
        assert (
            "You are a comprehensive test agent"
            in agent.system_prompt
        )
        assert agent.max_loops == 3
        assert agent.max_tokens == 4096
        assert agent.temperature == 0.7

    def test_create_agents_from_specs_empty_list(self):
        """Test that create_agents_from_specs handles empty agent list."""
        builder = AutoSwarmBuilder()

        specs = {"agents": []}

        agents = builder.create_agents_from_specs(specs)

        assert isinstance(agents, list)
        assert len(agents) == 0


if __name__ == "__main__":
    run_all_tests()