swarms/tests/structs/test_agent_router.py

"""
Simplified test suite for AgentRouter class using pytest.

This module contains focused tests for the core functionality of the AgentRouter class.
"""

import pytest
from unittest.mock import Mock, patch

from swarms.structs.agent_router import AgentRouter
from swarms.structs.agent import Agent


@pytest.fixture
def test_agent():
    """Create a real agent for testing."""
    with patch("swarms.structs.agent.LiteLLM") as mock_llm:
        mock_llm.return_value.run.return_value = "Test response"
        return Agent(
            agent_name="test_agent",
            agent_description="A test agent",
            system_prompt="You are a test agent",
            model_name="gpt-4o-mini",
            max_loops=1,
            verbose=False,
            print_on=False,
        )


def test_agent_router_initialization_default():
    """Test AgentRouter initialization with default parameters."""
    with patch("swarms.structs.agent_router.embedding"):
        router = AgentRouter()

        assert router.embedding_model == "text-embedding-ada-002"
        assert router.n_agents == 1
        assert router.api_key is None
        assert router.api_base is None
        assert router.agents == []
        assert router.agent_embeddings == []
        assert router.agent_metadata == []


def test_agent_router_initialization_custom():
    """Test AgentRouter initialization with custom parameters."""
    with patch("swarms.structs.agent_router.embedding"), patch(
        "swarms.structs.agent.LiteLLM"
    ) as mock_llm:
        mock_llm.return_value.run.return_value = "Test response"
        agents = [
            Agent(
                agent_name="test1",
                model_name="gpt-4o-mini",
                max_loops=1,
                verbose=False,
                print_on=False,
            ),
            Agent(
                agent_name="test2",
                model_name="gpt-4o-mini",
                max_loops=1,
                verbose=False,
                print_on=False,
            ),
        ]
        router = AgentRouter(
            embedding_model="custom-model",
            n_agents=3,
            api_key="custom_key",
            api_base="custom_base",
            agents=agents,
        )

        assert router.embedding_model == "custom-model"
        assert router.n_agents == 3
        assert router.api_key == "custom_key"
        assert router.api_base == "custom_base"
        assert len(router.agents) == 2


def test_cosine_similarity_identical_vectors():
    """Test cosine similarity with identical vectors."""
    router = AgentRouter()
    vec1 = [1.0, 0.0, 0.0]
    vec2 = [1.0, 0.0, 0.0]

    result = router._cosine_similarity(vec1, vec2)
    assert result == 1.0


def test_cosine_similarity_orthogonal_vectors():
    """Test cosine similarity with orthogonal vectors."""
    router = AgentRouter()
    vec1 = [1.0, 0.0, 0.0]
    vec2 = [0.0, 1.0, 0.0]

    result = router._cosine_similarity(vec1, vec2)
    assert result == 0.0


def test_cosine_similarity_opposite_vectors():
    """Test cosine similarity with opposite vectors."""
    router = AgentRouter()
    vec1 = [1.0, 0.0, 0.0]
    vec2 = [-1.0, 0.0, 0.0]

    result = router._cosine_similarity(vec1, vec2)
    assert result == -1.0


def test_cosine_similarity_different_lengths():
    """Test cosine similarity with vectors of different lengths."""
    router = AgentRouter()
    vec1 = [1.0, 0.0]
    vec2 = [1.0, 0.0, 0.0]

    with pytest.raises(
        ValueError, match="Vectors must have the same length"
    ):
        router._cosine_similarity(vec1, vec2)


@patch("swarms.structs.agent_router.embedding")
def test_generate_embedding_success(mock_embedding):
    """Test successful embedding generation."""
    mock_embedding.return_value.data = [
        Mock(embedding=[0.1, 0.2, 0.3, 0.4])
    ]

    router = AgentRouter()
    result = router._generate_embedding("test text")

    assert result == [0.1, 0.2, 0.3, 0.4]
    mock_embedding.assert_called_once()


@patch("swarms.structs.agent_router.embedding")
def test_generate_embedding_error(mock_embedding):
    """Test embedding generation error handling."""
    mock_embedding.side_effect = Exception("API Error")

    router = AgentRouter()

    with pytest.raises(Exception, match="API Error"):
        router._generate_embedding("test text")


@patch("swarms.structs.agent_router.embedding")
def test_add_agent_success(mock_embedding, test_agent):
    """Test successful agent addition."""
    mock_embedding.return_value.data = [
        Mock(embedding=[0.1, 0.2, 0.3])
    ]

    router = AgentRouter()
    router.add_agent(test_agent)

    assert len(router.agents) == 1
    assert len(router.agent_embeddings) == 1
    assert len(router.agent_metadata) == 1
    assert router.agents[0] == test_agent
    assert router.agent_embeddings[0] == [0.1, 0.2, 0.3]
    assert router.agent_metadata[0]["name"] == "test_agent"


@patch("swarms.structs.agent_router.embedding")
def test_add_agent_retry_error(mock_embedding, test_agent):
    """Test agent addition with retry mechanism failure."""
    mock_embedding.side_effect = Exception("Embedding error")

    router = AgentRouter()

    # Should raise RetryError after retries are exhausted
    with pytest.raises(Exception) as exc_info:
        router.add_agent(test_agent)

    # Check that it's a retry error or contains the original error
    assert "Embedding error" in str(
        exc_info.value
    ) or "RetryError" in str(exc_info.value)


