| 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 | |