Agent

def gather_handoff_ids(
    self,
    group_chats: list["WaldiezChat"],
    nested_chat_id: str,
) -> None:
    """Gather all the handoff IDs for this agent.

    This method will gather all the handoff IDs from the agent's data,
    including those that might not be passed in data.handoffs.

    Parameters
    ----------
    group_chats : list["WaldiezChat"]
        The list of group chats that this agent is part of.
    nested_chat_id : str
        The ID of the nested chat to include in handoffs if it exists.

    """
    existing_handoffs = set(self.data.handoffs)
    has_nested_chat = len(self.data.nested_chats) > 0 and any(
        chat.messages for chat in self.data.nested_chats
    )
    # Step 1: Add missing group chat handoffs
    # These are chats between group members (equivalent to groupEdges)
    for chat in group_chats:
        if chat.id not in existing_handoffs:
            self.data.handoffs.append(chat.id)
            existing_handoffs.add(chat.id)

    # Step 2: Add nested chat if it exists and is not already in handoffs
    if (
        has_nested_chat and nested_chat_id not in existing_handoffs
    ):  # pragma: no branch
        self.data.handoffs.append(nested_chat_id)
        existing_handoffs.add(nested_chat_id)

    # Step 3: Validate all handoffs still exist
    # Remove any handoffs that reference non-existent chats
    valid_chat_ids = {chat.id for chat in group_chats}
    if has_nested_chat:  # pragma: no branch
        valid_chat_ids.add(nested_chat_id)

    # Filter out invalid handoffs
    self.data.handoffs = [
        handoff
        for handoff in self.data.handoffs
        if handoff in valid_chat_ids
    ]

def gather_handoffs(
    self,
    all_agents: list["WaldiezAgent"],
    all_chats: list["WaldiezChat"],
) -> None:
    """Gather all the handoffs including.

    Including ones that might not be passed in data.handoffs.

    Parameters
    ----------
    all_agents : list["WaldiezAgent"]
        The list of all agents in the flow.
    all_chats : list["WaldiezChat"]
        The list of all chats in the flow.

    """
    self.gather_nested_chats(all_agents, all_chats)
    if not self.is_group_member or self._checked_handoffs:
        return
    nested_chat_id = "nested-chat"
    self._checked_handoffs = True
    group_chats, group_nested_chats = self._get_agent_chats(
        all_agents, all_chats
    )
    if group_nested_chats:  # pragma: no branch
        self._ensure_one_nested_chat(group_nested_chats)
    self.gather_handoff_ids(
        group_chats=group_chats, nested_chat_id=nested_chat_id
    )
    # generate the actual handoff instances
    for handoff_id in self.data.handoffs:
        if handoff_id == nested_chat_id:
            nested_chat_handoff = self._generate_handoff_from_nested(
                group_nested_chats
            )
            if nested_chat_handoff:  # pragma: no branch
                self._handoffs.append(nested_chat_handoff)
        else:
            chat = next(
                (chat for chat in group_chats if chat.id == handoff_id),
                None,
            )
            if chat:  # pragma: no branch
                self._handoffs.append(chat.as_handoff())

def gather_nested_chats(
    self,
    all_agents: list["WaldiezAgent"],
    all_chats: list["WaldiezChat"],
) -> None:
    """Gather the nested chats for the agent.

    Parameters
    ----------
    all_agents : list["WaldiezAgent"]
        All the agents in the flow.
    all_chats : list["WaldiezChat"]
        All the chats in the flow.
    """
    if self._checked_nested_chats:
        return
    self._checked_nested_chats = True
    all_chat_ids = {chat.id for chat in all_chats}
    all_agent_ids = {agent.id for agent in all_agents}
    # only use chats that do have messages and "triggered_by"
    nested_chats: list[WaldiezAgentNestedChat] = []
    for chat in self.data.nested_chats:
        if not chat.messages or not chat.triggered_by:  # pragma: no cover
            continue
        # make sure the ids exist
        chat.messages = [
            message
            for message in chat.messages
            if message.id in all_chat_ids
        ]
        chat.triggered_by = [
            agent_id
            for agent_id in chat.triggered_by
            if agent_id in all_agent_ids
        ]
        if chat.messages and chat.triggered_by:  # pragma: no branch
            nested_chats.append(chat)
    self.data.nested_chats = nested_chats

