| from base_agent import BaseAgent | |
| from models import AgentState | |
| from typing import AsyncGenerator | |
| import json | |
| class SequenceDiagramGenerator(BaseAgent): | |
| async def __call__(self, state: AgentState) -> AsyncGenerator[str, None]: | |
| prompt_template = """ | |
| Generate comprehensive PlantUML code for a sequence diagram based on: | |
| Workflow Analysis: {sequence_interactions} | |
| Instructions: | |
| 1. Use proper participant types: | |
| - actor for human users | |
| - boundary for UI elements | |
| - control for controllers | |
| - entity for domain objects | |
| - database for storage | |
| 2. Include all message types: | |
| - Synchronous: -> | |
| - Asynchronous: ->> | |
| - Return: --> | |
| - Self: -> self | |
| 3. Add activation bars with: | |
| - activate/deactivate | |
| - Nested activations where needed | |
| 4. Use advanced features: | |
| - Groups (alt/opt/loop/par) | |
| - Notes | |
| - Dividers (== Sections ==) | |
| - Stereotypes <<component>> | |
| - Arrows with: | |
| * Lost messages (->x) | |
| * Found messages (x->) | |
| * Parallel (->||) | |
| 5. Format for readability: | |
| - Proper indentation | |
| - Section comments | |
| - Color coding where helpful | |
| Example Structure (DO NOT TREAT AS VARIABLES): | |
| @startuml | |
| skinparam style strictuml | |
| skinparam lifelineStrategy solid | |
| actor Customer <<Human>> | |
| boundary "Web UI" as UI | |
| control OrderController | |
| entity Order | |
| database OrderDB | |
| == Order Submission == | |
| Customer -> UI: Submit Order | |
| activate UI | |
| UI -> OrderController: processOrder() | |
| activate OrderController | |
| group Payment Processing | |
| OrderController -> PaymentService <<External>>: authorizePayment() | |
| activate PaymentService | |
| PaymentService --> OrderController: status | |
| deactivate PaymentService | |
| end | |
| alt payment approved | |
| OrderController -> OrderDB: persistOrder() | |
| OrderController --> UI: confirmation | |
| else payment declined | |
| OrderController --> UI: error | |
| end | |
| deactivate OrderController | |
| UI --> Customer: Order Status | |
| deactivate UI | |
| @enduml | |
| """ | |
| async for chunk in self._stream_process( | |
| state=state, | |
| prompt_template=prompt_template, | |
| output_key="sequence_diagram", | |
| step_name="generate_sequence_diagram", | |
| sequence_interactions=state["sequence_interactions"] | |
| ): | |
| yield chunk | |