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