The Orchestration component provides tools for coordinating multiple agents in the atomic agents framework. It includes an orchestrator class that can register agents and execute workflows, with full integration with the BrainBlend-AI/atomic-agents framework.
The Orchestrator class coordinates the execution of multiple agents. It provides:
- Agent Registration: Register and unregister agents with the orchestrator
- Workflow Management: Create and execute workflows involving multiple agents
- Message Routing: Route messages between agents using a message broker
- Atomic-Agents Integration: Seamless integration with the atomic-agents framework
- Hybrid Execution: Support for both legacy and atomic agents in the same workflow
__init__(message_broker): Initialize the orchestrator with a message broker and create an atomic orchestratorstart(): Start the orchestrator and initialize all registered agentsstop(): Stop the orchestrator and all registered agentsregister_agent(agent): Register an agent with both the orchestrator and atomic orchestrator if applicableunregister_agent(agent_id): Unregister an agent from both orchestratorscreate_workflow(workflow_id, workflow_definition): Create a new workflow in both orchestratorsexecute_workflow(workflow_id, input_data): Execute a workflow with intelligent routing to atomic or legacy execution
The Orchestrator class is implemented in orchestrator.py:
from typing import Dict, List, Any, Optional
import asyncio
# Import from atomic-agents
from atomic_agents.lib.components.orchestrator import Orchestrator as AtomicOrchestrator
# Import from our framework
from core.agent_base.base_agent import BaseAgent
from core.messaging.message_broker import MessageBroker
from core.messaging.message import Message
class Orchestrator:
"""
Coordinates the execution of multiple agents.
This orchestrator leverages the atomic-agents framework's orchestration capabilities
while maintaining compatibility with our existing agent architecture.
"""
def __init__(self, message_broker: Optional[MessageBroker] = None):
"""
Initialize the orchestrator.
Args:
message_broker: Optional message broker to use for inter-agent communication
"""
self.agents: Dict[str, BaseAgent] = {}
self.message_broker = message_broker or MessageBroker()
self.workflows: Dict[str, Dict[str, Any]] = {}
# Create an atomic orchestrator instance
self.atomic_orchestrator = AtomicOrchestrator()
async def start(self):
"""
Start the orchestrator and initialize all registered agents.
"""
# Initialize all agents
for agent_id, agent in self.agents.items():
await agent.initialize()
async def stop(self):
"""
Stop the orchestrator and shut down all agents.
"""
# Shutdown all agents
for agent_id, agent in self.agents.items():
await agent.shutdown()
def register_agent(self, agent: BaseAgent):
"""
Register an agent with the orchestrator.
Args:
agent: The agent to register
"""
self.agents[agent.agent_id] = agent
# Set the message broker for the agent if it doesn't have one
if not agent.message_broker:
agent.message_broker = self.message_broker
# Subscribe the agent to receive messages
async def handle_message(message: Message):
await self._handle_agent_message(agent, message)
self.message_broker.subscribe(agent.agent_id, handle_message)
# Register with atomic orchestrator if the agent has an atomic_agent
if hasattr(agent, 'config') and agent.config and 'atomic_agent' in agent.config:
self.atomic_orchestrator.register_agent(agent.config['atomic_agent'])
def unregister_agent(self, agent_id: str):
"""
Unregister an agent from the orchestrator.
Args:
agent_id: ID of the agent to unregister
"""
if agent_id in self.agents:
agent = self.agents[agent_id]
# Unregister from atomic orchestrator if applicable
if hasattr(agent, 'config') and agent.config and 'atomic_agent' in agent.config:
self.atomic_orchestrator.unregister_agent(agent.config['atomic_agent'])
# Unsubscribe from message broker
self.message_broker.unsubscribe(agent_id)
# Remove from agents dict
del self.agents[agent_id]
async def _handle_agent_message(self, agent: BaseAgent, message: Message):
"""
Handle a message received by an agent.
Args:
agent: The agent that received the message
message: The message received
"""
try:
# Check if the message requires a response
if message.metadata.get("requires_response", False):
try:
# Process the message content with the agent
result = await agent.process(message.content)
# Create a response message
response = message.create_response(result)
# Publish the response
await self.message_broker.publish(response)
# If the agent has an atomic agent, also process with atomic orchestrator
if hasattr(agent, 'config') and agent.config and 'atomic_agent' in agent.config:
atomic_agent = agent.config['atomic_agent']
atomic_message = message.to_atomic_format()
await self.atomic_orchestrator.handle_message(atomic_agent, atomic_message)
except Exception as e:
# If there's an error, send an error response
error_response = message.create_response(
{"error": str(e)},
{"error": True}
)
await self.message_broker.publish(error_response)
else:
# For messages that don't require a response, just process them
await agent.process(message.content)
# If the agent has an atomic agent, also process with atomic orchestrator
if hasattr(agent, 'config') and agent.config and 'atomic_agent' in agent.config:
atomic_agent = agent.config['atomic_agent']
atomic_message = message.to_atomic_format()
await self.atomic_orchestrator.handle_message(atomic_agent, atomic_message)
except Exception as e:
print(f"Error handling message in agent {agent.agent_id}: {e}")
def create_workflow(self, workflow_id: str, workflow_definition: Dict[str, Any]):
"""
Create a new workflow.
Args:
workflow_id: Unique identifier for the workflow
workflow_definition: Definition of the workflow
"""
self.workflows[workflow_id] = workflow_definition
# Also register with atomic orchestrator
self.atomic_orchestrator.register_workflow(workflow_id, workflow_definition)
async def execute_workflow(self, workflow_id: str, input_data: Any) -> Any:
"""
Execute a workflow with the given input data.
Args:
workflow_id: ID of the workflow to execute
input_data: Input data for the workflow
Returns:
The result of the workflow execution
"""
if workflow_id not in self.workflows:
raise ValueError(f"Workflow {workflow_id} not found")
# First try to execute with atomic orchestrator if possible
try:
# Check if all agents in the workflow have atomic agents
workflow = self.workflows[workflow_id]
all_atomic = True
for step in workflow.get("steps", []):
agent_id = step.get("agent_id")
if agent_id not in self.agents:
raise ValueError(f"Agent {agent_id} not found")
agent = self.agents[agent_id]
if not (hasattr(agent, 'config') and agent.config and 'atomic_agent' in agent.config):
all_atomic = False
break
# If all agents have atomic implementations, use atomic orchestrator
if all_atomic:
return await self.atomic_orchestrator.execute_workflow(workflow_id, input_data)
except Exception as e:
print(f"Error executing workflow with atomic orchestrator: {e}")
# Fall back to legacy execution
pass
# Legacy execution path
workflow = self.workflows[workflow_id]
current_data = input_data
for step in workflow.get("steps", []):
agent_id = step.get("agent_id")
if agent_id not in self.agents:
raise ValueError(f"Agent {agent_id} not found")
agent = self.agents[agent_id]
current_data = await agent.process(current_data)
return current_dataA workflow is defined as a sequence of steps, where each step typically involves an agent processing some data. The result of each step is passed as input to the next step.
Example workflow definition:
workflow_definition = {
"steps": [
{
"agent_id": "data_agent_1",
"parameters": {
"output_format": "dict"
}
},
{
"agent_id": "search_agent_1",
"parameters": {
"max_results": 10
}
}
]
}import asyncio
from core.orchestration.orchestrator import Orchestrator
from core.messaging.message_broker import MessageBroker
from agents.data-agent.data_agent import DataAgent
from agents.search-agent.search_agent import SearchAgent
async def main():
# Create a message broker
broker = MessageBroker()
# Create an orchestrator
orchestrator = Orchestrator(broker)
# Start the orchestrator
await orchestrator.start()
try:
# Create and register agents
data_agent = DataAgent("data_agent_1", {
"output_format": "dict"
})
search_agent = SearchAgent("search_agent_1", {
"search_engine": "memory"
})
orchestrator.register_agent(data_agent)
orchestrator.register_agent(search_agent)
# Initialize agents
await data_agent.initialize()
await search_agent.initialize()
# Create a workflow
workflow_def = {
"steps": [
{
"agent_id": "data_agent_1"
},
{
"agent_id": "search_agent_1"
}
]
}
orchestrator.create_workflow("my_workflow", workflow_def)
# Execute the workflow
result = await orchestrator.execute_workflow("my_workflow", {
"data": [
{"id": 1, "text": "Sample text 1"},
{"id": 2, "text": "Sample text 2"}
]
})
print(f"Workflow result: {result}")
finally:
# Stop the orchestrator
await orchestrator.stop()
# Run the example
asyncio.run(main())When using the orchestration system:
- Always start the orchestrator before registering agents or executing workflows
- Initialize agents before registering them with the orchestrator
- Define workflows with clear steps and dependencies
- Handle exceptions during workflow execution
- Stop the orchestrator when it's no longer needed
- Use the
AgentFactoryto create agents with atomic-agents support - Leverage the adapter classes for bidirectional integration with atomic-agents
- Gradually migrate agents to use atomic-agents for improved functionality