uml / agents /code /use_case_diagram_generator.py
Mohammed Foud
Add application file
5a2d62e
from base_agent import BaseAgent
from models import AgentState
from typing import AsyncGenerator
import json
class UseCaseDiagramGenerator(BaseAgent):
async def __call__(self, state: AgentState) -> AsyncGenerator[str, None]:
prompt_template = """
Generate comprehensive PlantUML code for a use case diagram based on:
Project: {project_name}
Actors and Use Cases: {actors_use_cases}
Instructions:
1. Include all actors with proper typing:
- Primary actors on left
- Supporting actors on right
- System actors if applicable
2. Organize use cases into logical packages/subsystems when appropriate
3. Use proper relationship types:
- Association: --> (actor to use case)
- Include: .> (base use case includes another)
- Extend: .> (use case extends another with conditions)
4. Apply stereotypes where helpful (<<system>>, <<external>>, etc.)
5. Include notes for complex relationships when needed
6. Use left-to-right direction for better readability
7. Format with clear alignment and spacing
8. Include only raw PlantUML code without explanations
Example Structure (DO NOT TREAT AS VARIABLES):
@startuml
left to right direction
skinparam packageStyle rectangle
actor "Primary Actor" <<Human>> as user
actor "Supporting System" <<System>> as legacy
rectangle "Order Processing" {{
usecase "Place Order" as UC1
usecase "Process Payment" as UC2
usecase "Validate Order" as UC3
UC1 .> UC2 : include
UC1 .> UC3 : include
}}
rectangle "Inventory" {{
usecase "Check Stock" as UC4
}}
user --> UC1
legacy --> UC3
note right of UC1
This use case initiates the
order fulfillment workflow
end note
@enduml
"""
async for chunk in self._stream_process(
state=state,
prompt_template=prompt_template,
output_key="use_case_diagram",
step_name="generate_use_case_diagram",
actors_use_cases=state["actors_use_cases"],
project_name=state["project_name"]
):
yield chunk