tests for all swarming architectures

tests/models/mistral.py

@@ -2,6 +2,7 @@ import pytest
 from unittest.mock import patch, MagicMock
 from swarms.models.mistral import Mistral
 
+
 def test_mistral_initialization():
     mistral = Mistral(device="cpu")
     assert isinstance(mistral, Mistral)
@@ -14,28 +15,32 @@ def test_mistral_initialization():
     assert mistral.max_length == 100
     assert mistral.history == []
 
-@patch('your_module.AutoModelForCausalLM.from_pretrained')
-@patch('your_module.AutoTokenizer.from_pretrained')
+
+@patch("your_module.AutoModelForCausalLM.from_pretrained")
+@patch("your_module.AutoTokenizer.from_pretrained")
 def test_mistral_load_model(mock_tokenizer, mock_model):
     mistral = Mistral(device="cpu")
     mistral.load_model()
     mock_model.assert_called_once()
     mock_tokenizer.assert_called_once()
 
-@patch('your_module.Mistral.load_model')
+
+@patch("your_module.Mistral.load_model")
 def test_mistral_run(mock_load_model):
     mistral = Mistral(device="cpu")
     mistral.run("What's the weather in miami")
     mock_load_model.assert_called_once()
 
-@patch('your_module.Mistral.run')
+
+@patch("your_module.Mistral.run")
 def test_mistral_chat(mock_run):
     mistral = Mistral(device="cpu")
     mistral.chat("What's the weather in miami")
     mock_run.assert_called_once()
 
+
 def test_mistral__stream_response():
     mistral = Mistral(device="cpu")
     response = "It's sunny in Miami."
     tokens = list(mistral._stream_response(response))
-    assert tokens == ["It's", "sunny", "in", "Miami."]
+    assert tokens == ["It's", "sunny", "in", "Miami."]

tests/structs/nonlinear_workflow.py

@@ -2,6 +2,7 @@ import pytest
 from unittest.mock import patch, MagicMock
 from swarms.structs.nonlinear_workflow import NonLinearWorkflow, Task
 
+
 class MockTask(Task):
     def can_execute(self):
         return True
@@ -9,6 +10,7 @@ class MockTask(Task):
     def execute(self):
         return "Task executed"
 
+
 def test_nonlinearworkflow_initialization():
     agents = MagicMock()
     iters_per_task = MagicMock()
@@ -17,6 +19,7 @@ def test_nonlinearworkflow_initialization():
     assert workflow.agents == agents
     assert workflow.tasks == []
 
+
 def test_nonlinearworkflow_add():
     agents = MagicMock()
     iters_per_task = MagicMock()
@@ -25,8 +28,9 @@ def test_nonlinearworkflow_add():
     workflow.add(task)
     assert workflow.tasks == [task]
 
-@patch('your_module.NonLinearWorkflow.is_finished')
-@patch('your_module.NonLinearWorkflow.output_tasks')
+
+@patch("your_module.NonLinearWorkflow.is_finished")
+@patch("your_module.NonLinearWorkflow.output_tasks")
 def test_nonlinearworkflow_run(mock_output_tasks, mock_is_finished):
     agents = MagicMock()
     iters_per_task = MagicMock()
@@ -38,6 +42,7 @@ def test_nonlinearworkflow_run(mock_output_tasks, mock_is_finished):
     workflow.run()
     assert mock_output_tasks.called
 
+
 def test_nonlinearworkflow_output_tasks():
     agents = MagicMock()
     iters_per_task = MagicMock()
@@ -46,6 +51,7 @@ def test_nonlinearworkflow_output_tasks():
     workflow.add(task)
     assert workflow.output_tasks() == [task]
 
+
 def test_nonlinearworkflow_to_graph():
     agents = MagicMock()
     iters_per_task = MagicMock()
@@ -54,10 +60,11 @@ def test_nonlinearworkflow_to_graph():
     workflow.add(task)
     assert workflow.to_graph() == {"task1": set()}
 
+
 def test_nonlinearworkflow_order_tasks():
     agents = MagicMock()
     iters_per_task = MagicMock()
     workflow = NonLinearWorkflow(agents, iters_per_task)
     task = MockTask("task1")
     workflow.add(task)
-    assert workflow.order_tasks() == [task]
+    assert workflow.order_tasks() == [task]

tests/structs/workflow.py

