This blog post explores how to build a multi-agent AI system inspired by the UNIX philosophy of “do one thing and do it well,” using a practical security incident response implementation as an example. It demonstrates how specialised agents (like ATO and Phishing experts) can be coordinated by a central agent to handle complex incidents, while keeping each component focused and efficient.
Ever wondered what would happen if we took the “jack of all trades, master of none” saying and completely turned it on its head? Well, that’s exactly what we’re going to explore today with multi-agent AI systems!
Remember those Swiss Army knives that tried to do everything? Sure, they’re cool, but have you ever tried actually using that tiny scissors to cut something important? Yeah, not the best experience. That’s kind of what we’re dealing with when we try to make AI models do everything at once.
The UNIX Philosophy Meets AI
The UNIX philosophy of “do one thing and do it well” still rocks after all these years. Now, what if we applied this same principle to AI agents? Instead of having one massive model trying to handle everything from phishing attacks to account takeovers (and probably dropping the ball somewhere), we could have specialised agents that are absolute experts in their domains.
Show Me The Code!
Let’s look at an example implementation. We’ll create a system with three main components:
- A coordinator agent (our traffic controller)
- An Account Takeover (ATO) specialist
- A Phishing incident expert
Here’s the cool part — each agent knows exactly what it’s good at and isn’t shy about it. Check out this base structure:
@dataclass
class AgentCapabilities:
"""Defines what an agent can handle and when it should be called."""
agent_id: str
description: str
triggers: List[str] # Keywords or patterns that should trigger this agent
examples: List[str] # Example scenarios this agent should handle
class BaseAgent(ABC):
"""Abstract base class for all agents."""
def __init__(self, agent_id: str, system_prompt: str):
self.agent_id = agent_id
self.system_prompt = system_prompt
self.client = OpenAI()
The Coordinator: The Traffic Controller of Our AI Highway
Think of our coordinator agent as that one friend who knows exactly who to call for every situation. Leaky pipe? They’ve got a plumber on speed dial. Computer issues? They know just the right tech whiz.
class CoordinatorAgent(BaseAgent):
def process_message(self, message: Message) -> Message:
# Determine which agents should handle the incident
analysis = self._call_llm(
f"Analyse this security incident and determine which agents should handle it. "
f"Explain your reasoning based on the agents' capabilities and triggers: {message.content}"
)
required_agents = self._parse_required_agents(analysis)
responses = []
for agent_id in required_agents:
if agent_id in self.agents:
response = self.agents[agent_id].process_message(message)
responses.append(response)
Why This Approach Rocks
Reduced Complexity: Each agent only needs to know about its specific domain. Our phishing expert doesn’t need to know anything about account takeovers, just like your dentist doesn’t need to know how to fix your car.
Cost Efficiency: Here’s a fun fact — by using cached prompts and specialised agents, we can significantly reduce our API costs when we use certain models (such as OpenAI’s GPT models).
Easy to Scale: Want to add a new type of security handling? Just create a new agent! It’s like adding a new tool to your toolbox — no need to rebuild the whole thing.
The Power of Combination
But here’s where it gets really interesting — what happens when we combine these specialised agents? Just like how combining grep, sed, and awkcan create powerful text processing pipelines, combining our security agents can handle complex, multi-faceted incidents.
For example, consider this incident:
incident = Message(
sender="security_monitoring",
content="""
Suspicious activity detected:
- Multiple failed login attempts from IP 192.0.2.1
- Successful login from new location
- Suspicious email forwarding rule created
- Mass email sent to external addresses
""",
metadata={"priority": "high"}
)
This single incident triggers both our ATO and Phishing agents, and the coordinator makes sure they work together seamlessly. Cool, right?
The Future of Multi-Agent Systems
Remember how microservices revolutionised how we build applications? I have a feeling specialised AI agents might do the same for AI systems. Instead of building massive, complex models that try to do everything, we might be moving towards ecosystems of specialised agents working together.
Want to implement this yourself? The full code is available here and ready to be adapted to your needs. Just remember: like any good specialist, make sure each agent is really good at its job, even if that job is something as specific as detecting suspicious email forwarding rules!
from abc import ABC, abstractmethod
from typing import List, Dict, Any, Optional, Tuple
from dataclasses import dataclass
import json
from openai import OpenAI
from datetime import datetime
@dataclass
class Message:
"""Standard message format for inter-agent communication."""
sender: str
content: str
metadata: Dict[str, Any]
timestamp: str = str(datetime.now())
@dataclass
class AgentCapabilities:
"""Defines what an agent can handle and when it should be called."""
agent_id: str
description: str
triggers: List[str] # Keywords or patterns that should trigger this agent
examples: List[str] # Example scenarios this agent should handle
class BaseAgent(ABC):
"""Abstract base class for all agents."""
def __init__(self, agent_id: str, system_prompt: str):
self.agent_id = agent_id
self.system_prompt = system_prompt
self.client = OpenAI()
@abstractmethod
def process_message(self, message: Message) -> Message:
"""Process incoming message and return response."""
pass
@abstractmethod
def get_capabilities(self) -> AgentCapabilities:
"""Return the agent's capabilities for coordinator awareness."""
pass
def _call_llm(self, prompt: str) -> str:
"""Make API call to GPT-4-mini."""
try:
response = self.client.chat.completions.create(
model="gpt-4o-mini",
messages=[
{"role": "system", "content": [{"type": "text", "text": self.system_prompt}]},
{"role": "user", "content": [{"type": "text", "text": prompt}]}
],
temperature=0.7
)
return response.choices[0].message.content
except Exception as e:
return f"Error in LLM call: {str(e)}"
class CoordinatorAgent(BaseAgent):
"""Coordinator agent that routes messages to appropriate specialised agents."""
BASE_PROMPT = """You are a highly skilled coordinator agent responsible for:
1. Analysing incoming security incidents
2. Determining which specialised agent(s) should handle the incident
3. Coordinating responses when multiple agents are needed
4. Ensuring all necessary information is shared between agents
You have access to the following specialised agents:
{agent_capabilities}
When analysing incidents, carefully consider the capabilities and triggers
of each agent to determine the most appropriate handler(s).
You can and should involve multiple agents when an incident spans multiple domains."""
def __init__(self):
self.agents: Dict[str, BaseAgent] = {}
# Initialise with empty capabilities, will be updated on agent registration
super().__init__("coordinator", self._build_system_prompt())
def _build_system_prompt(self) -> str:
"""Dynamically build system prompt based on registered agents' capabilities."""
if not self.agents:
agent_capabilities = "NO AGENTS CURRENTLY REGISTERED"
else:
capabilities_text = []
for agent in self.agents.values():
caps = agent.get_capabilities()
agent_text = f"""
Agent ID: {caps.agent_id}
Description: {caps.description}
Triggers: {', '.join(caps.triggers)}
Example scenarios:
{chr(10).join('- ' + example for example in caps.examples)}
"""
capabilities_text.append(agent_text)
agent_capabilities = "\n".join(capabilities_text)
return self.BASE_PROMPT.format(agent_capabilities=agent_capabilities)
def register_agent(self, agent: BaseAgent):
"""Register a new specialised agent and update coordinator's knowledge."""
self.agents[agent.agent_id] = agent
# Update system prompt with new agent's capabilities
self.system_prompt = self._build_system_prompt()
def process_message(self, message: Message) -> Message:
# Determine which agents should handle the incident
analysis = self._call_llm(
f"Analyse this security incident and determine which agents should handle it. "
f"Explain your reasoning based on the agents' capabilities and triggers: {message.content}"
)
required_agents = self._parse_required_agents(analysis)
responses = []
for agent_id in required_agents:
if agent_id in self.agents:
response = self.agents[agent_id].process_message(message)
responses.append(response)
coordinated_response = self._coordinate_responses(responses, analysis)
return Message(
sender="coordinator",
content=coordinated_response,
metadata={
"involved_agents": required_agents,
"analysis": analysis
}
)
def _parse_required_agents(self, analysis: str) -> List[str]:
"""Parse the LLM analysis to determine required agents."""
# Ask LLM to explicitly list required agents based on analysis
agent_list = self._call_llm(
f"Based on this analysis:\n{analysis}\n\n"
f"List only the agent IDs that should be involved, separated by commas. "
f"Available agents: {', '.join(self.agents.keys())}"
)
return [agent_id.strip() for agent_id in agent_list.split(',')]
def _coordinate_responses(self, responses: List[Message], analysis: str) -> str:
"""Coordinate and combine responses from multiple agents."""
if not responses:
return "No specialised agents were required for this incident."
combined = "\n".join([f"{r.sender}: {r.content}" for r in responses])
return self._call_llm(
f"Initial incident analysis:\n{analysis}\n\n"
f"Agent responses:\n{combined}\n\n"
f"Coordinate these responses into a single coherent response, "
f"ensuring all necessary actions from each agent are preserved."
)
def get_capabilities(self) -> AgentCapabilities:
"""Return coordinator's capabilities."""
return AgentCapabilities(
agent_id="coordinator",
description="Central coordinator that analyses incidents and routes them to appropriate specialised agents",
triggers=["all incidents"], # Coordinator handles all incidents initially
examples=[
"Any security incident that needs to be routed to specialised agents",
"Complex incidents requiring multiple agents' involvement",
"Incidents requiring coordinated response across different domains"
]
)
class ATOAgent(BaseAgent):
"""Agent specialised in handling Account Takeover (ATO) incidents."""
ATO_PROMPT = """You are an Account Takeover (ATO) incident response specialist. Your responsibilities include:
1. Analysing potential ATO incidents
2. Following the ATO response runbook
3. Implementing immediate account security measures
4. Investigating the scope of the compromise
5. Recommending remediation steps
Your runbook includes:
- Immediate account lockdown procedures
- Authentication log analysis
- Suspicious activity pattern recognition
- Account recovery procedures
- Security control recommendations
Provide detailed, actionable responses based on the ATO runbook."""
def __init__(self):
super().__init__("ato_agent", self.ATO_PROMPT)
def get_capabilities(self) -> AgentCapabilities:
return AgentCapabilities(
agent_id="ato_agent",
description="Specialises in handling Account Takeover (ATO) incidents and implementing account security measures",
triggers=[
"failed login attempts",
"suspicious login location",
"password reset",
"unusual account activity",
"credential compromise",
"account lockout",
"unauthorised access",
"session hijacking"
],
examples=[
"Multiple failed login attempts from unusual IP addresses",
"Successful login from a new geographic location",
"Unusual account activity outside business hours",
"Multiple password reset attempts",
"Session token manipulation detected"
]
)
def process_message(self, message: Message) -> Message:
response = self._call_llm(
f"Handle this ATO incident according to the runbook: {message.content}"
)
return Message(
sender="ato_agent",
content=response,
metadata={"incident_type": "ATO"}
)
class PhishingAgent(BaseAgent):
"""Agent specialised in handling phishing campaign incidents."""
PHISHING_PROMPT = """You are a Phishing Campaign response specialist. Your responsibilities include:
1. Analysing potential phishing incidents
2. Following the phishing response playbook
3. Identifying phishing campaign patterns
4. Implementing email security measures
5. Coordinating with email security systems
Your playbook includes:
- Email header analysis
- URL and attachment analysis
- Campaign pattern recognition
- Email quarantine procedures
- User notification templates
Provide detailed, actionable responses based on the phishing playbook."""
def __init__(self):
super().__init__("phishing_agent", self.PHISHING_PROMPT)
def get_capabilities(self) -> AgentCapabilities:
return AgentCapabilities(
agent_id="phishing_agent",
description="Specialises in handling phishing campaigns and email-based attacks",
triggers=[
"suspicious email",
"phishing link",
"malicious attachment",
"email campaign",
"credential harvesting",
"spoofed sender",
"mass mailing",
"suspicious forwarding rules"
],
examples=[
"Mass phishing email detected across organization",
"Suspicious attachment with executable content",
"Credential harvesting website reported",
"Unusual email forwarding rules created",
"Executive impersonation attempt detected"
]
)
def process_message(self, message: Message) -> Message:
response = self._call_llm(
f"Handle this phishing incident according to the playbook: {message.content}"
)
return Message(
sender="phishing_agent",
content=response,
metadata={"incident_type": "PHISHING"}
)
# Example usage
def main():
# Initialise agents
coordinator = CoordinatorAgent()
ato_agent = ATOAgent()
phishing_agent = PhishingAgent()
# Register specialised agents with coordinator
coordinator.register_agent(ato_agent)
coordinator.register_agent(phishing_agent)
# Example incident
incident = Message(
sender="security_monitoring",
content="""
Suspicious activity detected:
- Multiple failed login attempts from IP 192.0.2.1
- Successful login from new location
- Suspicious email forwarding rule created
- Mass email sent to external addresses
""",
metadata={"priority": "high"}
)
# Process incident
response = coordinator.process_message(incident)
print(f"Coordinator Analysis:\n{response.metadata['analysis']}\n")
print(f"Coordinated Response:\n{response.content}\n")
print(f"Involved Agents: {response.metadata['involved_agents']}")
if __name__ == "__main__":
main()
When I ran this, I got this output:
Coordinator Analysis:
In this security incident, we have several suspicious activities that can be analyzed based on the triggers of the specialized agents available.
1. **Multiple failed login attempts from IP 192.0.2.1**: This is a clear indicator of suspicious login activity, which falls under the triggers for the **ato_agent**. The agent specializes in handling Account Takeover (ATO) incidents and would be well-equipped to address the failed login attempts and assess the potential for an account compromise.
2. **Successful login from a new location**: This is also a trigger for the **ato_agent**. A successful login from an unfamiliar geographic location is a red flag that may indicate unauthorized access to an account. The agent would need to investigate this successful login in conjunction with the failed attempts.
3. **Suspicious email forwarding rule created**: This activity suggests potential email compromise or malicious intent, which falls under the triggers for the **phishing_agent**. The creation of a suspicious email forwarding rule can indicate that an attacker is trying to exfiltrate information or redirect communications, warranting the expertise of the phishing agent.
4. **Mass email sent to external addresses**: This is also a trigger for the **phishing_agent**. Sending mass emails, especially to external addresses, can be indicative of a phishing campaign or other malicious email activities. Therefore, this aspect requires investigation by the phishing agent.
Based on the analysis of the incident, we should involve both specialized agents to address the different components of the incident effectively:
- **ato_agent**: To investigate the failed login attempts and the successful login from a new location, assessing potential account takeover risks.
- **phishing_agent**: To examine the suspicious email forwarding rule and the mass email sent, as these activities suggest a possible phishing campaign or compromise of email integrity.
In summary, both agents should be involved in the response due to the multifaceted nature of the incident, ensuring that all necessary information is shared and addressed appropriately.
Coordinated Response:
### Coordinated Response to Security Incident
In response to the recent security incident involving multiple suspicious activities, we will take a comprehensive approach by combining the efforts of both the **ato_agent** and the **phishing_agent**. The coordinated response will address the various components of the incident, ensuring thorough investigation and remediation.
#### **1. Immediate Account Lockdown Procedures (ato_agent)**
- **Lock the Affected Account:**
- Immediately lock the account associated with suspicious login attempts to prevent further unauthorized access.
- Change the password and require a password reset for the user upon recovery.
- **Disable Email Forwarding Rules:**
- Access the email account settings to disable any suspicious email forwarding rules that have been created.
- Review the account's settings for any other unauthorized changes.
#### **2. Email Quarantine Procedures (phishing_agent)**
- **Quarantine Suspicious Emails:**
- Use the email security system to quarantine any emails sent to external addresses, especially those originating from the affected account.
- Temporarily block or disable the accounts involved in the forwarding rule to prevent further exfiltration of information.
#### **3. Authentication Log Analysis (ato_agent)**
- **Review Login Attempts:**
- Analyze authentication logs for all successful and failed login attempts, focusing on IP address (192.0.2.1) and any notable locations.
- **Geolocation Analysis:**
- Determine geographical locations of login attempts to assess consistency with the user's normal behavior.
#### **4. Email Header Analysis (phishing_agent)**
- **Analyze Email Headers:**
- Identify and analyze the headers of the suspicious emails to check for sender authenticity and any discrepancies in the sender's domain.
- Look for unusual routing paths or IP addresses that do not match legitimate services.
#### **5. Suspicious Activity Pattern Recognition (ato_agent & phishing_agent)**
- **Identify Patterns Across Incidents:**
- Look for patterns in login attempts, including time of day, frequency, and consistency with the user’s typical behavior.
- Document the details of the incident, including multiple failed login attempts, successful login from a new location, and the creation of suspicious email forwarding rules.
#### **6. URL and Attachment Analysis (phishing_agent)**
- **Examine URLs and Attachments:**
- Identify any URLs in the suspicious emails and use a URL scanning tool to check for links to known phishing sites.
- Examine any attachments for malware or suspicious content, utilizing antivirus software for scanning.
#### **7. Investigating the Scope of the Compromise (ato_agent)**
- **Check for Unauthorized Access:**
- Review the account for any unauthorized actions taken, such as changes to account settings or sent emails.
- Assess whether any sensitive information has been accessed or exfiltrated.
#### **8. User Notification (phishing_agent)**
- **Notify Affected Users:**
- Prepare and send notification templates to inform users about the suspicious activity detected on their accounts.
- Advise users to change their passwords immediately and review account settings for unauthorized changes, including email forwarding rules.
#### **9. Recommending Remediation Steps (ato_agent)**
- **User Education and Security Controls Implementation:**
- Inform the affected user about the incident and advise on best practices for account security, including enabling multi-factor authentication (MFA).
- Recommend implementing IP whitelisting for the affected account to limit access to known locations only.
#### **10. Coordinate with IT Security (phishing_agent)**
- **Share Findings:**
- Share findings with the IT security team for further investigation and implement additional monitoring on the affected accounts and the suspicious IP address.
#### **11. Post-Incident Review (phishing_agent)**
- **Review of Incident:**
- Conduct a review of this incident to improve detection and response strategies.
- Update the phishing response playbook if necessary based on lessons learned.
### Conclusion
By executing this coordinated response, both agents will effectively address the multifaceted nature of the incident. We will ensure that all necessary information is shared and that appropriate actions are taken to mitigate risks and prevent future occurrences. Continuous monitoring will be established to detect any further suspicious activities on the affected accounts.
Involved Agents: ['ato_agent', 'phishing_agent']
P.S. And yes, our coordinator agent is basically an AI traffic controller wearing a high-vis vest, making sure every incident gets to the right specialist. How’s that for a mental image? 😄
Disclaimer: The perspectives shared here are my own and do not necessarily represent those of my employer. I use GenAI as a tool to help me compose and structure my articles.