From a403ed3c3b5993ccd75c961817273cf3ea046f43 Mon Sep 17 00:00:00 2001 From: Kye Date: Tue, 23 Jan 2024 13:52:27 -0500 Subject: [PATCH] [BUFG][Worker] --- company_example.py | 38 +++++++++++++++++++ pyproject.toml | 2 +- swarms/agents/__init__.py | 4 +- swarms/structs/company.py | 29 ++++++++++++--- tests/structs/test_company.py | 70 +++++++++++++++++++++++++++++++++++ 5 files changed, 135 insertions(+), 8 deletions(-) create mode 100644 company_example.py create mode 100644 tests/structs/test_company.py diff --git a/company_example.py b/company_example.py new file mode 100644 index 00000000..df2d2506 --- /dev/null +++ b/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() diff --git a/pyproject.toml b/pyproject.toml index 22f09946..7cb65d39 100644 --- a/pyproject.toml +++ b/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 "] diff --git a/swarms/agents/__init__.py b/swarms/agents/__init__.py index 9c855436..4b786bf4 100644 --- a/swarms/agents/__init__.py +++ b/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", ] diff --git a/swarms/structs/company.py b/swarms/structs/company.py index b01c0a00..80fe3eef 100644 --- a/swarms/structs/company.py +++ b/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) + + diff --git a/tests/structs/test_company.py b/tests/structs/test_company.py new file mode 100644 index 00000000..0b1ec105 --- /dev/null +++ b/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)