Swarm: User-defined AfterWork string for agent selection using LLM

autogen/agentchat/__init__.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2023 - 2024, Owners of https://github.com/ag2ai
+# Copyright (c) 2023 - 2025, Owners of https://github.com/ag2ai
 #
 # SPDX-License-Identifier: Apache-2.0
 #
@@ -17,14 +17,16 @@
 from .contrib.swarm_agent import (
     AFTER_WORK,
     ON_CONDITION,
-    UPDATE_SYSTEM_MESSAGE,
+    AfterWork,
     AfterWorkOption,
-    SwarmAgent,
+    OnCondition,
     SwarmResult,
+    UpdateCondition,
     a_initiate_swarm_chat,
     initiate_swarm_chat,
+    register_hand_off,
 )
-from .conversable_agent import ConversableAgent, register_function
+from .conversable_agent import UPDATE_SYSTEM_MESSAGE, ConversableAgent, UpdateSystemMessage, register_function
 from .groupchat import GroupChat, GroupChatManager
 from .user_proxy_agent import UserProxyAgent
 from .utils import gather_usage_summary
@@ -42,11 +44,15 @@
     "ChatResult",
     "initiate_swarm_chat",
     "a_initiate_swarm_chat",
-    "SwarmAgent",
     "SwarmResult",
     "ON_CONDITION",
+    "OnCondition",
+    "UpdateCondition",
     "AFTER_WORK",
+    "AfterWork",
     "AfterWorkOption",
+    "register_hand_off",
+    "UpdateSystemMessage",
     "UPDATE_SYSTEM_MESSAGE",
     "ReasoningAgent",
     "visualize_tree",

autogen/agentchat/groupchat.py

@@ -33,6 +33,15 @@
 
 logger = logging.getLogger(__name__)
 
+__SELECT_SPEAKER_MESSAGE_TEMPLATE__ = """You are in a role play game. The following roles are available:
+{roles}.
+Read the following conversation.
+Then select the next role from {agentlist} to play. Only return the role."""
+
+__SELECT_SPEAKER_PROMPT_TEMPLATE__ = (
+    "Read the above conversation. Then select the next role from {agentlist} to play. Only return the role."
+)
+
 
 @dataclass
 class GroupChat:
