From b5df48afd7d0dd4de72ea00cf93694bdbe83a8cc Mon Sep 17 00:00:00 2001 From: Kye Date: Tue, 10 Oct 2023 16:57:55 -0400 Subject: [PATCH] examples fix Former-commit-id: 454ee70d2266a840a0fe2ef934d69609f33a0475 --- multi_agent_debate.py | 2 +- playground/swarms/dialogue_simulator.py | 32 +++++++++++-- playground/swarms/godmode.py | 35 +++++++++++--- playground/swarms/group_chat.py | 1 - tests/swarms/dialogue_simulator.py | 15 +++--- tests/swarms/godmode.py | 15 ++++-- tests/swarms/multi_agent_collab.py | 61 +++++++++++++++++-------- tests/swarms/multi_agent_debate.py | 44 ++++++++++++++---- tests/swarms/orchestrate.py | 35 +++++++++----- tests/swarms/scalable_groupchat.py | 46 +++++++++++++------ tests/swarms/simple_swarm.py | 31 +++++++++---- 11 files changed, 233 insertions(+), 84 deletions(-) delete mode 100644 playground/swarms/group_chat.py diff --git a/multi_agent_debate.py b/multi_agent_debate.py index 067d3e89..922c0599 100644 --- a/multi_agent_debate.py +++ b/multi_agent_debate.py @@ -29,7 +29,7 @@ worker3 = Worker( temperature=0.5, ) -#init agents +# init agents agents = [worker1, worker2, worker3] # Initialize multi-agent debate with the selection function diff --git a/playground/swarms/dialogue_simulator.py b/playground/swarms/dialogue_simulator.py index a0800c49..755c502e 100644 --- a/playground/swarms/dialogue_simulator.py +++ b/playground/swarms/dialogue_simulator.py @@ -1,7 +1,33 @@ -from swarms import DialogueSimulator, Worker +from swarms.swarms import DialogueSimulator +from swarms.workers.worker import Worker +from langchain.llms import OpenAIChat -worker1 = Worker(ai_name="Plinus", openai_api_key="") -worker2 = Worker(ai_name="Optimus Prime", openai_api_key="") +llm = OpenAIChat(model_name="gpt-4", openai_api_key="api-key", temperature=0.5) + +worker1 = Worker( + llm=llm, + ai_name="Bumble Bee", + ai_role="Worker in a swarm", + external_tools=None, + human_in_the_loop=False, + temperature=0.5, +) +worker2 = Worker( + llm=llm, + ai_name="Optimus Prime", + ai_role="Worker in a swarm", + external_tools=None, + human_in_the_loop=False, + temperature=0.5, +) +worker3 = Worker( + llm=llm, + ai_name="Megatron", + ai_role="Worker in a swarm", + external_tools=None, + human_in_the_loop=False, + temperature=0.5, +) collab = DialogueSimulator( [worker1, worker2], diff --git a/playground/swarms/godmode.py b/playground/swarms/godmode.py index 1031d94e..f0adb2f2 100644 --- a/playground/swarms/godmode.py +++ b/playground/swarms/godmode.py @@ -1,14 +1,37 @@ -from langchain.models import Anthropic, GooglePalm, OpenAIChat +from langchain.models import OpenAIChat from swarms.swarms import GodMode +from swarms.workers.worker import Worker -claude = Anthropic(anthropic_api_key="") -palm = GooglePalm(google_api_key="") -gpt = OpenAIChat(openai_api_key="") +llm = OpenAIChat(model_name="gpt-4", openai_api_key="api-key", temperature=0.5) +worker1 = Worker( + llm=llm, + ai_name="Bumble Bee", + ai_role="Worker in a swarm", + external_tools=None, + human_in_the_loop=False, + temperature=0.5, +) +worker2 = Worker( + llm=llm, + ai_name="Optimus Prime", + ai_role="Worker in a swarm", + external_tools=None, + human_in_the_loop=False, + temperature=0.5, +) +worker3 = Worker( + llm=llm, + ai_name="Megatron", + ai_role="Worker in a swarm", + external_tools=None, + human_in_the_loop=False, + temperature=0.5, +) # Usage -llms = [claude, palm, gpt] +agents = [worker1, worker2, worker3] -god_mode = GodMode(llms) +god_mode = GodMode(agents) task = "What are the biggest risks facing humanity?" diff --git a/playground/swarms/group_chat.py b/playground/swarms/group_chat.py deleted file mode 100644 index a78ad0e5..00000000 --- a/playground/swarms/group_chat.py +++ /dev/null @@ -1 +0,0 @@ -from swarms.swarms import GroupChat diff --git a/tests/swarms/dialogue_simulator.py b/tests/swarms/dialogue_simulator.py index 634c29d1..ffe1cdb1 100644 --- a/tests/swarms/dialogue_simulator.py +++ b/tests/swarms/dialogue_simulator.py @@ -2,19 +2,22 @@ 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) + dialoguesimulator = DialogueSimulator(agents=[Worker] * 5) assert isinstance(dialoguesimulator, DialogueSimulator) assert len(dialoguesimulator.agents) == 5 -@patch('swarms.workers.worker.Worker.run') + +@patch("swarms.workers.worker.Worker.run") def test_dialoguesimulator_run(mock_run): - dialoguesimulator = DialogueSimulator(agents=[Worker]*5) + 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.Worker.run') + +@patch("swarms.workers.worker.Worker.run") def test_dialoguesimulator_run_without_name_and_message(mock_run): - dialoguesimulator = DialogueSimulator(agents=[Worker]*5) + dialoguesimulator = DialogueSimulator(agents=[Worker] * 5) dialoguesimulator.run(max_iters=5) - assert mock_run.call_count == 25 \ No newline at end of file + assert mock_run.call_count == 25 diff --git a/tests/swarms/godmode.py b/tests/swarms/godmode.py index 44e97b25..1894333b 100644 --- a/tests/swarms/godmode.py +++ b/tests/swarms/godmode.py @@ -2,25 +2,30 @@ 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) + 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) + godmode = GodMode(llms=[LLM] * 5) responses = godmode.run("task1") assert len(responses) == 5 assert responses == ["response", "response", "response", "response", "response"] -@patch('builtins.print') + +@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 = GodMode(llms=[LLM] * 5) godmode.print_responses("task1") - assert mock_print.call_count == 1 \ No newline at end of file + assert mock_print.call_count == 1 diff --git a/tests/swarms/multi_agent_collab.py b/tests/swarms/multi_agent_collab.py index 62b04dbf..b4464606 100644 --- a/tests/swarms/multi_agent_collab.py +++ b/tests/swarms/multi_agent_collab.py @@ -1,52 +1,77 @@ import pytest from unittest.mock import patch, MagicMock -from swarms.swarms.multi_agent_collab import MultiAgentCollaboration, Worker, select_next_speaker +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) + 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') + +@patch("swarms.workers.Worker.reset") def test_multiagentcollaboration_reset(mock_reset): - multiagentcollaboration = MultiAgentCollaboration(agents=[Worker]*5, selection_function=select_next_speaker) + multiagentcollaboration = MultiAgentCollaboration( + agents=[Worker] * 5, selection_function=select_next_speaker + ) multiagentcollaboration.reset() assert mock_reset.call_count == 5 -@patch('swarms.workers.Worker.run') + +@patch("swarms.workers.Worker.run") def test_multiagentcollaboration_inject(mock_run): - multiagentcollaboration = MultiAgentCollaboration(agents=[Worker]*5, selection_function=select_next_speaker) + 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') + +@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) + 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') + +@patch("swarms.workers.Worker.bid") def test_multiagentcollaboration_ask_for_bid(mock_bid): - multiagentcollaboration = MultiAgentCollaboration(agents=[Worker]*5, selection_function=select_next_speaker) + 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') + +@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) + 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') + +@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 = 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 \ No newline at end of file + assert mock_receive.call_count == 150 diff --git a/tests/swarms/multi_agent_debate.py b/tests/swarms/multi_agent_debate.py index f01e6a72..f5fd6199 100644 --- a/tests/swarms/multi_agent_debate.py +++ b/tests/swarms/multi_agent_debate.py @@ -2,37 +2,61 @@ 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) + 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') + +@patch("swarms.workers.Worker.reset") def test_multiagentdebate_reset_agents(mock_reset): - multiagentdebate = MultiAgentDebate(agents=[Worker]*5, selection_func=select_speaker) + 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 = MultiAgentDebate( + agents=[Worker] * 5, selection_func=select_speaker + ) multiagentdebate.inject_agent(Worker) assert len(multiagentdebate.agents) == 6 -@patch('swarms.workers.Worker.run') + +@patch("swarms.workers.Worker.run") def test_multiagentdebate_run(mock_run): - multiagentdebate = MultiAgentDebate(agents=[Worker]*5, selection_func=select_speaker) + 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 = 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!"}] + 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!" \ No newline at end of file + assert ( + formatted_results + == "Agent Agent 1 responded: Hello, world!\nAgent Agent 2 responded: Goodbye, world!" + ) diff --git a/tests/swarms/orchestrate.py b/tests/swarms/orchestrate.py index af33d479..7b5880aa 100644 --- a/tests/swarms/orchestrate.py +++ b/tests/swarms/orchestrate.py @@ -4,59 +4,70 @@ 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=[]) + 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 = 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=[]) + 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=[]) + 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=[]) + 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=[]) + 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 = 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=[]) + 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 = 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 = 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 = Orchestrator(agent=Worker, agent_list=[Worker] * 5, task_queue=[]) orchestrator.remove_agents(3) - assert orchestrator.agents.qsize() == 2 \ No newline at end of file + assert orchestrator.agents.qsize() == 2 diff --git a/tests/swarms/scalable_groupchat.py b/tests/swarms/scalable_groupchat.py index 583514a7..050a334a 100644 --- a/tests/swarms/scalable_groupchat.py +++ b/tests/swarms/scalable_groupchat.py @@ -2,43 +2,61 @@ import pytest from unittest.mock import patch, MagicMock from swarms.swarms.scalable_groupchat import ScalableGroupChat, Worker + def test_scalablegroupchat_initialization(): - scalablegroupchat = ScalableGroupChat(worker_count=5, collection_name="swarm", api_key="api_key") + scalablegroupchat = ScalableGroupChat( + worker_count=5, collection_name="swarm", api_key="api_key" + ) assert isinstance(scalablegroupchat, ScalableGroupChat) assert len(scalablegroupchat.workers) == 5 assert scalablegroupchat.collection_name == "swarm" assert scalablegroupchat.api_key == "api_key" -@patch('chromadb.utils.embedding_functions.OpenAIEmbeddingFunction') + +@patch("chromadb.utils.embedding_functions.OpenAIEmbeddingFunction") def test_scalablegroupchat_embed(mock_openaiembeddingfunction): - scalablegroupchat = ScalableGroupChat(worker_count=5, collection_name="swarm", api_key="api_key") + scalablegroupchat = ScalableGroupChat( + worker_count=5, collection_name="swarm", api_key="api_key" + ) result = scalablegroupchat.embed("input", "model_name") assert mock_openaiembeddingfunction.call_count == 1 -@patch('swarms.swarms.scalable_groupchat.ScalableGroupChat.collection.query') + +@patch("swarms.swarms.scalable_groupchat.ScalableGroupChat.collection.query") def test_scalablegroupchat_retrieve_results(mock_query): - scalablegroupchat = ScalableGroupChat(worker_count=5, collection_name="swarm", api_key="api_key") + scalablegroupchat = ScalableGroupChat( + worker_count=5, collection_name="swarm", api_key="api_key" + ) result = scalablegroupchat.retrieve_results(1) assert mock_query.call_count == 1 -@patch('swarms.swarms.scalable_groupchat.ScalableGroupChat.collection.add') + +@patch("swarms.swarms.scalable_groupchat.ScalableGroupChat.collection.add") def test_scalablegroupchat_update_vector_db(mock_add): - scalablegroupchat = ScalableGroupChat(worker_count=5, collection_name="swarm", api_key="api_key") + scalablegroupchat = ScalableGroupChat( + worker_count=5, collection_name="swarm", api_key="api_key" + ) scalablegroupchat.update_vector_db({"vector": "vector", "task_id": "task_id"}) assert mock_add.call_count == 1 -@patch('swarms.swarms.scalable_groupchat.ScalableGroupChat.collection.add') + +@patch("swarms.swarms.scalable_groupchat.ScalableGroupChat.collection.add") def test_scalablegroupchat_append_to_db(mock_add): - scalablegroupchat = ScalableGroupChat(worker_count=5, collection_name="swarm", api_key="api_key") + scalablegroupchat = ScalableGroupChat( + worker_count=5, collection_name="swarm", api_key="api_key" + ) scalablegroupchat.append_to_db("result") assert mock_add.call_count == 1 -@patch('swarms.swarms.scalable_groupchat.ScalableGroupChat.collection.add') -@patch('swarms.swarms.scalable_groupchat.ScalableGroupChat.embed') -@patch('swarms.swarms.scalable_groupchat.ScalableGroupChat.run') + +@patch("swarms.swarms.scalable_groupchat.ScalableGroupChat.collection.add") +@patch("swarms.swarms.scalable_groupchat.ScalableGroupChat.embed") +@patch("swarms.swarms.scalable_groupchat.ScalableGroupChat.run") def test_scalablegroupchat_chat(mock_run, mock_embed, mock_add): - scalablegroupchat = ScalableGroupChat(worker_count=5, collection_name="swarm", api_key="api_key") + scalablegroupchat = ScalableGroupChat( + worker_count=5, collection_name="swarm", api_key="api_key" + ) scalablegroupchat.chat(sender_id=1, receiver_id=2, message="Hello, Agent 2!") assert mock_embed.call_count == 1 assert mock_add.call_count == 1 - assert mock_run.call_count == 1 \ No newline at end of file + assert mock_run.call_count == 1 diff --git a/tests/swarms/simple_swarm.py b/tests/swarms/simple_swarm.py index c168dc7f..38598861 100644 --- a/tests/swarms/simple_swarm.py +++ b/tests/swarms/simple_swarm.py @@ -2,36 +2,51 @@ 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") + 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 = 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.Worker.run') + +@patch("swarms.workers.worker.Worker.run") def test_simpleswarm_process_task(mock_run): - simpleswarm = SimpleSwarm(num_workers=5, openai_api_key="api_key", ai_name="ai_name") + 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 = 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') + +@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") + 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 \ No newline at end of file + assert mock_run.call_count == 5