Flow

@classmethod
def default(cls) -> "WaldiezFlow":
    """Get the default flow.

    Returns
    -------
    WaldiezFlow
        The default flow.
    """
    an_id = get_id()
    return cls(
        id=an_id,
        storage_id=an_id,
        created_at=now(),
        updated_at=now(),
        type="flow",
        name="Default Flow",
        description="Default Flow",
        tags=[],
        requirements=[],
        data=WaldiezFlowData.default(),
    )

def get_agent_by_id(self, agent_id: str) -> WaldiezAgent:
    """Get the agent by ID.

    Parameters
    ----------
    agent_id : str
        The ID of the agent.

    Returns
    -------
    WaldiezAgent
        The agent.

    Raises
    ------
    ValueError
        If the agent with the given ID is not found.
    """
    for agent in self.data.agents.members:
        if agent.id == agent_id:
            return agent
    raise ValueError(f"Agent with ID {agent_id} not found.")

def get_agent_connections(
    self, agent_id: str, all_chats: bool = True
) -> list[str]:
    """Get the agent connections.

    Parameters
    ----------
    agent_id : str
        The ID of the agent.
    all_chats : bool, optional
        If True, get the connections from all the chats, otherwise
        get the connections from the ordered flow (main chat flow).

    Returns
    -------
    list[str]
        The list of agent ids that the agent with the given ID connects to.
    """
    connections: list[str] = []
    if all_chats:
        for chat in self.data.chats:
            if chat.source == agent_id:
                connections.append(chat.target)
            if chat.target == agent_id:
                connections.append(chat.source)
    else:
        for entry in self.ordered_flow:
            source = entry["source"]
            target = entry["target"]
            if source.id == agent_id:
                connections.append(target.id)
            if target.id == agent_id:
                connections.append(source.id)
    return connections

def get_group_chat_members(
    self, group_manager_id: str
) -> list[WaldiezAgent]:
    """Get the group chat members.

    Parameters
    ----------
    group_manager_id : str
        The ID of the group manager.

    Returns
    -------
    list[WaldiezAgent]
        The list of group chat members.
    """
    agent = self.get_agent_by_id(group_manager_id)
    if not agent.is_group_manager or not isinstance(
        agent, WaldiezGroupManager
    ):
        return []
    members = [
        agent
        for agent in self.data.agents.members
        if agent.data.parent_id == group_manager_id
    ]
    if agent.data.speakers.selection_method != "round_robin":
        return members
    ordered_ids = agent.get_speakers_order()
    if not ordered_ids:
        return members
    members_dict = {member.id: member for member in members}
    ordered_members = [
        members_dict[member_id]
        for member_id in ordered_ids
        if member_id in members_dict
    ]
    return ordered_members

def get_group_members(self, group_id: str) -> list[WaldiezAgent]:
    """Get the group members.

    Parameters
    ----------
    group_id : str
        The ID of the group.

    Returns
    -------
    list[WaldiezAgent]
        The list of group members.
    """
    return [
        agent
        for agent in self.data.agents.members
        if agent.data.parent_id == group_id
    ]

def get_root_group_manager(self) -> WaldiezGroupManager:
    """Get the root group manager.

    Returns
    -------
    WaldiezGroupManager
        The root group manager.

    Raises
    ------
    ValueError
        If no group manager is found.
    """
    for agent in self.data.agents.groupManagerAgents:
        if agent.data.parent_id is None:
            return agent
    raise ValueError("No group manager found.")

@model_validator(mode="after")
def validate_flow(self) -> Self:
    """Flow validation.

    - unique node ids
    - there are at least two agents
        - (or a single agent but not a group manager)
    - all the agents connect to at least one other agent
    - all the linked agent tools are found in the flow
    - all the linked agent models are found in the flow
    - all the managers have at least one member in the chat group
    - the ordered flow (chats with position >=0) is not empty
    - all agents' code execution config functions exist in the flow tools
    - if group chat flow, there is at least one group manager agent
    - if group chat flow, there is an initial group member agent

    Returns
    -------
    WaldiezFlow
        The validated flow.

    Raises
    ------
    ValueError
        If the ordered flow is empty.
        If the model IDs are not unique.
        If the tool IDs are not unique.
        If the agents do not connect to any other node.
        If the manager's group chat has no members.
    """
    all_members = list(self.data.agents.members)
    all_chats = list(self.data.chats)
    for agent in all_members:
        agent.gather_nested_chats(
            all_agents=all_members, all_chats=all_chats
        )
        agent.gather_handoffs(all_agents=all_members, all_chats=all_chats)
    self._validate_group_chat(all_members)
    if len(all_members) == 1:
        return self._validate_single_agent_mode(all_members[0])
    ordered_flow = self.ordered_flow  # could be empty (if group chat)
    if not ordered_flow and self._ordered_flow is None:
        raise ValueError("The ordered flow is empty.")
    model_ids = self._validate_flow_models()
    tools_ids = self._validate_flow_tools()
    self.data.agents.validate_flow(model_ids, tools_ids)
    self._validate_agent_connections()
    return self

