| from base_agent import BaseAgent | |
| from models import AgentState | |
| from typing import AsyncGenerator | |
| import json | |
| class ClassDiagramGenerator(BaseAgent): | |
| async def __call__(self, state: AgentState) -> AsyncGenerator[str, None]: | |
| prompt_template = """ | |
| Generate comprehensive PlantUML code for a class diagram based on the following domain analysis: | |
| {entities_classes} | |
| Instructions: | |
| 1. Include all classes with their complete attributes (visibility, type, constraints) | |
| 2. Use proper UML relationships with correct cardinality and labels: | |
| - Inheritance: <|-- | |
| - Implementation: <|.. | |
| - Composition: *-- | |
| - Aggregation: o-- | |
| - Dependency: --> | |
| 3. For complex relationships, use association classes when needed | |
| 4. Include proper visibility markers: | |
| - Public: + | |
| - Private: - | |
| - Protected: # | |
| - Package: ~ | |
| 5. Use advanced PlantUML features where appropriate: | |
| - Abstract classes/interfaces | |
| - Enumerations | |
| - Stereotypes | |
| - Notes and constraints | |
| - Method signatures | |
| 6. Organize classes into packages if the domain suggests logical groupings | |
| 7. Use proper formatting with alignment and spacing | |
| 8. Include only the raw PlantUML code without explanations | |
| Example Structure (DO NOT TREAT AS VARIABLES): | |
| @startuml | |
| package "Domain Model" {{ | |
| abstract class AbstractEntity {{ | |
| + id: UUID {{readonly}} | |
| # createdAt: DateTime | |
| + equals(other: Object): boolean {{abstract}} | |
| }} | |
| class User {{ | |
| + username: String {{unique}} | |
| + email: String | |
| - passwordHash: String | |
| + authenticate(password: String): boolean | |
| }} | |
| enum UserRole {{ | |
| ADMIN | |
| MEMBER | |
| GUEST | |
| }} | |
| User "1" *-- "0..*" Order : places > | |
| User o-- UserRole : has > | |
| (User, Order) .. OrderAssociation | |
| class OrderAssociation {{ | |
| + orderDate: DateTime | |
| + status: OrderStatus | |
| }} | |
| }} | |
| note top of User | |
| User represents all system users | |
| with authentication capabilities | |
| end note | |
| @enduml | |
| """ | |
| async for chunk in self._stream_process( | |
| state=state, | |
| prompt_template=prompt_template, | |
| output_key="class_diagram", | |
| step_name="generate_class_diagram", | |
| entities_classes=state["entities_classes"] | |
| ): | |
| yield chunk | |