@field_validator("agent_type")
@classmethod
def validate_agent_type(cls, v: WaldiezAgentType) -> WaldiezAgentType:
    """Validate the agent type.

    Parameters
    ----------
    v : WaldiezAgentType
        The agent type.

    Returns
    -------
    WaldiezAgentType
        The validated agent type.

    Raises
    ------
    ValueError
        If the agent type is not valid.
    """

    def _get_warning_message(old_type: str, new_type: str) -> str:
        return (
            f"The agent type '{old_type}' is deprecated. "
            f"Use '{new_type}' instead."
        )

    if v == "user":
        warnings.warn(
            _get_warning_message("user", "user_proxy"),
            DeprecationWarning,
            stacklevel=2,
        )
        return "user_proxy"
    if v == "rag_user":
        warnings.warn(
            _get_warning_message("rag_user", "rag_user_proxy"),
            DeprecationWarning,
            stacklevel=2,
        )
        return "rag_user_proxy"
    return v

def validate_code_execution(self, tool_ids: list[str]) -> None:
    """Validate the code execution config.

    Parameters
    ----------
    tool_ids : list[str]
        The list of tool IDs.

    Raises
    ------
    ValueError
        If a function is not found
    """
    # if the config dict has functions, make sure they can be found
    if isinstance(
        self.data.code_execution_config, WaldiezAgentCodeExecutionConfig
    ):
        for function in self.data.code_execution_config.functions:
            if function not in tool_ids:
                raise ValueError(
                    f"Function '{function}' not found in tools"
                )

def validate_linked_models(self, model_ids: list[str]) -> None:
    """Validate the models.

    Parameters
    ----------
    model_ids : List[str]
        The list of model IDs.

    Raises
    ------
    ValueError
        If a model is not found
    """
    # if the config dict has models, make sure they can be found
    for model in self.data.model_ids:
        if model not in model_ids:
            raise ValueError(
                f"Model '{model}' not found in agent's {self.id} models"
            )

def validate_linked_tools(
    self, tool_ids: list[str], agent_ids: list[str]
) -> None:
    """Validate the tools.

    Parameters
    ----------
    tool_ids : list[str]
        The list of tool IDs.
    agent_ids : list[str]
        The list of agent IDs.

    Raises
    ------
    ValueError
        If a tool or agent is not found
    """
    # if the config dict has tools, make sure they can be found
    for tool in self.data.tools:
        if tool.id not in tool_ids:
            raise ValueError(
                f"Tool '{tool.id}' not found in agent's {self.id} tools"
            )
        if tool.executor_id not in agent_ids:
            raise ValueError(
                f"Agent '{tool.executor_id}' not found in agents"
            )

@model_validator(mode="after")
def update_context_variables(self) -> Self:
    """Update the context variables.

    Returns
    -------
    Self
        The updated instance of the class.
    """
    context_vars = update_dict(self.context_variables)
    self.context_variables = context_vars
    return self

def get_termination_function(
    self,
    name_prefix: Optional[str] = None,
    name_suffix: Optional[str] = None,
) -> tuple[str, str]:
    """Get the termination function.

    Parameters
    ----------
    name_prefix : str
        The function name prefix.
    name_suffix : str
        The function name suffix.

    Returns
    -------
    tuple[str, str]
        The termination function and the function name.
    """
    function_name = "is_termination_message"
    if name_prefix:
        function_name = f"{name_prefix}_{function_name}"
    if name_suffix:
        function_name = f"{function_name}_{name_suffix}"
    if self.type in ("none", "keyword"):
        return self.string, function_name
    return (
        generate_function(
            function_name=function_name,
            function_args=IS_TERMINATION_MESSAGE_ARGS,
            function_types=IS_TERMINATION_MESSAGE_TYPES,
            function_body=self.string,
        ),
        function_name,
    )

@model_validator(mode="after")
def validate_termination_message(self) -> Self:
    """Validate the termination message configuration.

    Raises
    ------
    ValueError
        If the configuration is invalid.

    Returns
    -------
    WaldiezAgentTerminationMessage
        The validated termination message configuration.
    """
    if self.type == "method":
        self._validate_method_content()
    if self.type == "keyword":
        self._validate_keyword()
    if self.type == "none":
        self._string = "None"
    return self