@classmethod
def default(cls) -> "WaldiezFlowData":
    """Create a default flow data.

    Returns
    -------
    WaldiezFlowData
        The default flow data.
    """
    return cls(
        nodes=[],
        edges=[],
        viewport={},
        agents=WaldiezAgents(
            userProxyAgents=[],
            assistantAgents=[
                WaldiezAssistant(
                    id="assistant1",
                    name="Assistant 1",
                    created_at=now(),
                    updated_at=now(),
                    data=WaldiezAssistantData(
                        termination=WaldiezAgentTerminationMessage()
                    ),
                ),
                WaldiezAssistant(
                    id="assistant2",
                    name="Assistant 2",
                    created_at=now(),
                    updated_at=now(),
                    data=WaldiezAssistantData(
                        # is_multimodal=True,  # we need an api key for this
                        termination=WaldiezAgentTerminationMessage(),
                    ),
                ),
            ],
            ragUserProxyAgents=[],
            reasoningAgents=[],
            captainAgents=[],
        ),
        models=[],
        tools=[],
        chats=[
            WaldiezChat(
                id="chat1",
                type="chat",
                source="assistant1",
                target="assistant2",
                data=WaldiezChatData(
                    name="Chat 1",
                    order=0,
                    position=0,
                    source_type="assistant",
                    target_type="assistant",
                    summary=WaldiezChatSummary(),
                    message=WaldiezChatMessage(
                        type="string",
                        content="Hello, how can I help you?",
                    ),
                    condition=WaldiezDefaultCondition.create(),
                    available=WaldiezTransitionAvailability(),
                    nested_chat=WaldiezChatNested(),
                ),
            ),
            WaldiezChat(
                id="chat2",
                type="chat",
                source="assistant2",
                target="assistant1",
                data=WaldiezChatData(
                    name="Chat 2",
                    order=1,
                    position=1,
                    source_type="assistant",
                    target_type="assistant",
                    summary=WaldiezChatSummary(),
                    message=WaldiezChatMessage(
                        type="string",
                        content="Hello, I need some help.",
                    ),
                    condition=WaldiezDefaultCondition.create(),
                    available=WaldiezTransitionAvailability(),
                    nested_chat=WaldiezChatNested(),
                    prerequisites=["chat1"],
                ),
            ),
        ],
        is_async=False,
        cache_seed=42,
    )

@model_validator(mode="after")
def validate_flow_chats(self) -> Self:
    """Validate the flow chats.

    Returns
    -------
    WaldiezFlowData
        The flow data.

    Raises
    ------
    ValueError
        If there is a chat with a prerequisite that does not exist.
    """
    self.chats.sort(key=lambda x: x.order)
    # in async, ag2 uses the "chat_id" field (and it must be an int):
    # ```
    #    prerequisites = []
    #    for chat_info in chat_queue:
    #        if "chat_id" not in chat_info:
    #            raise ValueError(
    #               "Each chat must have a unique id for "
    #               "async multi-chat execution."
    #            )
    #     chat_id = chat_info["chat_id"]
    #     pre_chats = chat_info.get("prerequisites", [])
    #     for pre_chat_id in pre_chats:
    #         if not isinstance(pre_chat_id, int):
    #             raise ValueError("Prerequisite chat id is not int.")
    #         prerequisites.append((chat_id, pre_chat_id))
    #    return prerequisites
    # ```
    id_to_chat_id: dict[str, int] = {}
    for index, chat in enumerate(self.chats):
        id_to_chat_id[chat.id] = index
        chat.set_chat_id(index)
    if not self.is_async:
        return self
    # also update the chat prerequisites (if async)
    #  we have ids(str), not chat_ids(int)
    for chat in self.chats:
        chat_prerequisites: list[int] = []
        for chat_id in chat.data.prerequisites:
            if chat_id not in id_to_chat_id:  # pragma: no cover
                raise ValueError(
                    f"Chat with id {chat_id} not found in the flow."
                )
            chat_prerequisites.append(id_to_chat_id[chat_id])
        chat.set_prerequisites(chat_prerequisites)
    return self