Swarm: Support AfterWorkOption in SwarmResult agent parameter (#358)

* Added support for AfterWorkOption for the agent parameter of a SwarmResult Signed-off-by: Mark Sze <[email protected]> * Added documentation Signed-off-by: Mark Sze <[email protected]> * Added SWARM_MANAGER to documentation Signed-off-by: Mark Sze <[email protected]> * LangChain version updated to 0.3.14 * LangChain test rewriten * LangChain version relaxed * Update langchain to 0.3.14 in setup.py * Revert langchain setup.py * Restore lang chain test Signed-off-by: Mark Sze <[email protected]> * Updated documentation. Signed-off-by: Mark Sze <[email protected]> * Fix imports * Pre-commit tidy Signed-off-by: Mark Sze <[email protected]> --------- Signed-off-by: Mark Sze <[email protected]> Co-authored-by: Davor Runje <[email protected]>
ag2ai · Jan 20, 2025 · 212d04d · 212d04d
1 parent 6b61f67
commit 212d04d
Show file tree

Hide file tree

Showing 3 changed files with 94 additions and 25 deletions.
diff --git a/autogen/agentchat/contrib/swarm_agent.py b/autogen/agentchat/contrib/swarm_agent.py
@@ -297,45 +297,54 @@ def _determine_next_agent(
     if "tool_calls" in groupchat.messages[-1]:
         return tool_execution
 
+    after_work_condition = None
+
     if tool_execution._next_agent is not None:
         next_agent = tool_execution._next_agent
         tool_execution._next_agent = None
 
-        # Check for string, access agent from group chat.
+        if not isinstance(next_agent, AfterWorkOption):
+            # Check for string, access agent from group chat.
 
-        if isinstance(next_agent, str):
-            if next_agent in swarm_agent_names:
-                next_agent = groupchat.agent_by_name(name=next_agent)
-            else:
-                raise ValueError(f"No agent found with the name '{next_agent}'. Ensure the agent exists in the swarm.")
+            if isinstance(next_agent, str):
+                if next_agent in swarm_agent_names:
+                    next_agent = groupchat.agent_by_name(name=next_agent)
+                else:
+                    raise ValueError(
+                        f"No agent found with the name '{next_agent}'. Ensure the agent exists in the swarm."
+                    )
 
-        return next_agent
+            return next_agent
+        else:
+            after_work_condition = next_agent
 
     # get the last swarm agent
     last_swarm_speaker = None
     for message in reversed(groupchat.messages):
-        if "name" in message and message["name"] in swarm_agent_names:
+        if "name" in message and message["name"] in swarm_agent_names and message["name"] != __TOOL_EXECUTOR_NAME__:
             agent = groupchat.agent_by_name(name=message["name"])
             if isinstance(agent, SwarmAgent):
                 last_swarm_speaker = agent
                 break
     if last_swarm_speaker is None:
         raise ValueError("No swarm agent found in the message history")
 
-    # If the user last spoke, return to the agent prior
-    if (user_agent and last_speaker == user_agent) or groupchat.messages[-1]["role"] == "tool":
-        return last_swarm_speaker
+    if after_work_condition is None:
+        # If the user last spoke, return to the agent prior
+        if (user_agent and last_speaker == user_agent) or groupchat.messages[-1]["role"] == "tool":
+            return last_swarm_speaker
 
-    # Resolve after_work condition (agent-level overrides global)
-    after_work_condition = (
-        last_swarm_speaker.after_work if last_swarm_speaker.after_work is not None else swarm_after_work
-    )
-    if isinstance(after_work_condition, AFTER_WORK):
-        after_work_condition = after_work_condition.agent
+        # Resolve after_work condition (agent-level overrides global)
+        after_work_condition = (
+            last_swarm_speaker.after_work if last_swarm_speaker.after_work is not None else swarm_after_work
+        )
+
+        if isinstance(after_work_condition, AFTER_WORK):
+            after_work_condition = after_work_condition.agent
 
-    # Evaluate callable after_work
-    if isinstance(after_work_condition, Callable):
-        after_work_condition = after_work_condition(last_speaker, groupchat.messages, groupchat)
+        # Evaluate callable after_work
+        if isinstance(after_work_condition, Callable):
+            after_work_condition = after_work_condition(last_swarm_speaker, groupchat.messages, groupchat)
 
     if isinstance(after_work_condition, str):  # Agent name in a string
         if after_work_condition in swarm_agent_names:
@@ -350,7 +359,7 @@ def _determine_next_agent(
         elif after_work_condition == AfterWorkOption.REVERT_TO_USER:
             return None if user_agent is None else user_agent
         elif after_work_condition == AfterWorkOption.STAY:
-            return last_speaker
+            return last_swarm_speaker
         elif after_work_condition == AfterWorkOption.SWARM_MANAGER:
             return "auto"
     else:
@@ -565,7 +574,7 @@ class SwarmResult(BaseModel):
     """
 
     values: str = ""
-    agent: Optional[Union["SwarmAgent", str]] = None
+    agent: Optional[Union["SwarmAgent", str, AfterWorkOption]] = None
     context_variables: dict[str, Any] = {}
 
     class Config:  # Add this inner class

diff --git a/test/agentchat/contrib/test_swarm.py b/test/agentchat/contrib/test_swarm.py
@@ -1,11 +1,12 @@
 # Copyright (c) 2023 - 2024, Owners of https://github.com/ag2ai
 #
 # SPDX-License-Identifier: Apache-2.0
-from typing import Any, Union
+from typing import Any, Optional, Union
 from unittest.mock import MagicMock, patch
 
 import pytest
 
+from autogen.agentchat.agent import Agent
 from autogen.agentchat.contrib.swarm_agent import (
     AFTER_WORK,
     ON_CONDITION,
@@ -16,6 +17,7 @@
     SwarmResult,
     _cleanup_temp_user_messages,
     _create_nested_chats,
+    _determine_next_agent,
     _prepare_swarm_agents,
     _process_initial_messages,
     _setup_context_variables,
@@ -990,5 +992,48 @@ async def mock_a_generate_oai_reply(*args, **kwargs):
     assert context_vars == test_context
 
 
+def test_swarmresult_afterworkoption():
+    """Tests processing of the return of an AfterWorkOption in a SwarmResult. This is put in the tool executors _next_agent attribute."""
+
+    def call_determine_next_agent(
+        next_agent_afterworkoption: AfterWorkOption, swarm_afterworkoption: AfterWorkOption
+    ) -> Optional[Agent]:
+        last_speaker_agent = SwarmAgent("dummy_1")
+        tool_executor = SwarmAgent(__TOOL_EXECUTOR_NAME__)
+        user = UserProxyAgent("User")
+        groupchat = GroupChat(
+            agents=[last_speaker_agent],
+            messages=[
+                {"tool_calls": "", "role": "tool", "content": "Test message"},
+                {"role": "tool", "content": "Test message 2", "name": "dummy_1"},
+            ],
+        )
+
+        tool_executor._next_agent = next_agent_afterworkoption
+
+        return _determine_next_agent(
+            last_speaker=last_speaker_agent,
+            groupchat=groupchat,
+            initial_agent=last_speaker_agent,
+            use_initial_agent=False,
+            tool_execution=tool_executor,
+            swarm_agent_names=["dummy_1"],
+            user_agent=user,
+            swarm_after_work=swarm_afterworkoption,
+        )
+
+    next_speaker = call_determine_next_agent(AfterWorkOption.TERMINATE, AfterWorkOption.STAY)
+    assert next_speaker is None, "Expected None as the next speaker for AfterWorkOption.TERMINATE"
+
+    next_speaker = call_determine_next_agent(AfterWorkOption.STAY, AfterWorkOption.TERMINATE)
+    assert next_speaker.name == "dummy_1", "Expected the last speaker as the next speaker for AfterWorkOption.TERMINATE"
+
+    next_speaker = call_determine_next_agent(AfterWorkOption.REVERT_TO_USER, AfterWorkOption.TERMINATE)
+    assert next_speaker.name == "User", "Expected the user agent as the next speaker for AfterWorkOption.REVERT_TO_USER"
+
+    next_speaker = call_determine_next_agent(AfterWorkOption.SWARM_MANAGER, AfterWorkOption.TERMINATE)
+    assert next_speaker == "auto", "Expected the auto speaker selection mode for AfterWorkOption.SWARM_MANAGER"
+
+
 if __name__ == "__main__":
     pytest.main([__file__])
diff --git a/website/docs/topics/swarm.ipynb b/website/docs/topics/swarm.ipynb
@@ -154,6 +154,7 @@
     "- `TERMINATE`: Terminate the chat    \n",
     "- `STAY`: Stay at the current agent    \n",
     "- `REVERT_TO_USER`: Revert to the user agent. Only if a user agent is passed in when initializing. (See below for more details)\n",
+    "- `SWARM_MANAGER`: Use the internal group chat's `auto` speaker selection method\n",
     "\n",
     "The callable function signature is:\n",
     "`def my_after_work_func(last_speaker: SwarmAgent, messages: List[Dict[str, Any]], groupchat: GroupChat) -> Union[AfterWorkOption, SwarmAgent, str]:`\n",
@@ -187,11 +188,25 @@
     "    ...\n",
     "    after_work=AfterWorkOption.TERMINATE # Or an agent or Callable\n",
     ")\n",
-    "\n",
-    "```\n",
     "```"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### SwarmResult\n",
+    "\n",
+    "When tools are called, a `SwarmResult` can be returned and that can be used to specify the next agent to speak through the `SwarmResult`'s `agent` parameter.\n",
+    "\n",
+    "The `agent` property can be an agent object, an agent's name (string), an `AfterWorkOption`, or `None`.\n",
+    "- If it is an agent object or agent name, that agent will be the next speaker.\n",
+    "- If `None` it will return to the previous speaker.\n",
+    "- If an `AfterWorkOption`, it will follow the rules noted in the previous section.\n",
+    "\n",
+    "By using an `AfterWorkOption` you have additional flexibility, such as terminating the swarm at this point, or transferring to the swarm's user agent."
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},