[CLEANUP]

playground/structs/multi_agent_collaboration/swarm_network_example.py → playground/structs/swarm_network_example.py

@@ -2,12 +2,11 @@
 import os
 from swarms import (
     Agent,
-    OpenAIChat,
     SwarmNetwork,
-    Anthropic,
+    OpenAIChat,
     TogetherLLM,
 )
-from swarms.memory import ChromaDB
+from playground.memory.chromadb_example import ChromaDB
 from dotenv import load_dotenv
 
 # load the environment variables
@@ -22,7 +21,7 @@
 )
 
 # Initialize the Anthropic
-anthropic = Anthropic(max_tokens=3000)
+anthropic = OpenAIChat(max_tokens=3000)
 
 # TogeterLM
 together_llm = TogetherLLM(

swarms/agents/__init__.py

@@ -13,13 +13,10 @@
     check_success,
 )
 from swarms.agents.tool_agent import ToolAgent
-from swarms.agents.worker_agent import Worker
 
 __all__ = [
     "AbstractAgent",
     "ToolAgent",
-    "SimpleAgent",
-    "OmniModalAgent",
     "check_done",
     "check_finished",
     "check_complete",
@@ -30,6 +27,5 @@
     "check_cancelled",
     "check_exit",
     "check_end",
-    "Worker",
     "agent_wrapper",
 ]

swarms/structs/base_swarm.py

@@ -113,8 +113,10 @@ def __init__(
         self.collective_memory_system = collective_memory_system
         self.agent_ops_on = agent_ops_on
 
+        logger.info("Reliability checks activated.")
         # Ensure that agents is exists
         if self.agents is None:
+            logger.info("Agents must be provided.")
             raise ValueError("Agents must be provided.")
 
         # Ensure that agents is a list

swarms/structs/swarm_net.py

@@ -120,7 +120,7 @@ def __init__(
         *args,
         **kwargs,
     ):
-        super().__init__(*args, **kwargs)
+        super().__init__(agents=agents, *args, **kwargs)
         self.name = name
         self.description = description
         self.agents = agents

tests/structs/test_moa.py

@@ -0,0 +1,74 @@
+import pytest
+from unittest.mock import Mock, patch
+from swarms.structs.mixture_of_agents import MixtureOfAgents
+from swarms.structs.agent import Agent
+from swarms.memory.base_vectordb import BaseVectorDatabase
+
+
+def test_init():
+    with patch.object(
+        MixtureOfAgents, "agent_check"
+    ) as mock_agent_check, patch.object(
+        MixtureOfAgents, "final_agent_check"
+    ) as mock_final_agent_check, patch.object(
+        MixtureOfAgents, "swarm_initialization"
+    ) as mock_swarm_initialization, patch.object(
+        MixtureOfAgents, "communication_protocol"
+    ) as mock_communication_protocol:
+        agents = [Mock(spec=Agent)]
+        final_agent = Mock(spec=Agent)
+        scp = Mock(spec=BaseVectorDatabase)
+        MixtureOfAgents(agents=agents, final_agent=final_agent, scp=scp)
+        mock_agent_check.assert_called_once()
+        mock_final_agent_check.assert_called_once()
+        mock_swarm_initialization.assert_called_once()
+        mock_communication_protocol.assert_called_once()
+
+
+def test_communication_protocol():
+    agents = [Mock(spec=Agent)]
+    final_agent = Mock(spec=Agent)
+    scp = Mock(spec=BaseVectorDatabase)
+    swarm = MixtureOfAgents(
+        agents=agents, final_agent=final_agent, scp=scp
+    )
+    swarm.communication_protocol()
+    for agent in agents:
+        agent.long_term_memory.assert_called_once_with(scp)
+
+
+def test_agent_check():
+    final_agent = Mock(spec=Agent)
+    with pytest.raises(TypeError):
+        MixtureOfAgents(agents="not a list", final_agent=final_agent)
+    with pytest.raises(TypeError):
+        MixtureOfAgents(agents=["not an agent"], final_agent=final_agent)
+
+
+def test_final_agent_check():
+    agents = [Mock(spec=Agent)]
+    with pytest.raises(TypeError):
+        MixtureOfAgents(agents=agents, final_agent="not an agent")
+
+
+def test_swarm_initialization():
+    with patch("swarms.structs.mixture_of_agents.logger") as mock_logger:
+        agents = [Mock(spec=Agent)]
+        final_agent = Mock(spec=Agent)
+        swarm = MixtureOfAgents(agents=agents, final_agent=final_agent)
+        swarm.swarm_initialization()
+        assert mock_logger.info.call_count == 3
+
+
+def test_run():
+    with patch("swarms.structs.mixture_of_agents.logger"), patch(
+        "builtins.open", new_callable=Mock
+    ) as mock_open:
+        agents = [Mock(spec=Agent)]
+        final_agent = Mock(spec=Agent)
+        swarm = MixtureOfAgents(agents=agents, final_agent=final_agent)
+        swarm.run("task")
+        for agent in agents:
+            agent.run.assert_called_once()
+        final_agent.run.assert_called_once()
+        mock_open.assert_called_once_with(swarm.saved_file_name, "w")