Agent-Wiz

Agent-Wiz copied to clipboard

Overview

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.

Why Agent Wiz?

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

Core Features

https://github.com/user-attachments/assets/40231eae-9716-421e-a005-fd55d1d8cc71

Supported Frameworks

The following agent orchestration frameworks are currently supported:

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)

Security Analysis

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.

Installation

pip install repello-agent-wiz

Prerequisites

Before running any analysis commands, you must configure your OpenAI API key. You can do this in two ways:

Option 1: Environment Variable

export OPENAI_API_KEY=sk-...

Option 2: .env File (Recommended)

cp .env.sample .env

🚀 CLI Usage

1. Extract Agentic Workflow

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"
  }
}

2. Visualize the Agentic workflow

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.

3. Analyze against Threat Modeling

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

📈 Roadmap

Planned features (Not in any paricular order)

• [x] Build parsers for major agentic frameworks (Autogen, LangGraph, CrewAI, etc.)
• [x] Generate standardized JSON graph representations of agent flows
• [x] CLI interfaces
• [x] Security report generation
• [ ] Extend to STRIDE, PASTA, LINDDUN, etc.
• [ ] Agent simulation-based threat exploration

🤝 Contributing

We welcome contributions of all kinds!

⚠️ Please read CONTRIBUTING.md before submitting issues or PRs.

📜 Changelog

For recent changes and version history, see CHANGELOG.md.

📄 License

Licensed under the Apache 2.0 License. See LICENSE for full details.

Links

• Agent Wiz GitHub
• Issue Tracker
• PyPI Package

Attribution

Google ADK code examples are taken from Google ADK Samples

© 2025 Repello AI | Website