swarms/tests/structs/test_collaborative_utils.py at 6508d3e53b5d078a60bcffe4f4156b69e1175976 · aparekh02/swarms · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
import pytest
from swarms.structs.collaborative_utils import talk_to_agent
from swarms.structs.agent import Agent


def create_test_agent(name: str, description: str = "Test agent") -> Agent:
    """Create a real Agent instance for testing"""
    return Agent(
        agent_name=name,
        agent_description=description,
        system_prompt=f"You are {name}, a helpful test assistant.",
        model_name="gpt-4o-mini",
        max_loops=1,
        verbose=False,
    )


def test_talk_to_agent_success():
    """Test successful agent-to-agent communication"""
    current_agent = create_test_agent("CurrentAgent")
    target_agent = create_test_agent("TargetAgent")

    agents = [current_agent, target_agent]

    result = talk_to_agent(
        current_agent=current_agent,
        agents=agents,
        task="What is 2+2?",
        agent_name="TargetAgent",
        max_loops=1
    )

    assert result is not None
    # Result should be a list or string from the debate
    assert len(str(result)) > 0


def test_talk_to_agent_not_found():
    """Test error when target agent not found"""
    current_agent = create_test_agent("CurrentAgent")

    agents = [current_agent]

    with pytest.raises(ValueError, match="Agent 'NonExistent' not found"):
        talk_to_agent(
            current_agent=current_agent,
            agents=agents,
            task="Test task",
            agent_name="NonExistent"
        )


def test_talk_to_agent_with_max_loops():
    """Test talk_to_agent with custom max_loops"""
    current_agent = create_test_agent("CurrentAgent")
    target_agent = create_test_agent("TargetAgent")

    agents = [current_agent, target_agent]

    result = talk_to_agent(
        current_agent=current_agent,
        agents=agents,
        task="Discuss the weather briefly",
        agent_name="TargetAgent",
        max_loops=2
    )

    assert result is not None


def test_talk_to_agent_no_agent_name_attribute():
    """Test when target agent is not found by name"""
    current_agent = create_test_agent("CurrentAgent")
    # Create another agent with different name
    other_agent = create_test_agent("OtherAgent")

    agents = [current_agent, other_agent]

    with pytest.raises(ValueError, match="Agent 'TargetAgent' not found"):
        talk_to_agent(
            current_agent=current_agent,
            agents=agents,
            task="Test",
            agent_name="TargetAgent"
        )


def test_talk_to_agent_output_type():
    """Test talk_to_agent with custom output_type"""
    current_agent = create_test_agent("CurrentAgent")
    target_agent = create_test_agent("TargetAgent")

    agents = [current_agent, target_agent]

    result = talk_to_agent(
        current_agent=current_agent,
        agents=agents,
        task="Say hello",
        agent_name="TargetAgent",
        output_type="str",
        max_loops=1
    )

    assert result is not None