SwarmAgent

@model_validator(mode="after")
def validate_handoffs(self) -> Self:
    """Validate the hand offs.

    Returns
    -------
    Self
        The swarm agent data.

    Raises
    ------
    ValueError
        If there are more than one `AfterWork`s.
    """
    after_works: List[WaldiezSwarmAfterWork] = [
        hand_off
        for hand_off in self.handoffs
        if isinstance(hand_off, WaldiezSwarmAfterWork)
    ]
    if len(after_works) > 1:
        raise ValueError(
            "Each agent should have at most one `AfterWork` "
            "and (if any) it should be at the end of the list."
        )
    on_conditions: List[WaldiezSwarmOnCondition] = [
        hand_off
        for hand_off in self.handoffs
        if isinstance(hand_off, WaldiezSwarmOnCondition)
    ]
    on_conditions = sorted(on_conditions, key=lambda x: x.target.order)
    handoffs = on_conditions + after_works
    self.handoffs = handoffs
    return self

def get_available(
    self,
    name_prefix: Optional[str] = None,
    name_suffix: Optional[str] = None,
) -> Tuple[str, str]:
    """Get the available string.

    Parameters
    ----------
    name_prefix : str, optional
        The prefix to add to the function name. Default is None.
    name_suffix : str, optional
        The suffix to add to the function name. Default is None.
    Returns
    -------
    Tuple[str, str]
        The available string or function name and code if available.
    """
    return self.available.get_available(
        name_prefix,
        name_suffix,
    )

def get_recipient(
    self,
    agent_names: Dict[str, str],
    name_prefix: Optional[str] = None,
    name_suffix: Optional[str] = None,
) -> Tuple[str, str]:
    """Get the recipient string.

    Parameters
    ----------
    agent_names : Dict[str, str]
        A mapping of agent id to agent name.
    name_prefix : Optional[str], optional
        The prefix for the function name, by default None.
    name_suffix : Optional[str], optional
        The suffix for the function name, by default None.

    Returns
    -------
    Tuple[str, str]
        The recipient string and the function content if applicable.
    """
    if self.recipient_type == "option":
        return f"AfterWork(AfterWorkOption.{self.recipient})", ""
    if self.recipient_type == "agent":
        # the the recipient is passed as the agent name
        # (and not its id), care should be taken to ensure
        # the all the agents in the flow have unique names
        agent_instance = agent_names.get(self.recipient, self.recipient)
        return f"AfterWork({agent_instance})", ""

    function_name = CUSTOM_AFTER_WORK
    if name_prefix:
        function_name = f"{name_prefix}_{function_name}"
    if name_suffix:
        function_name = f"{function_name}_{name_suffix}"
    return (
        f"AfterWork({function_name})",
        generate_function(
            function_name=function_name,
            function_args=CUSTOM_AFTER_WORK_ARGS,
            function_body=self._recipient_string,
            function_types=CUSTOM_AFTER_WORK_TYPES,
        ),
    )

@model_validator(mode="after")
def validate_recipient(self) -> Self:
    """Validate the recipient.

    Returns
    -------
    WaldiezSwarmAfterWork
        The validated after work model.

    Raises
    ------
    ValueError
        If the validation fails.
    """
    self._recipient_string = self.recipient
    if self.recipient_type == "callable":
        is_valid, error_or_body = check_function(
            code_string=self.recipient,
            function_name=CUSTOM_AFTER_WORK,
            function_args=CUSTOM_AFTER_WORK_ARGS,
        )
        if not is_valid or not error_or_body:
            # pylint: disable=inconsistent-quotes
            raise ValueError(
                f"Invalid custom method: {error_or_body or 'no content'}"
            )
        self._recipient_string = error_or_body
    elif self.recipient_type == "option":
        if self.recipient not in [
            "TERMINATE",
            "REVERT_TO_USER",
            "STAY",
            "SWARM_MANAGER",
        ]:
            raise ValueError("Invalid option.")
    return self

def get_update_function(
    self,
    name_prefix: Optional[str] = None,
    name_suffix: Optional[str] = None,
) -> Tuple[str, str]:
    """Get the update function.

    Parameters
    ----------
    name_prefix : str, optional
        The prefix of the name, by default None
    name_suffix : str, optional
        The suffix of the name, by default None

    Returns
    -------
    Tuple[str, str]
        The update function and the function name.

    """
    function_name = CUSTOM_UPDATE_SYSTEM_MESSAGE
    if name_prefix:
        function_name = f"{name_prefix}_{function_name}"
    if name_suffix:
        function_name = f"{function_name}_{name_suffix}"
    return (
        generate_function(
            function_name=function_name,
            function_args=CUSTOM_UPDATE_SYSTEM_MESSAGE_ARGS,
            function_types=CUSTOM_UPDATE_SYSTEM_MESSAGE_TYPES,
            function_body=self._update_function,
        ),
        function_name,
    )

@model_validator(mode="after")
def validate_update_system_message(self) -> Self:
    """Validate the update system message function.

    Returns
    -------
    UpdateSystemMessage
        The validated update system message.

    Raises
    ------
    ValueError
        If the type is callable and the function is invalid.
        or if the function type is not 'string' or 'callable'.

    """
    self._update_function = self.update_function
    if self.update_function_type == "callable":
        valid, error_or_body = check_function(
            code_string=self.update_function,
            function_name=CUSTOM_UPDATE_SYSTEM_MESSAGE,
            function_args=CUSTOM_UPDATE_SYSTEM_MESSAGE_ARGS,
        )
        if not valid or not error_or_body:
            # pylint: disable=inconsistent-quotes
            raise ValueError(
                f"Invalid custom method: {error_or_body or 'no content'}"
            )
        self._update_function = error_or_body
    return self