Agent-Wiz

KO | EN

gitlite — search

#python #tailwindcss #css #html #javascript #typescript #nodejs #nextjs #reactjs #ai #docker #aws

★ 393 Open GitHub ↗

A CLI tool for threat modeling and visualizing AI agents built using popular frameworks like LangGraph, AutoGen, CrewAI, and more.

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

<br/>

![](https://github.com/Repello-AI/Agent-Wiz/raw/master/assets/example_vis.png)

## 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

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

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

## Supported Frameworks

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)

## Security Analysis

Agent Wiz currently supports [**MAESTRO**](https://cloudsecurityalliance.org/blog/2025/02/06/agentic-ai-threat-modeling-framework-maestro) as its primary threat modeling framework. It evaluates agent workflows against the following structure:

- **M**ission: Defining the system purpose and security objectives
- **A**ssets: Inventorying critical components (agents, tools, data flows)
- **E**ntrypoints: Mapping attack surfaces and access vectors
- **S**ecurity Controls: Evaluating existing defensive measures
- **T**hreats: Identifying potential vulnerabilities and attack scenarios
- **R**isks: Calculating impact and likelihood of security events
- **O**perations: 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

```bash
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**
```bash
export OPENAI_API_KEY=sk-...
```

**Option 2: .env File (Recommended)**
```bash
cp .env.sample .env
```

## 🚀 CLI Usage

### 1. Extract Agentic Workflow

```bash
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:

```json
{
  "nodes": [...],
  "edges": [...],
  "metadata": {
    "framework": "autogen"
  }
}
```

### 2. Visualize the Agentic workflow
```bash
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

```bash
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:__
```bash
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`](./CONTRIBUTING.md) before submitting issues or PRs.

## 📜 Changelog

For recent changes and version history, see [`CHANGELOG.md`](./CHANGELOG.md).

## 📄 License

Licensed under the **Apache 2.0 License**. See [`LICENSE`](./LICENSE) for full details.

## Links

- [Agent Wiz GitHub](https://github.com/Repello-AI/Agent-Wiz)
- [Issue Tracker](https://github.com/Repello-AI/Agent-Wiz/issues)
- [PyPI Package](https://pypi.org/project/repello-agent-wiz/)

## Attribution

[Google ADK code examples](https://github.com/Repello-AI/Agent-Wiz/tree/google-adk/examples/code/google_adk/agents) are taken from [Google ADK Samples](https://github.com/google/adk-samples)

Agent-Wiz

Explore Similar Repositories

Was this content helpful?

Recent Feedback

Download README