parent
c11a2f467c
commit
70d55e94bc
@ -1,12 +0,0 @@
|
||||
- pdf_to_text: "swarms/utils/pdf_to_text.md"
|
||||
- load_model_torch: "swarms/utils/load_model_torch.md"
|
||||
- metrics_decorator: "swarms/utils/metrics_decorator.md"
|
||||
- prep_torch_inference: "swarms/utils/prep_torch_inference.md"
|
||||
- find_image_path: "swarms/utils/find_image_path.md"
|
||||
- print_class_parameters: "swarms/utils/print_class_parameters.md"
|
||||
- extract_code_from_markdown: "swarms/utils/extract_code_from_markdown.md"
|
||||
- check_device: "swarms/utils/check_device.md"
|
||||
- display_markdown_message: "swarms/utils/display_markdown_message.md"
|
||||
- phoenix_tracer: "swarms/utils/phoenix_tracer.md"
|
||||
- limit_tokens_from_string: "swarms/utils/limit_tokens_from_string.md"
|
||||
- math_eval: "swarms/utils/math_eval.md"
|
@ -0,0 +1,48 @@
|
||||
import os
|
||||
|
||||
from dotenv import load_dotenv
|
||||
|
||||
# Import the OpenAIChat model and the Agent struct
|
||||
from swarms.models import OpenAIChat
|
||||
from swarms.structs import Agent
|
||||
from swarms.structs.swarm_net import SwarmNetwork
|
||||
|
||||
# Load the environment variables
|
||||
load_dotenv()
|
||||
|
||||
# Get the API key from the environment
|
||||
api_key = os.environ.get("OPENAI_API_KEY")
|
||||
|
||||
# Initialize the language model
|
||||
llm = OpenAIChat(
|
||||
temperature=0.5,
|
||||
openai_api_key=api_key,
|
||||
)
|
||||
|
||||
## Initialize the workflow
|
||||
agent = Agent(llm=llm, max_loops=1, agent_name="Social Media Manager")
|
||||
agent2 = Agent(llm=llm, max_loops=1, agent_name=" Product Manager")
|
||||
agent3 = Agent(llm=llm, max_loops=1, agent_name="SEO Manager")
|
||||
|
||||
|
||||
# Load the swarmnet with the agents
|
||||
swarmnet = SwarmNetwork(
|
||||
agents=[agent, agent2, agent3],
|
||||
)
|
||||
|
||||
# List the agents in the swarm network
|
||||
out = swarmnet.list_agents()
|
||||
print(out)
|
||||
|
||||
# Run the workflow on a task
|
||||
out = swarmnet.run_single_agent(
|
||||
agent2.id, "Generate a 10,000 word blog on health and wellness."
|
||||
)
|
||||
print(out)
|
||||
|
||||
|
||||
# Run all the agents in the swarm network on a task
|
||||
out = swarmnet.run_many_agents(
|
||||
"Generate a 10,000 word blog on health and wellness."
|
||||
)
|
||||
print(out)
|
@ -0,0 +1,320 @@
|
||||
import asyncio
|
||||
import logging
|
||||
import queue
|
||||
import threading
|
||||
from typing import List, Optional
|
||||
|
||||
from fastapi import FastAPI
|
||||
|
||||
from swarms.structs.agent import Agent
|
||||
from swarms.structs.base import BaseStructure
|
||||
|
||||
|
||||
class SwarmNetwork(BaseStructure):
|
||||
"""
|
||||
SwarmNetwork class
|
||||
|
||||
The SwarmNetwork class is responsible for managing the agents pool
|
||||
and the task queue. It also monitors the health of the agents and
|
||||
scales the pool up or down based on the number of pending tasks
|
||||
and the current load of the agents.
|
||||
|
||||
For example, if the number of pending tasks is greater than the
|
||||
number of agents in the pool, the SwarmNetwork will scale up the
|
||||
pool by adding new agents. If the number of pending tasks is less
|
||||
than the number of agents in the pool, the SwarmNetwork will scale
|
||||
down the pool by removing agents.
|
||||
|
||||
The SwarmNetwork class also provides a simple API for interacting
|
||||
with the agents pool. The API is implemented using the Flask
|
||||
framework and is enabled by default. The API can be disabled by
|
||||
setting the `api_enabled` parameter to False.
|
||||
|
||||
Features:
|
||||
- Agent pool management
|
||||
- Task queue management
|
||||
- Agent health monitoring
|
||||
- Agent pool scaling
|
||||
- Simple API for interacting with the agent pool
|
||||
- Simple API for interacting with the task queue
|
||||
- Simple API for interacting with the agent health monitor
|
||||
- Simple API for interacting with the agent pool scaler
|
||||
- Create APIs for each agent in the pool (optional)
|
||||
- Run each agent on it's own thread
|
||||
- Run each agent on it's own process
|
||||
- Run each agent on it's own container
|
||||
- Run each agent on it's own machine
|
||||
- Run each agent on it's own cluster
|
||||
|
||||
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
idle_threshold: float = 0.2,
|
||||
busy_threshold: float = 0.7,
|
||||
agents: List[Agent] = None,
|
||||
api_enabled: Optional[bool] = False,
|
||||
logging_enabled: Optional[bool] = False,
|
||||
*args,
|
||||
**kwargs,
|
||||
):
|
||||
self.task_queue = queue.Queue()
|
||||
self.idle_threshold = idle_threshold
|
||||
self.busy_threshold = busy_threshold
|
||||
self.lock = threading.Lock()
|
||||
self.agents = agents
|
||||
self.api_enabled = api_enabled
|
||||
self.logging_enabled = logging_enabled
|
||||
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
self.logger = logging.getLogger(__name__)
|
||||
|
||||
if api_enabled:
|
||||
self.api = FastAPI()
|
||||
|
||||
self.agent_pool = []
|
||||
|
||||
def add_task(self, task):
|
||||
"""Add task to the task queue
|
||||
|
||||
Args:
|
||||
task (_type_): _description_
|
||||
|
||||
Example:
|
||||
>>> from swarms.structs.agent import Agent
|
||||
>>> from swarms.structs.swarm_net import SwarmNetwork
|
||||
>>> agent = Agent()
|
||||
>>> swarm = SwarmNetwork(agents=[agent])
|
||||
>>> swarm.add_task("task")
|
||||
|
||||
"""
|
||||
self.logger.info(f"Adding task {task} to queue")
|
||||
try:
|
||||
self.task_queue.put(task)
|
||||
self.logger.info(f"Task {task} added to queue")
|
||||
except Exception as error:
|
||||
print(
|
||||
f"Error adding task to queue: {error} try again with"
|
||||
" a new task"
|
||||
)
|
||||
raise error
|
||||
|
||||
async def async_add_task(self, task):
|
||||
"""Add task to the task queue
|
||||
|
||||
Args:
|
||||
task (_type_): _description_
|
||||
|
||||
Example:
|
||||
>>> from swarms.structs.agent import Agent
|
||||
>>> from swarms.structs.swarm_net import SwarmNetwork
|
||||
>>> agent = Agent()
|
||||
>>> swarm = SwarmNetwork(agents=[agent])
|
||||
>>> swarm.add_task("task")
|
||||
|
||||
"""
|
||||
self.logger.info(
|
||||
f"Adding task {task} to queue asynchronously"
|
||||
)
|
||||
try:
|
||||
# Add task to queue asynchronously with asyncio
|
||||
loop = asyncio.get_running_loop()
|
||||
await loop.run_in_executor(
|
||||
None, self.task_queue.put, task
|
||||
)
|
||||
self.logger.info(f"Task {task} added to queue")
|
||||
except Exception as error:
|
||||
print(
|
||||
f"Error adding task to queue: {error} try again with"
|
||||
" a new task"
|
||||
)
|
||||
raise error
|
||||
|
||||
def run_single_agent(
|
||||
self, agent_id, task: Optional[str] = None, *args, **kwargs
|
||||
):
|
||||
"""Run agent the task on the agent id
|
||||
|
||||
Args:
|
||||
agent_id (_type_): _description_
|
||||
task (str, optional): _description_. Defaults to None.
|
||||
|
||||
Raises:
|
||||
ValueError: _description_
|
||||
|
||||
Returns:
|
||||
_type_: _description_
|
||||
"""
|
||||
self.logger.info(f"Running task {task} on agent {agent_id}")
|
||||
try:
|
||||
for agent in self.agents_pool:
|
||||
if agent.id == agent_id:
|
||||
return agent.run(task, *args, **kwargs)
|
||||
self.logger.info(f"No agent found with ID {agent_id}")
|
||||
raise ValueError(f"No agent found with ID {agent_id}")
|
||||
except Exception as error:
|
||||
print(f"Error running task on agent: {error}")
|
||||
raise error
|
||||
|
||||
def run_many_agents(
|
||||
self, task: Optional[str] = None, *args, **kwargs
|
||||
) -> List:
|
||||
"""Run the task on all agents
|
||||
|
||||
Args:
|
||||
task (str, optional): _description_. Defaults to None.
|
||||
|
||||
Returns:
|
||||
List: _description_
|
||||
"""
|
||||
self.logger.info(f"Running task {task} on all agents")
|
||||
try:
|
||||
return [
|
||||
agent.run(task, *args, **kwargs)
|
||||
for agent in self.agents_pool
|
||||
]
|
||||
except Exception as error:
|
||||
print(f"Error running task on agents: {error}")
|
||||
raise error
|
||||
|
||||
def list_agents(self):
|
||||
"""List all agents
|
||||
|
||||
Returns:
|
||||
List: _description_
|
||||
"""
|
||||
self.logger.info("[Listing all active agents]")
|
||||
try:
|
||||
# return [agent.id for agent in self.agents_pool]
|
||||
for agent in self.agents:
|
||||
num_agents = len(self.agents)
|
||||
self.logger.info(
|
||||
f"[Number of active agents: {num_agents}]"
|
||||
)
|
||||
return self.logger.info(
|
||||
f"[Agent] [ID: {agent.id}] [Name:"
|
||||
f" {agent.agent_name}] [Description:"
|
||||
f" {agent.agent_description}] [Status] [Running]"
|
||||
)
|
||||
except Exception as error:
|
||||
print(f"Error listing agents: {error}")
|
||||
raise error
|
||||
|
||||
def get_agent(self, agent_id):
|
||||
"""Get agent by id
|
||||
|
||||
Args:
|
||||
agent_id (_type_): _description_
|
||||
|
||||
Returns:
|
||||
_type_: _description_
|
||||
"""
|
||||
self.logger.info(f"Getting agent {agent_id}")
|
||||
|
||||
try:
|
||||
for agent in self.agents_pool:
|
||||
if agent.id == agent_id:
|
||||
return agent
|
||||
raise ValueError(f"No agent found with ID {agent_id}")
|
||||
except Exception as error:
|
||||
self.logger.error(f"Error getting agent: {error}")
|
||||
raise error
|
||||
|
||||
def add_agent(self, agent):
|
||||
"""Add agent to the agent pool
|
||||
|
||||
Args:
|
||||
agent (_type_): _description_
|
||||
"""
|
||||
self.logger.info(f"Adding agent {agent} to pool")
|
||||
try:
|
||||
self.agents_pool.append(agent)
|
||||
except Exception as error:
|
||||
print(f"Error adding agent to pool: {error}")
|
||||
raise error
|
||||
|
||||
def remove_agent(self, agent_id):
|
||||
"""Remove agent from the agent pool
|
||||
|
||||
Args:
|
||||
agent_id (_type_): _description_
|
||||
"""
|
||||
self.logger.info(f"Removing agent {agent_id} from pool")
|
||||
try:
|
||||
for agent in self.agents_pool:
|
||||
if agent.id == agent_id:
|
||||
self.agents_pool.remove(agent)
|
||||
return
|
||||
raise ValueError(f"No agent found with ID {agent_id}")
|
||||
except Exception as error:
|
||||
print(f"Error removing agent from pool: {error}")
|
||||
raise error
|
||||
|
||||
async def async_remove_agent(self, agent_id):
|
||||
"""Remove agent from the agent pool
|
||||
|
||||
Args:
|
||||
agent_id (_type_): _description_
|
||||
"""
|
||||
self.logger.info(f"Removing agent {agent_id} from pool")
|
||||
try:
|
||||
# Remove agent from pool asynchronously with asyncio
|
||||
loop = asyncio.get_running_loop()
|
||||
await loop.run_in_executor(
|
||||
None, self.remove_agent, agent_id
|
||||
)
|
||||
except Exception as error:
|
||||
print(f"Error removing agent from pool: {error}")
|
||||
raise error
|
||||
|
||||
def scale_up(self, num_agents: int = 1):
|
||||
"""Scale up the agent pool
|
||||
|
||||
Args:
|
||||
num_agents (int, optional): _description_. Defaults to 1.
|
||||
"""
|
||||
self.logger.info(f"Scaling up agent pool by {num_agents}")
|
||||
try:
|
||||
for _ in range(num_agents):
|
||||
self.agents_pool.append(Agent())
|
||||
except Exception as error:
|
||||
print(f"Error scaling up agent pool: {error}")
|
||||
raise error
|
||||
|
||||
def scale_down(self, num_agents: int = 1):
|
||||
"""Scale down the agent pool
|
||||
|
||||
Args:
|
||||
num_agents (int, optional): _description_. Defaults to 1.
|
||||
"""
|
||||
for _ in range(num_agents):
|
||||
self.agents_pool.pop()
|
||||
|
||||
# - Create APIs for each agent in the pool (optional) with fastapi
|
||||
def create_apis_for_agents(self):
|
||||
"""Create APIs for each agent in the pool (optional) with fastapi
|
||||
|
||||
Returns:
|
||||
_type_: _description_
|
||||
"""
|
||||
self.apis = []
|
||||
for agent in self.agents:
|
||||
self.api.get(f"/{agent.id}")
|
||||
|
||||
def run_agent(task: str, *args, **kwargs):
|
||||
return agent.run(task, *args, **kwargs)
|
||||
|
||||
self.apis.append(self.api)
|
||||
|
||||
def run(self):
|
||||
"""run the swarm network"""
|
||||
# Observe all agents in the pool
|
||||
self.logger.info("Starting the SwarmNetwork")
|
||||
|
||||
for agent in self.agents:
|
||||
self.logger.info(f"Starting agent {agent.id}")
|
||||
self.logger.info(
|
||||
f"[Agent][{agent.id}] [Status] [Running] [Awaiting"
|
||||
" Task]"
|
||||
)
|
@ -0,0 +1,95 @@
|
||||
import json
|
||||
from typing import List, Optional
|
||||
|
||||
from pydantic.v1 import BaseModel, Field, Json, root_validator
|
||||
|
||||
from swarms.structs.agent import Agent
|
||||
from swarms.structs.task import Task
|
||||
|
||||
|
||||
class Team(BaseModel):
|
||||
"""
|
||||
Class that represents a group of agents, how they should work together and
|
||||
their tasks.
|
||||
|
||||
Attributes:
|
||||
tasks (Optional[List[Task]]): List of tasks.
|
||||
agents (Optional[List[Agent]]): List of agents in this Team.
|
||||
architecture (str): Architecture that the Team will follow. Default is "sequential".
|
||||
verbose (bool): Verbose mode for the Agent Execution. Default is False.
|
||||
config (Optional[Json]): Configuration of the Team. Default is None.
|
||||
"""
|
||||
|
||||
tasks: Optional[List[Task]] = Field(description="List of tasks")
|
||||
agents: Optional[List[Agent]] = Field(description="List of agents in this Team.")
|
||||
architecture = Field(
|
||||
description="architecture that the Team will follow.", default="sequential"
|
||||
)
|
||||
verbose: bool = Field(
|
||||
description="Verbose mode for the Agent Execution", default=False
|
||||
)
|
||||
config: Optional[Json] = Field(
|
||||
description="Configuration of the Team.", default=None
|
||||
)
|
||||
|
||||
@root_validator(pre=True)
|
||||
def check_config(_cls, values):
|
||||
if not values.get("config") and (
|
||||
not values.get("agents") and not values.get("tasks")
|
||||
):
|
||||
raise ValueError("Either agents and task need to be set or config.")
|
||||
|
||||
if values.get("config"):
|
||||
config = json.loads(values.get("config"))
|
||||
if not config.get("agents") or not config.get("tasks"):
|
||||
raise ValueError("Config should have agents and tasks.")
|
||||
|
||||
values["agents"] = [Agent(**agent) for agent in config["agents"]]
|
||||
|
||||
tasks = []
|
||||
for task in config["tasks"]:
|
||||
task_agent = [
|
||||
agt for agt in values["agents"] if agt.role == task["agent"]
|
||||
][0]
|
||||
del task["agent"]
|
||||
tasks.append(Task(**task, agent=task_agent))
|
||||
|
||||
values["tasks"] = tasks
|
||||
return values
|
||||
|
||||
def run(self) -> str:
|
||||
"""
|
||||
Kickoff the Team to work on its tasks.
|
||||
|
||||
Returns:
|
||||
output (List[str]): Output of the Team for each task.
|
||||
"""
|
||||
if self.architecture == "sequential":
|
||||
return self.__sequential_loop()
|
||||
|
||||
def __sequential_loop(self) -> str:
|
||||
"""
|
||||
Loop that executes the sequential architecture.
|
||||
|
||||
Returns:
|
||||
output (str): Output of the Team.
|
||||
"""
|
||||
task_outcome = None
|
||||
for task in self.tasks:
|
||||
# Add delegation tools to the task if the agent allows it
|
||||
# if task.agent.allow_delegation:
|
||||
# tools = AgentTools(agents=self.agents).tools()
|
||||
# task.tools += tools
|
||||
|
||||
self.__log(f"\nWorking Agent: {task.agent.role}")
|
||||
self.__log(f"Starting Task: {task.description} ...")
|
||||
|
||||
task_outcome = task.execute(task_outcome)
|
||||
|
||||
self.__log(f"Task output: {task_outcome}")
|
||||
|
||||
return task_outcome
|
||||
|
||||
def __log(self, message):
|
||||
if self.verbose:
|
||||
print(message)
|
@ -0,0 +1,50 @@
|
||||
import pytest
|
||||
from unittest.mock import Mock, patch
|
||||
from swarms.structs.swarm_net import SwarmNetwork
|
||||
from swarms.structs.agent import Agent
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def swarm_network():
|
||||
agents = [Agent(id=f"Agent_{i}") for i in range(5)]
|
||||
return SwarmNetwork(agents=agents)
|
||||
|
||||
|
||||
def test_swarm_network_init(swarm_network):
|
||||
assert isinstance(swarm_network.agents, list)
|
||||
assert len(swarm_network.agents) == 5
|
||||
|
||||
|
||||
@patch("swarms.structs.swarm_net.SwarmNetwork.logger")
|
||||
def test_run(mock_logger, swarm_network):
|
||||
swarm_network.run()
|
||||
assert (
|
||||
mock_logger.info.call_count == 10
|
||||
) # 2 log messages per agent
|
||||
|
||||
|
||||
def test_run_with_mocked_agents(mocker, swarm_network):
|
||||
mock_agents = [Mock(spec=Agent) for _ in range(5)]
|
||||
mocker.patch.object(swarm_network, "agents", mock_agents)
|
||||
swarm_network.run()
|
||||
for mock_agent in mock_agents:
|
||||
assert mock_agent.run.called
|
||||
|
||||
|
||||
def test_swarm_network_with_no_agents():
|
||||
swarm_network = SwarmNetwork(agents=[])
|
||||
assert swarm_network.agents == []
|
||||
|
||||
|
||||
def test_swarm_network_add_agent(swarm_network):
|
||||
new_agent = Agent(id="Agent_5")
|
||||
swarm_network.add_agent(new_agent)
|
||||
assert len(swarm_network.agents) == 6
|
||||
assert swarm_network.agents[-1] == new_agent
|
||||
|
||||
|
||||
def test_swarm_network_remove_agent(swarm_network):
|
||||
agent_to_remove = swarm_network.agents[0]
|
||||
swarm_network.remove_agent(agent_to_remove)
|
||||
assert len(swarm_network.agents) == 4
|
||||
assert agent_to_remove not in swarm_network.agents
|
Loading…
Reference in new issue