round_robin.py

# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
#
# SPDX-License-Identifier: Apache-2.0

from typing import TYPE_CHECKING, Any, Optional

from ..context_variables import ContextVariables
from ..targets.transition_target import AgentTarget, TransitionTarget
from .pattern import Pattern

if TYPE_CHECKING:
    from ...conversable_agent import ConversableAgent
    from ...groupchat import GroupChat, GroupChatManager
    from ..group_tool_executor import GroupToolExecutor


class RoundRobinPattern(Pattern):
    """RoundRobinPattern implements a round robin with handoffs between agents."""

    def _generate_handoffs(
        self,
        initial_agent: "ConversableAgent",
        agents: list["ConversableAgent"],
        user_agent: Optional["ConversableAgent"],
    ) -> None:
        """Generate handoffs between agents in a round-robin fashion."""
        # Create a list of the agents and the user_agent but put the initial_agent first
        agent_list = [initial_agent]

        # Add the rest of the agents, excluding the initial_agent and user_agent
        for agent in agents:
            if agent != initial_agent and (user_agent is None or agent != user_agent):
                agent_list.append(agent)

        # Add the user_agent last if it exists
        if user_agent is not None:
            agent_list.append(user_agent)

        # Create handoffs in a round-robin fashion
        for i, agent in enumerate(agent_list):
            # Last agent hands off to the first agent
            # Otherwise agent hands off to the next one
            handoff_target = agent_list[0] if i == len(agent_list) - 1 else agent_list[i + 1]

            agent.handoffs.set_after_work(target=AgentTarget(agent=handoff_target))

    def prepare_group_chat(
        self,
        max_rounds: int,
        messages: list[dict[str, Any]] | str,
        isolate_agent_views: bool = False,
    ) -> tuple[
        list["ConversableAgent"],
        list["ConversableAgent"],
        Optional["ConversableAgent"],
        ContextVariables,
        "ConversableAgent",
        TransitionTarget,
        "GroupToolExecutor",
        "GroupChat",
        "GroupChatManager",
        list[dict[str, Any]],
        Any,
        list[str],
        list[Any],
    ]:
        """Prepare the group chat for organic agent selection.

        Ensures that:
        1. The group manager has a valid LLM config
        2. All agents have appropriate descriptions for the group manager to use

        Args:
            max_rounds: Maximum number of conversation rounds.
            messages: Initial message(s) to start the conversation.
            isolate_agent_views: If True, agents will only maintain their own message history
                and will not receive messages from other agents. When False (default), all agents
                receive all messages. When True, messages are still stored in groupchat.messages
                for the GroupChatManager's view, but are not broadcast to other agents.

        Returns:
            Tuple containing all necessary components for the group chat.
        """
        # Use the parent class's implementation to prepare the agents and group chat
        (
            agents,
            wrapped_agents,
            user_agent,
            context_variables,
            initial_agent,
            group_after_work,
            tool_executor,
            groupchat,
            manager,
            processed_messages,
            last_agent,
            group_agent_names,
            temp_user_list,
        ) = super().prepare_group_chat(
            max_rounds=max_rounds,
            messages=messages,
            isolate_agent_views=isolate_agent_views,
        )

        # Create the handoffs between agents
        self._generate_handoffs(initial_agent=initial_agent, agents=agents, user_agent=user_agent)

        # Return all components with our group_after_work
        return (
            agents,
            wrapped_agents,
            user_agent,
            context_variables,
            initial_agent,
            group_after_work,
            tool_executor,
            groupchat,
            manager,
            processed_messages,
            last_agent,
            group_agent_names,
            temp_user_list,
        )