from base_agent import BaseAgent
from models import AgentState
from typing import AsyncGenerator
import json

class StateAnalysis(BaseAgent):
    async def __call__(self, state: AgentState) -> AsyncGenerator[str, None]:
        prompt_template = """
        Analyze the following project and identify state transitions for key entities:

        Project: {project_name}
        Description: {project_description}
        Entities: {entities_classes}
        Sequence: {sequence_interactions}

        Instructions:
        1. For each key entity, identify its possible states
        2. List the events that trigger state transitions
        3. Format as a structured list of states and transitions
        """

        async for chunk in self._stream_process(
            state=state,
            prompt_template=prompt_template,
            output_key="state_transitions",
            step_name="extract_state_transitions",
            project_name=state["project_name"],
            project_description=state["project_description"],
            entities_classes=state["entities_classes"],
            sequence_interactions=state["sequence_interactions"]
        ):
            yield chunk