@@ -2,6 +2,7 @@ import pytest
 from unittest.mock import patch, MagicMock
 from swarms.structs.workflow import Workflow
 
+
 def test_workflow_initialization():
     agent = MagicMock()
     workflow = Workflow(agent)
@@ -10,6 +11,7 @@ def test_workflow_initialization():
     assert workflow.tasks == []
     assert workflow.parallel == False
 
+
 def test_workflow_add():
     agent = MagicMock()
     workflow = Workflow(agent)
@@ -21,6 +23,7 @@ def test_workflow_add():
     assert task.output == None
     assert task.structure == workflow
 
+
 def test_workflow_first_task():
     agent = MagicMock()
     workflow = Workflow(agent)
@@ -28,6 +31,7 @@ def test_workflow_first_task():
     workflow.add("What's the weather in miami")
     assert workflow.first_task().task == "What's the weather in miami"
 
+
 def test_workflow_last_task():
     agent = MagicMock()
     workflow = Workflow(agent)
@@ -35,7 +39,8 @@ def test_workflow_last_task():
     workflow.add("What's the weather in miami")
     assert workflow.last_task().task == "What's the weather in miami"
 
-@patch('your_module.Workflow.__run_from_task')
+
+@patch("your_module.Workflow.__run_from_task")
 def test_workflow_run(mock_run_from_task):
     agent = MagicMock()
     workflow = Workflow(agent)
@@ -43,17 +48,23 @@ def test_workflow_run(mock_run_from_task):
     workflow.run()
     mock_run_from_task.assert_called_once()
 
+
 def test_workflow_context():
     agent = MagicMock()
     workflow = Workflow(agent)
     task = workflow.add("What's the weather in miami")
-    assert workflow.context(task) == {"parent_output": None, "parent": None, "child": None}
+    assert workflow.context(task) == {
+        "parent_output": None,
+        "parent": None,
+        "child": None,
+    }
+
 
-@patch('your_module.Workflow.Task.execute')
+@patch("your_module.Workflow.Task.execute")
 def test_workflow___run_from_task(mock_execute):
     agent = MagicMock()
     workflow = Workflow(agent)
     task = workflow.add("What's the weather in miami")
     mock_execute.return_value = "Sunny"
     workflow.__run_from_task(task)
-    mock_execute.assert_called_once()
+    mock_execute.assert_called_once()

tests/swarms/autoscaler.py

@@ -0,0 +1,62 @@
+import pytest
+from unittest.mock import patch, MagicMock
+from swarms.swarms.autoscaler import AutoScaler, Worker
+
+
+def test_autoscaler_initialization():
+    autoscaler = AutoScaler(
+        initial_agents=5,
+        scale_up_factor=2,
+        idle_threshold=0.1,
+        busy_threshold=0.8,
+        agent=Worker,
+    )
+    assert isinstance(autoscaler, AutoScaler)
+    assert autoscaler.scale_up_factor == 2
+    assert autoscaler.idle_threshold == 0.1
+    assert autoscaler.busy_threshold == 0.8
+    assert len(autoscaler.agents_pool) == 5
+
+
+def test_autoscaler_add_task():
+    autoscaler = AutoScaler(agent=Worker)
+    autoscaler.add_task("task1")
+    assert autoscaler.task_queue.qsize() == 1
+
+
+def test_autoscaler_scale_up():
+    autoscaler = AutoScaler(initial_agents=5, scale_up_factor=2, agent=Worker)
+    autoscaler.scale_up()
+    assert len(autoscaler.agents_pool) == 10
+
+
+def test_autoscaler_scale_down():
+    autoscaler = AutoScaler(initial_agents=5, agent=Worker)
+    autoscaler.scale_down()
+    assert len(autoscaler.agents_pool) == 4
+
+
+@patch("your_module.AutoScaler.scale_up")
+@patch("your_module.AutoScaler.scale_down")
+def test_autoscaler_monitor_and_scale(mock_scale_down, mock_scale_up):
+    autoscaler = AutoScaler(initial_agents=5, agent=Worker)
+    autoscaler.add_task("task1")
+    autoscaler.monitor_and_scale()
+    mock_scale_up.assert_called_once()
+    mock_scale_down.assert_called_once()
+
+
+@patch("your_module.AutoScaler.monitor_and_scale")
+@patch("your_module.Worker.run")
+def test_autoscaler_start(mock_run, mock_monitor_and_scale):
+    autoscaler = AutoScaler(initial_agents=5, agent=Worker)
+    autoscaler.add_task("task1")
+    autoscaler.start()
+    mock_run.assert_called_once()
+    mock_monitor_and_scale.assert_called_once()
+
+
+def test_autoscaler_del_agent():
+    autoscaler = AutoScaler(initial_agents=5, agent=Worker)
+    autoscaler.del_agent()
+    assert len(autoscaler.agents_pool) == 4

