Agent Wiz is a Python CLI for extracting agentic workflows from popular AI frameworks and performing automated threat assessments using established threat modeling methodologies. Built for developers, researchers, and security teams - Agent Wiz brings visibility to complex LLM-based orchestration to visualize flows, map tool/agent interactions, and generate actionable security reports.
In modern LLM-powered systems, agentic workflows are becoming increasingly complex, often involving multiple autonomous agents, tools, and inter-agent communication chains. Agent Wiz helps you bring:
- Visibility: Clearly visualize complex agent graphs without manual tracing
- Structure: Map relationships between agents, tools, and data flows
- Security: Apply threat modeling frameworks to identify potential vulnerabilities
Capability | Description |
---|---|
Workflow Extraction | Extract agent-based workflows from code using AST-based static parsing |
Threat Vector Visualization | View agent-to-agent, agent-to-tool, and chained connections in an interactive graph |
Automated Threat Assessment | Generate comprehensive threat assessment report using established threat modeling frameworks for AI agents like MAESTRO |
Framework Agnostic | Works with all major LLM orchestration frameworks |
Developer Friendly | Simple CLI, extensible SDK, and clean JSON exports |
agent_wiz.mp4
The following agent orchestration frameworks are currently supported:
Framework | Status |
---|---|
Autogen (core) | ✅ |
AgentChat | ✅ |
CrewAI | ✅ |
LangGraph | ✅ |
LlamaIndex | ✅ |
n8n | ✅ |
OpenAI Agents | ✅ |
Pydantic-AI | ✅ |
Swarm | ✅ |
Google-ADK | ✅ |
Each framework has its own AST-based static parser to extract:
- Agents (class/function-based)
- Tool functions
- Agent-to-agent transitions
- Tool call chains
- Group agents (e.g., selector, round-robin)
Agent Wiz currently supports MAESTRO as its primary threat modeling framework. It evaluates agent workflows against the following structure:
- Mission: Defining the system purpose and security objectives
- Assets: Inventorying critical components (agents, tools, data flows)
- Entrypoints: Mapping attack surfaces and access vectors
- Security Controls: Evaluating existing defensive measures
- Threats: Identifying potential vulnerabilities and attack scenarios
- Risks: Calculating impact and likelihood of security events
- Operations: Assessing runtime security considerations
Sample threat modelling report generated:
You can also add this line to your .bashrc
, .zshrc
, or environment setup script for persistent use.
🧪 More threat models analysis (STRIDE, PASTA, LINDDUN, etc.) are under development.
pip install repello-agent-wiz
Before running any analysis commands, you must set your OpenAI API key as an environment variable:
export OPENAI_API_KEY=sk-...
agent-wiz extract --framework agent_chat --directory ./examples/code/agent_chat --output agentchat_graph.json
This will generate a graph JSON with the following structure:
{
"nodes": [...],
"edges": [...],
"metadata": {
"framework": "autogen"
}
}
agent-wiz visualize --input agentchat_graph.json --open
This will generate an html d3 based visualisation of the agentic workflow. The open
flag (optional) and automatically opens the visualization in your default browser.
agent-wiz analyze --input agentchat_graph.json
This will generate a report like: autogen_report.md
based on the provided graph and threat modeling frameworks.
Run agent-wiz --help for more info:
usage: agent-wiz [-h] {extract,analyze,visualize} ...
Agent Wiz CLI: Extract, Analyze, Visualize agentic workflows.
positional arguments:
{extract,analyze,visualize}
extract Extract graph from source code
analyze Run threat modeling analysis on extracted graph
visualize Generate HTML visualization from graph JSON
options:
-h, --help show this help message and exit
Planned features (Not in any paricular order)
- Build parsers for major agentic frameworks (Autogen, LangGraph, CrewAI, etc.)
- Generate standardized JSON graph representations of agent flows
- CLI interfaces
- Security report generation
- Extend to STRIDE, PASTA, LINDDUN, etc.
- Agent simulation-based threat exploration
We welcome contributions of all kinds!
CONTRIBUTING.md
before submitting issues or PRs.
For recent changes and version history, see CHANGELOG.md
.
Licensed under the Apache 2.0 License. See LICENSE
for full details.
Google ADK code examples are taken from Google ADK Samples
© 2025 Repello AI | Website