@@ -130,13 +139,8 @@ def custom_speaker_selection_func(
     speaker_transitions_type: Literal["allowed", "disallowed", None] = None
     enable_clear_history: bool = False
     send_introductions: bool = False
-    select_speaker_message_template: str = """You are in a role play game. The following roles are available:
-                {roles}.
-                Read the following conversation.
-                Then select the next role from {agentlist} to play. Only return the role."""
-    select_speaker_prompt_template: str = (
-        "Read the above conversation. Then select the next role from {agentlist} to play. Only return the role."
-    )
+    select_speaker_message_template: str = __SELECT_SPEAKER_MESSAGE_TEMPLATE__
+    select_speaker_prompt_template: str = __SELECT_SPEAKER_PROMPT_TEMPLATE__
     select_speaker_auto_multiple_template: str = """You provided more than one name in your text, please return just the name of the next speaker. To determine the speaker use these prioritised rules:
     1. If the context refers to themselves as a speaker e.g. "As the..." , choose that speaker's name
     2. If it refers to the "next" speaker name, choose that name
@@ -375,7 +379,7 @@ def select_speaker_prompt(self, agents: Optional[list[Agent]] = None) -> str:
 
         agentlist = f"{[agent.name for agent in agents]}"
 
-        return_prompt = self.select_speaker_prompt_template.format(agentlist=agentlist)
+        return_prompt = f"{self.select_speaker_prompt_template}".replace("{agentlist}", agentlist)
         return return_prompt
 
     def introductions_msg(self, agents: Optional[list[Agent]] = None) -> str:

autogen/agentchat/realtime_agent/realtime_agent.py

@@ -8,7 +8,6 @@
 import anyio
 from asyncer import create_task_group, syncify
 
-from ... import SwarmAgent
 from ...tools import Tool, get_function_schema
 from ..agent import Agent
 from ..conversable_agent import ConversableAgent
@@ -93,8 +92,8 @@ def __init__(
         self._answer_event: anyio.Event = anyio.Event()
         self._answer: str = ""
         self._start_swarm_chat = False
-        self._initial_agent: Optional[SwarmAgent] = None
-        self._agents: Optional[list[SwarmAgent]] = None
+        self._initial_agent: Optional[ConversableAgent] = None
+        self._agents: Optional[list[ConversableAgent]] = None
 
     @property
     def logger(self) -> Logger:
@@ -122,15 +121,15 @@ def register_observer(self, observer: RealtimeObserver) -> None:
     def register_swarm(
         self,
         *,
-        initial_agent: SwarmAgent,
-        agents: list[SwarmAgent],
+        initial_agent: ConversableAgent,
+        agents: list[ConversableAgent],
         system_message: Optional[str] = None,
     ) -> None:
         """Register a swarm of agents with the Realtime Agent.
 
         Args:
-            initial_agent (SwarmAgent): The initial agent.
-            agents (list[SwarmAgent]): The agents in the swarm.
+            initial_agent (ConversableAgent): The initial agent.
+            agents (list[ConversableAgent]): The agents in the swarm.
             system_message (str): The system message for the agent.
         """
         logger = self.logger

notebook/agentchat_realtime_swarm.ipynb

@@ -14,6 +14,17 @@
     "In this notebook, we implement OpenAI's [airline customer service example](https://github.com/openai/swarm/tree/main/examples/airline) in AG2 using the RealtimeAgent for enhanced interaction."
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "````mdx-code-block\n",
+    ":::note\n",
+    "This notebook has been updated as swarms can now accommodate any ConversableAgent.\n",
+    ":::\n",
+    "````"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -270,22 +281,22 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
-    "from autogen import ON_CONDITION, SwarmAgent\n",
+    "from autogen import ConversableAgent, OnCondition, register_hand_off\n",
     "\n",
     "# Triage Agent\n",
-    "triage_agent = SwarmAgent(\n",
+    "triage_agent = ConversableAgent(\n",
     "    name=\"Triage_Agent\",\n",
     "    system_message=triage_instructions(context_variables=context_variables),\n",
     "    llm_config=llm_config,\n",
     "    functions=[non_flight_enquiry],\n",
     ")\n",
     "\n",
     "# Flight Modification Agent\n",
-    "flight_modification = SwarmAgent(\n",
+    "flight_modification = ConversableAgent(\n",
     "    name=\"Flight_Modification_Agent\",\n",
     "    system_message=\"\"\"You are a Flight Modification Agent for a customer service airline.\n",
     "      Your task is to determine if the user wants to cancel or change their flight.\n",
@@ -295,23 +306,23 @@
     ")\n",
     "\n",
     "# Flight Cancel Agent\n",
-    "flight_cancel = SwarmAgent(\n",
+    "flight_cancel = ConversableAgent(\n",
     "    name=\"Flight_Cancel_Traversal\",\n",
     "    system_message=STARTER_PROMPT + FLIGHT_CANCELLATION_POLICY,\n",
     "    llm_config=llm_config,\n",
     "    functions=[initiate_refund, initiate_flight_credits, case_resolved, escalate_to_agent],\n",
     ")\n",
     "\n",
     "# Flight Change Agent\n",
-    "flight_change = SwarmAgent(\n",
+    "flight_change = ConversableAgent(\n",
     "    name=\"Flight_Change_Traversal\",\n",
     "    system_message=STARTER_PROMPT + FLIGHT_CHANGE_POLICY,\n",
     "    llm_config=llm_config,\n",
     "    functions=[valid_to_change_flight, change_flight, case_resolved, escalate_to_agent],\n",
     ")\n",
     "\n",
     "# Lost Baggage Agent\n",
-    "lost_baggage = SwarmAgent(\n",
+    "lost_baggage = ConversableAgent(\n",
     "    name=\"Lost_Baggage_Traversal\",\n",
     "    system_message=STARTER_PROMPT + LOST_BAGGAGE_POLICY,\n",
     "    llm_config=llm_config,\n",
@@ -325,33 +336,35 @@
    "source": [
     "### Register Handoffs\n",
     "\n",
-    "Now we register the handoffs for the agents. Note that you don't need to define the transfer functions and pass them in. Instead, you can directly register the handoffs using the `ON_CONDITION` class."
+    "Now we register the handoffs for the agents. Note that you don't need to define the transfer functions and pass them in. Instead, you can directly register the handoffs using the `OnCondition` class."
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "# Register hand-offs\n",
-    "triage_agent.register_hand_off(\n",
-    "    [\n",
-    "        ON_CONDITION(flight_modification, \"To modify a flight\"),\n",
-    "        ON_CONDITION(lost_baggage, \"To find lost baggage\"),\n",
-    "    ]\n",
+    "register_hand_off(\n",
+    "    agent=triage_agent,\n",
+    "    hand_to=[\n",
+    "        OnCondition(flight_modification, \"To modify a flight\"),\n",
+    "        OnCondition(lost_baggage, \"To find lost baggage\"),\n",
+    "    ],\n",
     ")\n",
     "\n",
-    "flight_modification.register_hand_off(\n",
-    "    [\n",
-    "        ON_CONDITION(flight_cancel, \"To cancel a flight\"),\n",
-    "        ON_CONDITION(flight_change, \"To change a flight\"),\n",
-    "    ]\n",
+    "register_hand_off(\n",
+    "    agent=flight_modification,\n",
+    "    hand_to=[\n",
+    "        OnCondition(flight_cancel, \"To cancel a flight\"),\n",
+    "        OnCondition(flight_change, \"To change a flight\"),\n",
+    "    ],\n",
     ")\n",
     "\n",
     "transfer_to_triage_description = \"Call this function when a user needs to be transferred to a different agent and a different policy.\\nFor instance, if a user is asking about a topic that is not handled by the current agent, call this function.\"\n",
     "for agent in [flight_modification, flight_cancel, flight_change, lost_baggage]:\n",
-    "    agent.register_hand_off(ON_CONDITION(triage_agent, transfer_to_triage_description))"
+    "    register_hand_off(agent=agent, hand_to=OnCondition(triage_agent, transfer_to_triage_description))"
    ]
   },
   {

notebook/agentchat_swarm.ipynb

@@ -16,6 +16,17 @@
     "In this notebook, we implement OpenAI's [airline customer service example](https://github.com/openai/swarm/tree/main/examples/airline) in AG2."
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "````mdx-code-block\n",
+    ":::note\n",
+    "This notebook has been updated as swarms can now accommodate any ConversableAgent.\n",
+    ":::\n",
+    "````"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -222,18 +233,18 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from autogen import ON_CONDITION, AfterWorkOption, SwarmAgent, initiate_swarm_chat\n",
+    "from autogen import AfterWorkOption, ConversableAgent, OnCondition, initiate_swarm_chat, register_hand_off\n",
     "\n",
     "# Triage Agent\n",
-    "triage_agent = SwarmAgent(\n",
+    "triage_agent = ConversableAgent(\n",
     "    name=\"Triage_Agent\",\n",
     "    system_message=triage_instructions(context_variables=context_variables),\n",
     "    llm_config=llm_config,\n",
     "    functions=[non_flight_enquiry],\n",
     ")\n",
     "\n",
     "# Flight Modification Agent\n",
-    "flight_modification = SwarmAgent(\n",
+    "flight_modification = ConversableAgent(\n",
     "    name=\"Flight_Modification_Agent\",\n",
     "    system_message=\"\"\"You are a Flight Modification Agent for a customer service airline.\n",
     "      Your task is to determine if the user wants to cancel or change their flight.\n",
@@ -243,23 +254,23 @@
     ")\n",
     "\n",
     "# Flight Cancel Agent\n",
-    "flight_cancel = SwarmAgent(\n",
+    "flight_cancel = ConversableAgent(\n",
     "    name=\"Flight_Cancel_Traversal\",\n",
     "    system_message=STARTER_PROMPT + FLIGHT_CANCELLATION_POLICY,\n",
     "    llm_config=llm_config,\n",
     "    functions=[initiate_refund, initiate_flight_credits, case_resolved, escalate_to_agent],\n",
     ")\n",
     "\n",
     "# Flight Change Agent\n",
-    "flight_change = SwarmAgent(\n",
+    "flight_change = ConversableAgent(\n",
     "    name=\"Flight_Change_Traversal\",\n",
     "    system_message=STARTER_PROMPT + FLIGHT_CHANGE_POLICY,\n",
     "    llm_config=llm_config,\n",
     "    functions=[valid_to_change_flight, change_flight, case_resolved, escalate_to_agent],\n",
     ")\n",
     "\n",
     "# Lost Baggage Agent\n",
-    "lost_baggage = SwarmAgent(\n",
+    "lost_baggage = ConversableAgent(\n",
     "    name=\"Lost_Baggage_Traversal\",\n",
     "    system_message=STARTER_PROMPT + LOST_BAGGAGE_POLICY,\n",
     "    llm_config=llm_config,\n",
@@ -273,7 +284,7 @@
    "source": [
     "### Register Handoffs\n",
     "\n",
-    "Now we register the handoffs for the agents. Note that you don't need to define the transfer functions and pass them in. Instead, you can directly register the handoffs using the `ON_CONDITION` class."
+    "Now we register the handoffs for the agents. Note that you don't need to define the transfer functions and pass them in. Instead, you can directly register the handoffs using the `OnCondition` class."
    ]
   },
   {
@@ -283,23 +294,25 @@
    "outputs": [],
    "source": [
     "# Register hand-offs\n",
-    "triage_agent.register_hand_off(\n",
-    "    [\n",
-    "        ON_CONDITION(flight_modification, \"To modify a flight\"),\n",
-    "        ON_CONDITION(lost_baggage, \"To find lost baggage\"),\n",
-    "    ]\n",
+    "register_hand_off(\n",
+    "    agent=triage_agent,\n",
+    "    hand_to=[\n",
+    "        OnCondition(flight_modification, \"To modify a flight\"),\n",
+    "        OnCondition(lost_baggage, \"To find lost baggage\"),\n",
+    "    ],\n",
     ")\n",
     "\n",
-    "flight_modification.register_hand_off(\n",
-    "    [\n",
-    "        ON_CONDITION(flight_cancel, \"To cancel a flight\"),\n",
-    "        ON_CONDITION(flight_change, \"To change a flight\"),\n",
-    "    ]\n",
+    "register_hand_off(\n",
+    "    agent=flight_modification,\n",
+    "    hand_to=[\n",
+    "        OnCondition(flight_cancel, \"To cancel a flight\"),\n",
+    "        OnCondition(flight_change, \"To change a flight\"),\n",
+    "    ],\n",
     ")\n",
     "\n",
     "transfer_to_triage_description = \"Call this function when a user needs to be transferred to a different agent and a different policy.\\nFor instance, if a user is asking about a topic that is not handled by the current agent, call this function.\"\n",
     "for agent in [flight_modification, flight_cancel, flight_change, lost_baggage]:\n",
-    "    agent.register_hand_off(ON_CONDITION(triage_agent, transfer_to_triage_description))"
+    "    register_hand_off(agent=agent, hand_to=OnCondition(triage_agent, transfer_to_triage_description))"
    ]
   },
   {
@@ -373,7 +386,7 @@
       "\u001b[33mTool_Execution\u001b[0m (to chat_manager):\n",
       "\n",
       "\u001b[32m***** Response from calling tool (call_Qgji9KAw1e3ktxykLU8v1wg7) *****\u001b[0m\n",
-      "SwarmAgent --> Flight_Modification_Agent\n",
+      "Swarm agent --> Flight_Modification_Agent\n",
       "\u001b[32m**********************************************************************\u001b[0m\n",
       "\n",
       "--------------------------------------------------------------------------------\n",
@@ -397,7 +410,7 @@
       "\u001b[33mTool_Execution\u001b[0m (to chat_manager):\n",
       "\n",
       "\u001b[32m***** Response from calling tool (call_QYu7uBko1EaEZ7VzxPwx2jNO) *****\u001b[0m\n",
-      "SwarmAgent --> Flight_Cancel_Traversal\n",
+      "Swarm agent --> Flight_Cancel_Traversal\n",
       "\u001b[32m**********************************************************************\u001b[0m\n",
       "\n",
       "--------------------------------------------------------------------------------\n",