tests/swarms/dialogue_simulator.py

@@ -0,0 +1,20 @@
+import pytest
+from unittest.mock import patch, MagicMock
+from swarms.swarms.dialogue_simulator import DialogueSimulator, Worker
+
+def test_dialoguesimulator_initialization():
+    dialoguesimulator = DialogueSimulator(agents=[Worker]*5)
+    assert isinstance(dialoguesimulator, DialogueSimulator)
+    assert len(dialoguesimulator.agents) == 5
+
+@patch('swarms.workers.Worker.run')
+def test_dialoguesimulator_run(mock_run):
+    dialoguesimulator = DialogueSimulator(agents=[Worker]*5)
+    dialoguesimulator.run(max_iters=5, name="Agent 1", message="Hello, world!")
+    assert mock_run.call_count == 30
+
+@patch('swarms.workers.Worker.run')
+def test_dialoguesimulator_run_without_name_and_message(mock_run):
+    dialoguesimulator = DialogueSimulator(agents=[Worker]*5)
+    dialoguesimulator.run(max_iters=5)
+    assert mock_run.call_count == 25

tests/swarms/godmode.py

@@ -0,0 +1,26 @@
+import pytest
+from unittest.mock import patch, MagicMock
+from swarms.swarms.god_mode import GodMode, LLM
+
+def test_godmode_initialization():
+    godmode = GodMode(llms=[LLM]*5)
+    assert isinstance(godmode, GodMode)
+    assert len(godmode.llms) == 5
+
+def test_godmode_run(monkeypatch):
+    def mock_llm_run(self, task):
+        return "response"
+    monkeypatch.setattr(LLM, "run", mock_llm_run)
+    godmode = GodMode(llms=[LLM]*5)
+    responses = godmode.run("task1")
+    assert len(responses) == 5
+    assert responses == ["response", "response", "response", "response", "response"]
+
+@patch('builtins.print')
+def test_godmode_print_responses(mock_print, monkeypatch):
+    def mock_llm_run(self, task):
+        return "response"
+    monkeypatch.setattr(LLM, "run", mock_llm_run)
+    godmode = GodMode(llms=[LLM]*5)
+    godmode.print_responses("task1")
+    assert mock_print.call_count == 1

tests/swarms/multi_agent_collab.py

@@ -0,0 +1,52 @@
+import pytest
+from unittest.mock import patch, MagicMock
+from swarms.swarms.multi_agent_collab import MultiAgentCollaboration, Worker, select_next_speaker
+
+def test_multiagentcollaboration_initialization():
+    multiagentcollaboration = MultiAgentCollaboration(agents=[Worker]*5, selection_function=select_next_speaker)
+    assert isinstance(multiagentcollaboration, MultiAgentCollaboration)
+    assert len(multiagentcollaboration.agents) == 5
+    assert multiagentcollaboration._step == 0
+
+@patch('swarms.workers.Worker.reset')
+def test_multiagentcollaboration_reset(mock_reset):
+    multiagentcollaboration = MultiAgentCollaboration(agents=[Worker]*5, selection_function=select_next_speaker)
+    multiagentcollaboration.reset()
+    assert mock_reset.call_count == 5
+
+@patch('swarms.workers.Worker.run')
+def test_multiagentcollaboration_inject(mock_run):
+    multiagentcollaboration = MultiAgentCollaboration(agents=[Worker]*5, selection_function=select_next_speaker)
+    multiagentcollaboration.inject("Agent 1", "Hello, world!")
+    assert multiagentcollaboration._step == 1
+    assert mock_run.call_count == 5
+
+@patch('swarms.workers.Worker.send')
+@patch('swarms.workers.Worker.receive')
+def test_multiagentcollaboration_step(mock_receive, mock_send):
+    multiagentcollaboration = MultiAgentCollaboration(agents=[Worker]*5, selection_function=select_next_speaker)
+    result = multiagentcollaboration.step()
+    assert multiagentcollaboration._step == 1
+    assert mock_send.call_count == 5
+    assert mock_receive.call_count == 25
+
+@patch('swarms.workers.Worker.bid')
+def test_multiagentcollaboration_ask_for_bid(mock_bid):
+    multiagentcollaboration = MultiAgentCollaboration(agents=[Worker]*5, selection_function=select_next_speaker)
+    result = multiagentcollaboration.ask_for_bid(Worker)
+    assert isinstance(result, int)
+
+@patch('swarms.workers.Worker.bid')
+def test_multiagentcollaboration_select_next_speaker(mock_bid):
+    multiagentcollaboration = MultiAgentCollaboration(agents=[Worker]*5, selection_function=select_next_speaker)
+    result = multiagentcollaboration.select_next_speaker(1, [Worker]*5)
+    assert isinstance(result, int)
+
+@patch('swarms.workers.Worker.send')
+@patch('swarms.workers.Worker.receive')
+def test_multiagentcollaboration_run(mock_receive, mock_send):
+    multiagentcollaboration = MultiAgentCollaboration(agents=[Worker]*5, selection_function=select_next_speaker)
+    multiagentcollaboration.run(max_iters=5)
+    assert multiagentcollaboration._step == 6
+    assert mock_send.call_count == 30
+    assert mock_receive.call_count == 150

