LLM-Based UML Generation Agent

A Python agent that parses C++ source files with Tree-sitter, lets an LLM generate PlantUML code, validates it with plantuml.jar, and outputs a .puml diagram.

Prerequisites

Item	Version / Notes
Python	(virtual environment recommended)
Java	JRE/JDK 8+ (required by plantuml.jar)
OpenAI API key	Set in .env file
PlantUML JAR	Already included as plantuml.jar in the repo

Quick Start

# 1️⃣ Clone the repo (if not already)
git clone <repo-url>
cd uml-agent

# 2️⃣ Create & activate a virtual environment
python -m venv .venv
source .venv/bin/activate

# 3️⃣ Install dependencies
pip install -r requirements.txt

# 4️⃣ Create a .env file
echo "OPENAI_API_KEY=sk-..." > .env

# 5️⃣ Run the agent on a C++ file

# Generate Class Diagram (Default)
python -m src.main example.cpp

# Generate Sequence Diagram
# Use --target to specify the starting function (e.g. main, User::login)
python -m src.main example.cpp --type sequence --target main

Expected output

• example.puml will be generated in the project root.
• The console prints a short status message:

Analyzing example.cpp...
Successfully generated PlantUML to example.puml

Rendering the diagram

java -jar plantuml.jar example.puml

This creates example.png (or .svg depending on PlantUML settings) in the same directory.

Running the test suite

⚠️ Important: Always activate the virtual environment before running tests!

# Activate virtual environment first
source .venv/bin/activate

# Run tests
python -m pytest tests/

Or in one line:

source .venv/bin/activate && python -m pytest tests/

All tests should pass:

============ 4 passed in X.XXs ============

• tests/test_analyzer.py checks class and inheritance extraction.
• tests/test_validator.py verifies PlantUML syntax validation.

Project Structure

uml-agent/
├── src/
│   ├── __init__.py
│   ├── analyzer.py      # Tree-sitter based C++ parser
│   ├── validator.py     # PlantUML syntax validator
│   ├── workflow.py      # LangGraph state machine
│   └── main.py          # CLI entry point
├── tests/
│   ├── test_analyzer.py
│   └── test_validator.py
├── example.cpp          # Sample source file
├── plantuml.jar         # PlantUML runtime
├── requirements.txt
└── README.md            # This file

FAQ

Q: "plantuml" command not found

A: The repository ships plantuml.jar. Use java -jar plantuml.jar … or install PlantUML globally (brew install plantuml on macOS).

Q: "ModuleNotFoundError: No module named 'src'" when running tests

A: You forgot to activate the virtual environment. Run:

source .venv/bin/activate

Q: How to change the LLM provider?

A: Update src/workflow.py to use ChatAnthropic, AzureOpenAI, etc., and set the corresponding environment variables.