| from base_agent import BaseAgent | |
| from models import AgentState | |
| from typing import AsyncGenerator | |
| import json | |
| class TimingDiagramGenerator(BaseAgent): | |
| async def __call__(self, state: AgentState) -> AsyncGenerator[str, None]: | |
| prompt_template = """ | |
| Generate PlantUML code for a timing diagram based on the following project: | |
| Project: {project_name} | |
| Description: {project_description} | |
| Sequence: {sequence_interactions} | |
| Instructions: | |
| 1. Show the timing constraints between components | |
| 2. Include lifelines and state changes over time | |
| 3. Use standard PlantUML timing diagram syntax | |
| 4. Don't include any explanation, just the PlantUML code | |
| Example format (DO NOT TREAT THIS AS VARIABLES): | |
| '''plantuml | |
| @startuml | |
| clock "Clock" as C with period 1 | |
| binary "Signal" as S | |
| @0 | |
| S is low | |
| @5 | |
| S is high | |
| @10 | |
| S is low | |
| @enduml | |
| ''' | |
| """ | |
| async for chunk in self._stream_process( | |
| state=state, | |
| prompt_template=prompt_template, | |
| output_key="timing_diagram", | |
| step_name="generate_timing_diagram", | |
| project_name=state["project_name"], | |
| project_description=state["project_description"], | |
| sequence_interactions=state["sequence_interactions"] | |
| ): | |
| yield chunk | |