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.
115 lines
3.9 KiB
115 lines
3.9 KiB
import os
|
|
from dotenv import load_dotenv
|
|
|
|
# Swarm imports
|
|
from swarms.structs.agent import Agent
|
|
from swarms.structs.hiearchical_swarm import (
|
|
HierarchicalSwarm,
|
|
SwarmSpec,
|
|
OrganizationalUnit,
|
|
)
|
|
from swarms.utils.function_caller_model import OpenAIFunctionCaller
|
|
|
|
# Load environment variables
|
|
load_dotenv()
|
|
|
|
# Create the agents first
|
|
research_manager = Agent(
|
|
agent_name="Research Manager",
|
|
agent_description="Manages research operations and coordinates research tasks",
|
|
system_prompt="You are a research manager responsible for overseeing research projects and coordinating research efforts.",
|
|
model_name="gpt-4o",
|
|
)
|
|
|
|
data_analyst = Agent(
|
|
agent_name="Data Analyst",
|
|
agent_description="Analyzes data and generates insights",
|
|
system_prompt="You are a data analyst specializing in processing and analyzing data to extract meaningful insights.",
|
|
model_name="gpt-4o",
|
|
)
|
|
|
|
research_assistant = Agent(
|
|
agent_name="Research Assistant",
|
|
agent_description="Assists with research tasks and data collection",
|
|
system_prompt="You are a research assistant who helps gather information and support research activities.",
|
|
model_name="gpt-4o",
|
|
)
|
|
|
|
development_manager = Agent(
|
|
agent_name="Development Manager",
|
|
agent_description="Manages development projects and coordinates development tasks",
|
|
system_prompt="You are a development manager responsible for overseeing software development projects and coordinating development efforts.",
|
|
model_name="gpt-4o",
|
|
)
|
|
|
|
software_engineer = Agent(
|
|
agent_name="Software Engineer",
|
|
agent_description="Develops and implements software solutions",
|
|
system_prompt="You are a software engineer specializing in building and implementing software solutions.",
|
|
model_name="gpt-4o",
|
|
)
|
|
|
|
qa_engineer = Agent(
|
|
agent_name="QA Engineer",
|
|
agent_description="Tests and ensures quality of software",
|
|
system_prompt="You are a QA engineer responsible for testing software and ensuring its quality.",
|
|
model_name="gpt-4o",
|
|
)
|
|
|
|
# Create organizational units with the agents
|
|
research_unit = OrganizationalUnit(
|
|
name="Research Unit",
|
|
description="Handles research and analysis tasks",
|
|
manager=research_manager,
|
|
members=[data_analyst, research_assistant],
|
|
)
|
|
|
|
development_unit = OrganizationalUnit(
|
|
name="Development Unit",
|
|
description="Handles development and implementation tasks",
|
|
manager=development_manager,
|
|
members=[software_engineer, qa_engineer],
|
|
)
|
|
|
|
# Initialize the director agent
|
|
director = OpenAIFunctionCaller(
|
|
model_name="gpt-4o",
|
|
system_prompt=(
|
|
"As the Director of this Hierarchical Agent Swarm, you are responsible for:\n"
|
|
"1. Analyzing tasks and breaking them down into subtasks\n"
|
|
"2. Assigning tasks to appropriate organizational units\n"
|
|
"3. Coordinating communication between units\n"
|
|
"4. Ensuring tasks are completed efficiently and effectively\n"
|
|
"5. Providing feedback and guidance to units as needed\n\n"
|
|
"Your decisions should be based on the capabilities of each unit and the requirements of the task."
|
|
),
|
|
api_key=os.getenv("OPENAI_API_KEY"),
|
|
temperature=0.5,
|
|
base_model=SwarmSpec,
|
|
max_tokens=10000,
|
|
)
|
|
|
|
# Initialize the hierarchical swarm with the organizational units
|
|
swarm = HierarchicalSwarm(
|
|
name="Example Hierarchical Swarm",
|
|
description="A hierarchical swarm demonstrating multi-unit collaboration",
|
|
director=director,
|
|
organizational_units=[research_unit, development_unit],
|
|
max_loops=2, # Allow for feedback and iteration
|
|
output_type="dict",
|
|
)
|
|
|
|
# Example task to run through the swarm
|
|
task = """
|
|
Develop a comprehensive market analysis for a new AI-powered productivity tool.
|
|
The analysis should include:
|
|
1. Market research and competitor analysis
|
|
2. User needs and pain points
|
|
3. Technical feasibility assessment
|
|
4. Implementation recommendations
|
|
"""
|
|
|
|
# Run the task through the swarm
|
|
result = swarm.run(task)
|
|
print("Swarm Results:", result)
|