diff --git a/swarms/swarms/__init__.py b/swarms/swarms/__init__.py index 6aa76db3..38ced622 100644 --- a/swarms/swarms/__init__.py +++ b/swarms/swarms/__init__.py @@ -1,10 +1,11 @@ from swarms.structs.autoscaler import AutoScaler from swarms.swarms.god_mode import GodMode from swarms.swarms.multi_agent_collab import MultiAgentCollaboration - +from swarms.swarms.base import AbstractSwarm __all__ = [ "AutoScaler", "GodMode", "MultiAgentCollaboration", + "AbstractSwarm", ] diff --git a/swarms/swarms/base.py b/swarms/swarms/base.py index 2f65bf86..d5c6eea6 100644 --- a/swarms/swarms/base.py +++ b/swarms/swarms/base.py @@ -2,9 +2,12 @@ import asyncio import concurrent.futures import time from abc import ABC, abstractmethod -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Optional, Callable from swarms.structs.agent import Agent from swarms.agents.base import AbstractWorker +from concurrent.futures import ThreadPoolExecutor, as_completed +import logging +from termcolor import colored class AbstractSwarm(ABC): @@ -360,5 +363,83 @@ class AbstractSwarm(ABC): # Assign task to agent by their agent id agent = self.select_agent(agent_id) return agent.run(task, *args, **kwargs) - - \ No newline at end of file + + def task_assignment_by_name( + self, task: str, agent_name: str, *args, **kwargs + ): + """ + Assign a task to an agent + """ + # Assign task to agent by their agent id + agent = self.select_agent_by_name(agent_name) + return agent.run(task, *args, **kwargs) + + def concurrent_run(self, task: str) -> List[str]: + """Synchronously run the task on all llms and collect responses""" + with ThreadPoolExecutor() as executor: + future_to_llm = { + executor.submit(agent, task): agent + for agent in self.agents + } + responses = [] + for future in as_completed(future_to_llm): + try: + responses.append(future.result()) + except Exception as error: + print( + f"{future_to_llm[future]} generated an" + f" exception: {error}" + ) + self.last_responses = responses + self.task_history.append(task) + return responses + + def add_llm(self, agent: Callable): + """Add an llm to the god mode""" + self.agents.append(agent) + + def remove_llm(self, agent: Callable): + """Remove an llm from the god mode""" + self.agents.remove(agent) + + def add_agent(self, agent: Agent = None, *args, **kwargs): + """Add an agent to the swarm + + Args: + agent (Agent, optional): _description_. Defaults to None. + + Returns: + _type_: _description_ + """ + self.agents.append(agent) + return agent + + def run_all(self, task: str = None, *args, **kwargs): + """Run all agents + + Args: + task (str, optional): _description_. Defaults to None. + + Returns: + _type_: _description_ + """ + responses = [] + for agent in self.agents: + responses.append(agent(task, *args, **kwargs)) + return responses + + def run_on_all_agents(self, task: str = None, *args, **kwargs): + """Run on all agents + + Args: + task (str, optional): _description_. Defaults to None. + + Returns: + _type_: _description_ + """ + with ThreadPoolExecutor() as executor: + responses = executor.map( + lambda agent: agent(task, *args, **kwargs), + self.agents, + ) + return list(responses) diff --git a/swarms/swarms/groupchat.py b/swarms/swarms/groupchat.py index 76f287bc..f3677023 100644 --- a/swarms/swarms/groupchat.py +++ b/swarms/swarms/groupchat.py @@ -98,6 +98,11 @@ class GroupChat: return self.next_agent(last_speaker) def _participant_roles(self): + """Print the roles of the participants. + + Returns: + _type_: _description_ + """ return "\n".join( [ f"{agent.name}: {agent.system_message}" @@ -106,6 +111,14 @@ class GroupChat: ) def format_history(self, messages: List[Dict]) -> str: + """Format the history of the messages. + + Args: + messages (List[Dict]): _description_ + + Returns: + str: _description_ + """ formatted_messages = [] for message in messages: formatted_message = ( @@ -137,6 +150,14 @@ class GroupChatManager: self.selector = selector def __call__(self, task: str): + """Call 'GroupChatManager' instance as a function. + + Args: + task (str): _description_ + + Returns: + _type_: _description_ + """ self.groupchat.messages.append( {"role": self.selector.name, "content": task} ) diff --git a/tests/models/test_gpt4_vision_api.py b/tests/models/test_gpt4_vision_api.py index c8485e83..dfd03e27 100644 --- a/tests/models/test_gpt4_vision_api.py +++ b/tests/models/test_gpt4_vision_api.py @@ -92,7 +92,6 @@ def test_initialization_with_custom_key(): assert api.openai_api_key == custom_key - def test_run_with_exception(gpt_api): task = "What is in the image?" img_url = img