File size: 2,682 Bytes
5a2d62e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 |
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
|