@patch("swarms.structs.agent_router.embedding")
def test_add_agents_multiple(mock_embedding):
    """Test adding multiple agents."""
    mock_embedding.return_value.data = [
        Mock(embedding=[0.1, 0.2, 0.3])
    ]

    with patch("swarms.structs.agent.LiteLLM") as mock_llm:
        mock_llm.return_value.run.return_value = "Test response"
        router = AgentRouter()
        agents = [
            Agent(
                agent_name="agent1",
                model_name="gpt-4o-mini",
                max_loops=1,
                verbose=False,
                print_on=False,
            ),
            Agent(
                agent_name="agent2",
                model_name="gpt-4o-mini",
                max_loops=1,
                verbose=False,
                print_on=False,
            ),
            Agent(
                agent_name="agent3",
                model_name="gpt-4o-mini",
                max_loops=1,
                verbose=False,
                print_on=False,
            ),
        ]

        router.add_agents(agents)

        assert len(router.agents) == 3
        assert len(router.agent_embeddings) == 3
        assert len(router.agent_metadata) == 3


@patch("swarms.structs.agent_router.embedding")
def test_find_best_agent_success(mock_embedding):
    """Test successful best agent finding."""
    # Mock embeddings for agents and task
    mock_embedding.side_effect = [
        Mock(data=[Mock(embedding=[0.1, 0.2, 0.3])]),  # agent1
        Mock(data=[Mock(embedding=[0.4, 0.5, 0.6])]),  # agent2
        Mock(data=[Mock(embedding=[0.7, 0.8, 0.9])]),  # task
    ]

    with patch("swarms.structs.agent.LiteLLM") as mock_llm:
        mock_llm.return_value.run.return_value = "Test response"
        router = AgentRouter()
        agent1 = Agent(
            agent_name="agent1",
            agent_description="First agent",
            system_prompt="Prompt 1",
            model_name="gpt-4o-mini",
            max_loops=1,
            verbose=False,
            print_on=False,
        )
        agent2 = Agent(
            agent_name="agent2",
            agent_description="Second agent",
            system_prompt="Prompt 2",
            model_name="gpt-4o-mini",
            max_loops=1,
            verbose=False,
            print_on=False,
        )

        router.add_agent(agent1)
        router.add_agent(agent2)

        # Mock the similarity calculation to return predictable results
        with patch.object(
            router, "_cosine_similarity"
        ) as mock_similarity:
            mock_similarity.side_effect = [
                0.8,
                0.6,
            ]  # agent1 more similar

            result = router.find_best_agent("test task")

            assert result == agent1


def test_find_best_agent_no_agents():
    """Test finding best agent when no agents are available."""
    with patch("swarms.structs.agent_router.embedding"):
        router = AgentRouter()

        result = router.find_best_agent("test task")

        assert result is None


@patch("swarms.structs.agent_router.embedding")
def test_find_best_agent_retry_error(mock_embedding):
    """Test error handling in find_best_agent with retry mechanism."""
    mock_embedding.side_effect = Exception("API Error")

    with patch("swarms.structs.agent.LiteLLM") as mock_llm:
        mock_llm.return_value.run.return_value = "Test response"
        router = AgentRouter()
        router.agents = [
            Agent(
                agent_name="agent1",
                model_name="gpt-4o-mini",
                max_loops=1,
                verbose=False,
                print_on=False,
            )
        ]
        router.agent_embeddings = [[0.1, 0.2, 0.3]]

        # Should raise RetryError after retries are exhausted
        with pytest.raises(Exception) as exc_info:
            router.find_best_agent("test task")

        # Check that it's a retry error or contains the original error
        assert "API Error" in str(
            exc_info.value
        ) or "RetryError" in str(exc_info.value)


@patch("swarms.structs.agent_router.embedding")
def test_update_agent_history_success(mock_embedding, test_agent):
    """Test successful agent history update."""
    mock_embedding.return_value.data = [
        Mock(embedding=[0.1, 0.2, 0.3])
    ]

    router = AgentRouter()
    router.add_agent(test_agent)

    # Update agent history
    router.update_agent_history("test_agent")

    # Verify the embedding was regenerated
    assert (
        mock_embedding.call_count == 2
    )  # Once for add, once for update


def test_update_agent_history_agent_not_found():
    """Test updating history for non-existent agent."""
    with patch(
        "swarms.structs.agent_router.embedding"
    ) as mock_embedding:
        mock_embedding.return_value.data = [
            Mock(embedding=[0.1, 0.2, 0.3])
        ]
        router = AgentRouter()

        # Should not raise an exception, just log a warning
        router.update_agent_history("non_existent_agent")


@patch("swarms.structs.agent_router.embedding")
def test_agent_metadata_structure(mock_embedding, test_agent):
    """Test the structure of agent metadata."""
    mock_embedding.return_value.data = [
        Mock(embedding=[0.1, 0.2, 0.3])
    ]

    router = AgentRouter()
    router.add_agent(test_agent)

    metadata = router.agent_metadata[0]
    assert "name" in metadata
    assert "text" in metadata
    assert metadata["name"] == "test_agent"
    assert (
        "test_agent A test agent You are a test agent"
        in metadata["text"]
    )


def test_agent_router_edge_cases():
    """Test various edge cases."""
    with patch(
        "swarms.structs.agent_router.embedding"
    ) as mock_embedding:
        mock_embedding.return_value.data = [
            Mock(embedding=[0.1, 0.2, 0.3])
        ]

        router = AgentRouter()

        # Test with empty string task
        result = router.find_best_agent("")
        assert result is None

        # Test with very long task description
        long_task = "test " * 1000
        result = router.find_best_agent(long_task)
        assert result is None


if __name__ == "__main__":
    pytest.main([__file__])