[BUFG][Worker]

company_example.py

@@ -0,0 +1,38 @@
+# Example
+
+import os
+
+from dotenv import load_dotenv
+
+from swarms import Agent, OpenAIChat
+from swarms.structs.company import Company
+
+load_dotenv()
+
+llm = OpenAIChat(
+    openai_api_key=os.getenv("OPENAI_API_KEY"), max_tokens=4000
+)
+
+ceo = Agent(llm=llm, ai_name="CEO")
+dev = Agent(llm=llm, ai_name="Developer")
+va = Agent(llm=llm, ai_name="VA")
+
+# Create a company
+company = Company(
+    org_chart = [[dev, va]],
+    shared_instructions="Do your best",
+    ceo=ceo,
+)
+
+# Add agents to the company
+hr = Agent(llm=llm, name="HR")
+company.add(hr)
+
+# Get an agent from the company
+hr = company.get("CEO")
+
+# Remove an agent from the company
+company.remove(hr)
+
+# Run the company
+company.run()

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "swarms"
-version = "3.7.5"
+version = "3.7.6"
 description = "Swarms - Pytorch"
 license = "MIT"
 authors = ["Kye Gomez <kye@apac.ai>"]

swarms/agents/__init__.py

@@ -14,7 +14,7 @@ from swarms.agents.stopping_conditions import (
     check_success,
 )
 from swarms.agents.tool_agent import ToolAgent
-from swarms.agents.worker_agent import WorkerAgent
+from swarms.agents.worker_agent import Worker
 
 __all__ = [
     "AbstractAgent",
@@ -31,5 +31,5 @@ __all__ = [
     "check_cancelled",
     "check_exit",
     "check_end",
-    "WorkerAgent",
+    "Worker",
 ]

swarms/structs/company.py

@@ -1,21 +1,23 @@
-from typing import List, Optional, Union, Dict
 from dataclasses import dataclass, field
+from typing import Dict, List, Optional, Union
+
 from swarms.structs.agent import Agent
 from swarms.utils.logger import logger
-
+from swarms.structs.conversation import Conversation
 
 @dataclass
 class Company:
     """
     Represents a company with a hierarchical organizational structure.
     """
-    org_chart: Union[List[Agent], List[List[Agent]]]
+    org_chart: List[List[Agent]]
     shared_instructions: str = None
     ceo: Optional[Agent] = None
     agents: List[Agent] = field(default_factory=list)
     agent_interactions: Dict[str, List[str]] = field(
         default_factory=dict
     )
+    history: Conversation = field(default_factory=Conversation)
 
     def __post_init__(self):
         self._parse_org_chart(self.org_chart)
@@ -81,7 +83,13 @@ class Company:
         Args:
             agent (Agent): The agent to be removed.
         """
-        self.agents.remove(agent)
+        try:
+            self.agents.remove(agent)
+        except Exception as error:
+            logger.error(
+                f"[ERROR][CLASS: Company][METHOD: remove] {error}"
+            )
+            raise error
 
     def _parse_org_chart(
         self, org_chart: Union[List[Agent], List[List[Agent]]]
@@ -96,7 +104,6 @@ class Company:
         Raises:
             ValueError: If more than one CEO is found in the org chart or if an invalid
                 agent is encountered.
-
         """
         try:
             for node in org_chart:
@@ -132,6 +139,16 @@ class Company:
         agent1: Agent,
         agent2: Agent,
     ) -> None:
+        """
+        Initializes the interaction between two agents.
+
+        Args:
+            agent1 (Agent): The first agent involved in the interaction.
+            agent2 (Agent): The second agent involved in the interaction.
+
+        Returns:
+            None
+        """
         if agent1.ai_name not in self.agents_interactions:
             self.agents_interactions[agent1.ai_name] = []
         self.agents_interactions[agent1.ai_name].append(
@@ -154,3 +171,5 @@ class Company:
                 )
                 print(f"{task_description} is being executed")
                 agent.run(task_description)
+
+

tests/structs/test_company.py

@@ -0,0 +1,70 @@
+import pytest
+from swarms.structs.agent import Agent
+from swarms.structs.company import Company
+from swarms import OpenAIChat
+
+# Mock OpenAIChat instance
+llm = OpenAIChat(openai_api_key="test_key", max_tokens=4000)
+
+# Mock Agents
+ceo = Agent(llm=llm, name="CEO")
+dev = Agent(llm=llm, name="Developer")
+va = Agent(llm=llm, name="VA")
+hr = Agent(llm=llm, name="HR")
+shared_instructions = "Listen to your boss"
+
+
+def test_add_agent():
+    company = Company(
+        org_chart=[[ceo, [dev, va]]],
+        shared_instructions=shared_instructions,
+    )
+    company.add(hr)
+    assert hr in company.agents
+
+
+def test_get_agent():
+    company = Company(
+        org_chart=[[ceo, [dev, va]]],
+        shared_instructions=shared_instructions,
+    )
+    company.add(hr)
+    assert company.get("HR") == hr
+
+
+def test_remove_agent():
+    company = Company(
+        org_chart=[[ceo, [dev, va]]],
+        shared_instructions=shared_instructions,
+    )
+    company.add(hr)
+    company.remove(hr)
+    assert hr not in company.agents
+
+
+def test_add_existing_agent():
+    company = Company(
+        org_chart=[[ceo, [dev, va]]],
+        shared_instructions=shared_instructions,
+    )
+    company.add(hr)
+    with pytest.raises(ValueError):
+        company.add(hr)
+
+
+def test_get_nonexistent_agent():
+    company = Company(
+        org_chart=[[ceo, [dev, va]]],
+        shared_instructions=shared_instructions,
+    )
+    with pytest.raises(ValueError):
+        company.get("Nonexistent")
+
+
+def test_remove_nonexistent_agent():
+    company = Company(
+        org_chart=[[ceo, [dev, va]]],
+        shared_instructions=shared_instructions,
+    )
+    with pytest.raises(ValueError):
+        company.remove(hr)