tests/swarms/multi_agent_debate.py

@@ -0,0 +1,38 @@
+import pytest
+from unittest.mock import patch, MagicMock
+from swarms.swarms.multi_agent_debate import MultiAgentDebate, Worker, select_speaker
+
+def test_multiagentdebate_initialization():
+    multiagentdebate = MultiAgentDebate(agents=[Worker]*5, selection_func=select_speaker)
+    assert isinstance(multiagentdebate, MultiAgentDebate)
+    assert len(multiagentdebate.agents) == 5
+    assert multiagentdebate.selection_func == select_speaker
+
+@patch('swarms.workers.Worker.reset')
+def test_multiagentdebate_reset_agents(mock_reset):
+    multiagentdebate = MultiAgentDebate(agents=[Worker]*5, selection_func=select_speaker)
+    multiagentdebate.reset_agents()
+    assert mock_reset.call_count == 5
+
+def test_multiagentdebate_inject_agent():
+    multiagentdebate = MultiAgentDebate(agents=[Worker]*5, selection_func=select_speaker)
+    multiagentdebate.inject_agent(Worker)
+    assert len(multiagentdebate.agents) == 6
+
+@patch('swarms.workers.Worker.run')
+def test_multiagentdebate_run(mock_run):
+    multiagentdebate = MultiAgentDebate(agents=[Worker]*5, selection_func=select_speaker)
+    results = multiagentdebate.run("Write a short story.")
+    assert len(results) == 5
+    assert mock_run.call_count == 5
+
+def test_multiagentdebate_update_task():
+    multiagentdebate = MultiAgentDebate(agents=[Worker]*5, selection_func=select_speaker)
+    multiagentdebate.update_task("Write a short story.")
+    assert multiagentdebate.task == "Write a short story."
+
+def test_multiagentdebate_format_results():
+    multiagentdebate = MultiAgentDebate(agents=[Worker]*5, selection_func=select_speaker)
+    results = [{"agent": "Agent 1", "response": "Hello, world!"}, {"agent": "Agent 2", "response": "Goodbye, world!"}]
+    formatted_results = multiagentdebate.format_results(results)
+    assert formatted_results == "Agent Agent 1 responded: Hello, world!\nAgent Agent 2 responded: Goodbye, world!"

tests/swarms/orchestrate.py

