You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
69 lines
1.8 KiB
69 lines
1.8 KiB
from typing import Dict, Any, List
|
|
from swarms.structs.agent import Agent
|
|
|
|
|
|
# Helper functions for manager/corporate agents
|
|
def parse_tasks(
|
|
task: str = None,
|
|
) -> Dict[str, Any]:
|
|
"""Parse tasks
|
|
|
|
Args:
|
|
task (str, optional): _description_. Defaults to None.
|
|
|
|
Returns:
|
|
Dict[str, Any]: _description_
|
|
"""
|
|
tasks = {}
|
|
for line in task.split("\n"):
|
|
if line.startswith("<agent_id>") and line.endwith(
|
|
"</agent_id>"
|
|
):
|
|
agent_id, task = line[10:-11].split("><")
|
|
tasks[agent_id] = task
|
|
return tasks
|
|
|
|
|
|
def find_agent_by_id(
|
|
agent_id: str = None, agents: List[Agent] = None, *args, **kwargs
|
|
) -> Agent:
|
|
"""Find agent by id
|
|
|
|
Args:
|
|
agent_id (str, optional): _description_. Defaults to None.
|
|
agents (List[Agent], optional): _description_. Defaults to None.
|
|
|
|
Returns:
|
|
Agent: _description_
|
|
"""
|
|
for agent in agents:
|
|
if agent.id == agent_id:
|
|
return agent
|
|
return None
|
|
|
|
|
|
def distribute_tasks(
|
|
task: str = None, agents: List[Agent] = None, *args, **kwargs
|
|
):
|
|
"""Distribute tasks to agents
|
|
|
|
Args:
|
|
task (str, optional): _description_. Defaults to None.
|
|
agents (List[Agent], optional): _description_. Defaults to None.
|
|
"""
|
|
# Parse the task to extract tasks and agent id
|
|
tasks = parse_tasks(task)
|
|
|
|
# Distribute tasks to agents
|
|
for agent_id, task in tasks.item():
|
|
assigned_agent = find_agent_by_id(agent_id, agents)
|
|
if assigned_agent:
|
|
print(f"Assigning task {task} to agent {agent_id}")
|
|
output = assigned_agent.run(task, *args, **kwargs)
|
|
print(f"Output from agent {agent_id}: {output}")
|
|
else:
|
|
print(
|
|
f"No agent found with ID {agent_id}. Task '{task}' is"
|
|
" not assigned."
|
|
)
|