🧩AI Agent Based UML Diagram Generator
Agents help me to have my cake and eat it too.
As a Software and Data Architect, a significant part of my role involves creating extensive documentation; This includes designing myriad diagrams that illustrate system architectures, data flows, and intricate interactions between various components. These visual representations serve as a blueprint, providing a clear and structured view of how the system is expected to function from end to end.
It also means (or meant) that a considerable amount of time was spent creating such diagrams by hand. Drawling them meticulously in a tool like Visio or Draw.IO — this took up a lot of my time.
I have always wished for a faster and more effortless way to create such diagrams — A magic wand that I could wave, make a wish, and voilà! — a fully formed sequence diagram would instantly appear, ready to be shared with my team and stakeholders.
However, reality is far from magical. In the real world, as architects, we painstakingly craft documentation to give structure and credibility to our designs. This process requires meticulous attention to detail, ensuring that every diagram accurately reflects the system’s architecture, workflows, and dependencies. It’s a time-consuming yet essential task that bridges the gap between an idea and its execution.
At least that was before.
Before what?
Generative AI and the power of LLM Based Agentic Systems
Let’s not pretend like Generative AI isn’t a game changer. Sure, people are still finding credible use cases to apply Large Language Models but we can all agree that for certain use cases, LLMs are a productive and viable option.
2 Clear use cases among the winners are
1️⃣ Coding
2️⃣ Outlining requirements and diagramming
I say 2 because I have been experimenting with LLMs to see how good they are at understanding requirements in terms of user stories.
These experiments have led me to design an Agentic framework to create UML diagrams in Plant UML code.
UML diagrams, such as Class Diagrams, Sequence Diagrams, and others, serve as invaluable tools for visually representing system architectures and design flows. In fact, I would argue that they provide an essential perspective on a system — whether during its initial conceptualization or post-implementation.
Comprehensive documentation proves indispensable in the long run, particularly in dynamic development environments where teams experience constant flux; Old engineers leave to give way to new ones; Legacy systems accumulate technical debt.
A well-maintained set of UML diagrams not only enhances clarity but also ensures continuity, facilitating smoother transitions and more effective system understanding and subsequently — maintenance.
Ultimately, the significance of UML diagrams cannot be overstated — they offer an intuitive and structured approach to documenting complex systems, making them an essential asset for architects, developers, and stakeholders alike.
Thus, I decided to make my life easier and implement a small project that will use Atomic Agents to generate System diagrams for me. This project should help you too, so I am outlining the code here for you so that you can also create System diagrams quickly to save time and beat Rush hour.
Why spend time in traffic when you can spend it being terrific!
The Code
Let’s start with the requirements
▶️ Python ≥ 3.11
▶️ Atomic Agents
▶️ Open AI api
This is what the requirements.txt file will look like —
atomic-agents
openai
python-dotenvDotenv to load the Open AI api key
I have User Story file to feed user stories to the agents so that they can determine what UML diagram is to be created and also output the Plant UML code to a file. This is what I have within the userstory.txt file
Story 1:
As a user of XYZ Accounting application, I want to be able to enter my registered user id and password on a login screen so that I can log into the application.
Acceptance Criteria: On successful login, I will be issued a token to enable sso and I will be taken to the application dashboard.
Unsuccesful attempt with display an authentication error in a toast message on the login screen itself.
Story 2:
As a user of XYZ Accounting application, I want to see my account details in the dashboard.
Acceptance Criteria: On successful login, my account details will be displayed, including masked account number, account type, and account balance.This is what the directory structure of the project looks like —
Agents folder has the Agent definitions.
Main.py has the invocation code for the flow to complete.
Let’s talk about the Agents first. Using Atomic Agents, I have created 2 agents.
Agent 1 — The Analyzer
Analyse the user story and determine which sort of diagram would best describe the user story. Feed the instructions to the UML Generator agent to generate the code for the diagrams.
This is the code for the agent analyzer.py
import instructor
import openai
from pydantic import Field
from atomic_agents.agents.base_agent import BaseIOSchema, BaseAgent, BaseAgentConfig
from atomic_agents.lib.components.system_prompt_generator import SystemPromptGenerator
from .config import ChatConfig
class ChoiceAgentInputSchema(BaseIOSchema):
"""Input schema for the ChoiceAgent."""
user_message: str = Field(..., description="The user stories to be analyzed")
decision_type: str = Field(..., description="The type of UML diagram to be generated for each user story")
class ChoiceAgentOutputSchema(BaseIOSchema):
"""Output schema for the ChoiceAgent."""
reasoning: str = Field(..., description="Instructions on the UML diagrams to be generated")
decision: bool = Field(..., description="The type of Diagrams and PlantUML code to be generated")
choice_agent = BaseAgent(
BaseAgentConfig(
client=instructor.from_openai(openai.OpenAI(api_key=ChatConfig.api_key)),
model=ChatConfig.model,
system_prompt_generator=SystemPromptGenerator(
background=[
"You are a decision-making agent that examines the user's input and determines the next course of action.",
"Your primary role is to analyze the user input and determine what sort of Plant UML diagram would be most appropriate for each user story.",
"You must output instructions for the next agent on which Plant UML diagrams to generate based on the user stories.",
],
steps=[
"1. Analyze the user stories to determine the nature of UML diagrams requested",
"2. Determine the type of Plant UML diagrams to generate",
"3. Instruct the next agent on the type of Plant UML diagrams and code to be generated.",
],
output_instructions=[
"You should give clear instructions to the next agent on Which Plant UML diagrams and Plant UML codes are to be generated.",
],
),
input_schema=ChoiceAgentInputSchema,
output_schema=ChoiceAgentOutputSchema,
)
)The agent has been designed according to the specifications from Atomic Agents —
Agent 2 — The UML Generator
The second agent takes the instructions from the first agent (the analyzer) and generates the actual Plat UML code, then writes it to an output file. This is the definition for the UML Generator agent umlgenerator.py
from .config import UMLGeneratorConfig
class UMLGenerator:
def __init__(self, config: UMLGeneratorConfig):
self.config = config
def read_user_story(self) -> str:
with open(self.config.input_file, 'r') as file:
return file.read()
def generate_uml(self, instructions: str) -> str:
user_story = self.read_user_story()
response = self.config.client.chat.completions.create(
model=self.config.model,
temperature=self.config.temperature,
messages=[
{"role": "system", "content": """You are a UML diagram expert. Convert user stories into PlantUML code.
Focus on creating clear and accurate class diagrams, sequence diagrams or activity diagrams."""},
{"role": "user", "content": f"Instructions: {instructions}\nUser Story:\n{user_story}"}
]
)
uml_code = response.choices[0].message.content
# Save the PlantUML code
with open(self.config.output_file, 'w') as file:
file.write(uml_code)
return uml_code Within the Agents directory, we are also maintaining a configuration in the file — config.py
import os
from dataclasses import dataclass
from typing import Set, List, Any
from dotenv import load_dotenv
from pydantic import BaseModel, ConfigDict
from openai import OpenAI
load_dotenv()
def get_api_key() -> str:
"""Retrieve API key from environment or raise error"""
api_key = os.getenv("OPENAI_API_KEY")
if not api_key:
raise ValueError("API key not found. Please set the OPENAI_API_KEY environment variable.")
return api_key
@dataclass
class ChatConfig:
"""Configuration for the chat application"""
api_key: str = get_api_key() # This becomes a class variable
model: str = "gpt-4o-mini"
exit_commands: Set[str] = frozenset({"/exit", "/quit"})
def __init__(self):
# Prevent instantiation
raise TypeError("ChatConfig is not meant to be instantiated")
class AgentConfig(BaseModel):
model_config = ConfigDict(arbitrary_types_allowed=True)
client: OpenAI
model: str
system_prompt: str
temperature: float = 0.7
class UMLGeneratorConfig(AgentConfig):
input_file: str = "userstory.txt"
output_file: str = "diagram.puml"
temperature: float = 0.7In the config file, we initialize the Open AI client and also mention the source file for the user stories as well as the output file to write the result to.
Main
Next, we have the main.py file where all the magic happens —
from openai import OpenAI
import os
from dotenv import load_dotenv
from pydantic import BaseModel, Field
from typing import List
from agent import BaseAgent, BaseAgentConfig, SystemPromptGenerator, AgentMemory
from agents.umlgenerator import UMLGenerator
from agents.analyzer import choice_agent, ChoiceAgentInputSchema
from agents.config import UMLGeneratorConfig
load_dotenv()
client = OpenAI(
api_key=os.getenv("OPENAI_API_KEY")
)
# Define a custom output schema for the main agent
class CustomOutputSchema(BaseModel):
"""Schema for the main agent's output"""
chat_message: str = Field(..., description="Instructions for UML generation")
suggested_questions: List[str] = Field(..., description="Generate the appropriate PlantUML code with annotations")
# Set up the system prompt for the main agent
system_prompt_generator = SystemPromptGenerator(
background=[
"You are a UML analysis expert that helps convert user requirements into PlantUML code.",
"Your role is to analyze user stories and provide corresponsing diagram recommendations as well as the corresponding PlantUML code."
],
steps=[
"1. Analyze the user stories carefully",
"2. Identify the key components, relationships, and behaviors",
"3. Determine the most appropriate type of UML diagram for each user story",
"4. Generate the appropriate PlantUML code with annotations for each user story"
],
output_instructions=[
"Provide specific instructions about what elements should be included in the UML diagram",
"Specify the type of diagram (class, sequence, or activity) that would best represent the user story",
"Include suggestions for layout and organization of the diagram",
"Generate the appropriate PlantUML code with annotations"
]
)
def process_user_story(user_story_file: str = "userstory.txt") -> None:
# Read the user story
try:
with open(user_story_file, 'r') as file:
user_story = file.read()
except FileNotFoundError:
print(f"Error: {user_story_file} not found. Please create this file with your user story.")
return
# First, use the choice agent to analyze and determine diagram type
choice_response = choice_agent.run(
ChoiceAgentInputSchema(
user_message=user_story,
decision_type="uml_diagram_type"
)
)
print(f"\nAnalysis Results:\n{choice_response.reasoning}")
print(f"Recommended Diagram Type: {choice_response.decision}\n")
# Create config for UML generator
generator_config = UMLGeneratorConfig(
client=client,
model="gpt-4o",
system_prompt=f"""You are a UML expert. Generate clear and accurate {choice_response.decision} diagrams.
Follow the analysis and create appropriate PlantUML code.""",
input_file=user_story_file,
output_file="diagram.puml"
)
# Initialize the UML generator agent
uml_generator = UMLGenerator(config=generator_config)
# Generate the UML diagram based on the analysis
uml_code = uml_generator.generate_uml(
f"""Analysis: {choice_response.reasoning}
User Story: {user_story}
Please generate a {choice_response.decision} diagram using PlantUML syntax."""
)
print("\nGenerated PlantUML code:")
print("------------------------")
print(uml_code)
# The diagram will be saved to the output file specified in config
print(f"\nDiagram has been saved to {generator_config.output_file}")
if __name__ == "__main__":
try:
process_user_story()
except FileNotFoundError:
print("Error: userstory.txt file not found. Please create this file with your user story.")
except Exception as e:
print(f"An error occurred: {str(e)}")Running the Code
Running the agentic workflow yields the following output in the console—
Analysis Results:
The user stories describe functionalities related to user authentication and account details display in an application. For Story 1, a Use Case Diagram is appropriate to illustrate the login process and the interactions between the user and the system. For Story 2, a Class Diagram can be used to represent the structure of the account details that will be displayed on the dashboard, including attributes like account number, account type, and account balance.
Recommended Diagram Type: True
Generated PlantUML code:
------------------------
To address the user stories provided, I will create two diagrams using PlantUML syntax: a Use Case Diagram for Story 1, and a Class Diagram for Story 2.
### Story 1: Use Case Diagram
```plantuml
@startuml
actor User
usecase "Enter User ID and Password" as UC1
usecase "Authenticate User" as UC2
usecase "Issue Token for SSO" as UC3
usecase "Display Dashboard" as UC4
usecase "Display Authentication Error" as UC5
User --> UC1
UC1 --> UC2
UC2 --> UC3 : <<success>>
UC2 --> UC5 : <<failure>>
UC3 --> UC4
@enduml
```
### Story 2: Class Diagram
```plantuml
@startuml
class Dashboard {
+displayAccountDetails()
}
class AccountDetails {
-maskedAccountNumber: String
-accountType: String
-accountBalance: float
}
Dashboard "1" --> "1" AccountDetails
@enduml
```
These diagrams represent the interactions and structure as described in the user stories. The Use Case Diagram outlines the login process, while the Class Diagram details the structure of account information displayed in the dashboard.
Diagram has been saved to diagram.pumlAnd the following UML has been written to the output file — diagram.puml
To address the user stories provided, I will create two diagrams using PlantUML syntax: a Use Case Diagram for Story 1, and a Class Diagram for Story 2.
### Story 1: Use Case Diagram
```plantuml
@startuml
actor User
usecase "Enter User ID and Password" as UC1
usecase "Authenticate User" as UC2
usecase "Issue Token for SSO" as UC3
usecase "Display Dashboard" as UC4
usecase "Display Authentication Error" as UC5
User --> UC1
UC1 --> UC2
UC2 --> UC3 : <<success>>
UC2 --> UC5 : <<failure>>
UC3 --> UC4
@enduml
```
### Story 2: Class Diagram
```plantuml
@startuml
class Dashboard {
+displayAccountDetails()
}
class AccountDetails {
-maskedAccountNumber: String
-accountType: String
-accountBalance: float
}
Dashboard "1" --> "1" AccountDetails
@enduml
```
These diagrams represent the interactions and structure as described in the user stories. The Use Case Diagram outlines the login process, while the Class Diagram details the structure of account information displayed in the dashboard.Of course this is Plant UML code which can be visualized as a diagram in the PlanText website’s editor
Diagram 1 — The Use Case diagram for Login flow
Diagram 2 — The class diagram for Accounts Dashboard
And voila, my user stories now live and breathe as actual UML diagrams. Wasn’t that easy?
Some Observations
Agents Improvement
Should definitely build a RAG to provide existing UML diagram context to my AI agents. If the agents would have access to actual (working production grade) documentation, I’m sure the results would be quite phenomenal.
Also, it would help to build a human in the loop mechanism here to give the agents feedback.
Project Enhancements
Can definitely enhance the project to start looking at some actual code also and generate UML diagrams. That would be a great way to create documentation around legacy systems.
Choice of LLM provider
I have used Open AI 4o-mini but you can use just about any AI provider by updating the config.py file.
Follow me Ritesh Shergill
for more articles on
🤖AI/ML
👨💻 Tech
👩🎓 Career advice
📲 User Experience
🏆 Leadership
I also do
✅ Career Guidance counselling — https://topmate.io/ritesh_shergill/149890
✅ Mentor Startups as a Fractional CTO — https://topmate.io/ritesh_shergill/193786