@@ -0,0 +1,62 @@
+import numpy as np
+import pytest
+from unittest.mock import patch, MagicMock
+from swarms.swarms.orchestrate import Orchestrator, Worker
+import chromadb
+
+def test_orchestrator_initialization():
+    orchestrator = Orchestrator(agent=Worker, agent_list=[Worker]*5, task_queue=[])
+    assert isinstance(orchestrator, Orchestrator)
+    assert orchestrator.agents.qsize() == 5
+    assert orchestrator.task_queue.qsize() == 0
+
+def test_orchestrator_assign_task():
+    orchestrator = Orchestrator(agent=Worker, agent_list=[Worker]*5, task_queue=[])
+    orchestrator.assign_task(1, {"content": "task1"})
+    assert orchestrator.task_queue.qsize() == 1
+
+def test_orchestrator_embed():
+    orchestrator = Orchestrator(agent=Worker, agent_list=[Worker]*5, task_queue=[])
+    result = orchestrator.embed("Hello, world!", "api_key", "model_name")
+    assert isinstance(result, np.ndarray)
+
+def test_orchestrator_retrieve_results():
+    orchestrator = Orchestrator(agent=Worker, agent_list=[Worker]*5, task_queue=[])
+    result = orchestrator.retrieve_results(1)
+    assert isinstance(result, list)
+
+def test_orchestrator_update_vector_db():
+    orchestrator = Orchestrator(agent=Worker, agent_list=[Worker]*5, task_queue=[])
+    data = {"vector": np.array([1, 2, 3]), "task_id": 1}
+    orchestrator.update_vector_db(data)
+    assert orchestrator.collection.count() == 1
+
+def test_orchestrator_get_vector_db():
+    orchestrator = Orchestrator(agent=Worker, agent_list=[Worker]*5, task_queue=[])
+    result = orchestrator.get_vector_db()
+    assert isinstance(result, chromadb.Collection)
+
+def test_orchestrator_append_to_db():
+    orchestrator = Orchestrator(agent=Worker, agent_list=[Worker]*5, task_queue=[])
+    orchestrator.append_to_db("Hello, world!")
+    assert orchestrator.collection.count() == 1
+
+def test_orchestrator_run():
+    orchestrator = Orchestrator(agent=Worker, agent_list=[Worker]*5, task_queue=[])
+    result = orchestrator.run("Write a short story.")
+    assert isinstance(result, list)
+
+def test_orchestrator_chat():
+    orchestrator = Orchestrator(agent=Worker, agent_list=[Worker]*5, task_queue=[])
+    orchestrator.chat(1, 2, "Hello, Agent 2!")
+    assert orchestrator.collection.count() == 1
+
+def test_orchestrator_add_agents():
+    orchestrator = Orchestrator(agent=Worker, agent_list=[Worker]*5, task_queue=[])
+    orchestrator.add_agents(5)
+    assert orchestrator.agents.qsize() == 10
+
+def test_orchestrator_remove_agents():
+    orchestrator = Orchestrator(agent=Worker, agent_list=[Worker]*5, task_queue=[])
+    orchestrator.remove_agents(3)
+    assert orchestrator.agents.qsize() == 2

tests/swarms/simple_swarm.py

@@ -0,0 +1,37 @@
+import pytest
+from unittest.mock import patch, MagicMock
+from swarms.swarms.simple_swarm import SimpleSwarm, Worker
+
+def test_simpleswarm_initialization():
+    simpleswarm = SimpleSwarm(num_workers=5, openai_api_key="api_key", ai_name="ai_name")
+    assert isinstance(simpleswarm, SimpleSwarm)
+    assert len(simpleswarm.workers) == 5
+    assert simpleswarm.task_queue.qsize() == 0
+    assert simpleswarm.priority_queue.qsize() == 0
+
+def test_simpleswarm_distribute():
+    simpleswarm = SimpleSwarm(num_workers=5, openai_api_key="api_key", ai_name="ai_name")
+    simpleswarm.distribute("task1")
+    assert simpleswarm.task_queue.qsize() == 1
+    simpleswarm.distribute("task2", priority=1)
+    assert simpleswarm.priority_queue.qsize() == 1
+
+@patch('swarms.workers.Worker.run')
+def test_simpleswarm_process_task(mock_run):
+    simpleswarm = SimpleSwarm(num_workers=5, openai_api_key="api_key", ai_name="ai_name")
+    result = simpleswarm._process_task("task1")
+    assert mock_run.call_count == 5
+
+def test_simpleswarm_run():
+    simpleswarm = SimpleSwarm(num_workers=5, openai_api_key="api_key", ai_name="ai_name")
+    simpleswarm.distribute("task1")
+    simpleswarm.distribute("task2", priority=1)
+    results = simpleswarm.run()
+    assert len(results) == 2
+
+@patch('swarms.workers.Worker.run')
+def test_simpleswarm_run_old(mock_run):
+    simpleswarm = SimpleSwarm(num_workers=5, openai_api_key="api_key", ai_name="ai_name")
+    results = simpleswarm.run_old("task1")
+    assert len(results) == 5
+    assert mock_run.call_count == 5