File size: 1,476 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 |
from base_agent import BaseAgent
from models import AgentState
from typing import AsyncGenerator
import json
class ObjectDiagramGenerator(BaseAgent):
async def __call__(self, state: AgentState) -> AsyncGenerator[str, None]:
prompt_template = """
Generate PlantUML code for an object diagram based on the following project:
Project: {project_name}
Description: {project_description}
Entities: {entities_classes}
Instructions:
1. Show object instances and their relationships
2. Include attribute values for a specific scenario
3. Use standard PlantUML object diagram syntax
4. Don't include any explanation, just the PlantUML code
Example format (DO NOT TREAT THIS AS VARIABLES):
'''plantuml
@startuml
object "Object1" as o1 {{
attribute1 = value1
attribute2 = value2
}}
object "Object2" as o2 {{
attribute3 = value3
}}
o1 --> o2 : relationship
@enduml
'''
"""
async for chunk in self._stream_process(
state=state,
prompt_template=prompt_template,
output_key="object_diagram",
step_name="generate_object_diagram",
project_name=state["project_name"],
project_description=state["project_description"],
entities_classes=state["entities_classes"]
):
yield chunk
|