Solves issue with SwarmResult.agent AfterWorkOption -> str (#1155)

* Solves issue with SwarmResult.agent AfterWorkOption -> str

The previous constructor of SwarmResult converts AfterWorkOptions to a string.
This causes exceptions when a function sets an AfterWorkOption as the next agent.

* Adding unit test

* Update SwarmResult docstring

* Fix after work condition being ignored.

* Tests for tool call SwarmResult processing

* swarm_agent formatting

* Formatting test_swarm.py

* Update .secrets.baseline

---------

Co-authored-by: Mark Sze <[email protected]>

.secrets.baseline

@@ -823,7 +823,7 @@
         "filename": "test/agentchat/contrib/test_swarm.py",
         "hashed_secret": "957af7971d66cb7bfda1fff979038caeffd61025",
         "is_verified": false,
-        "line_number": 295,
+        "line_number": 299,
         "is_secret": false
       }
     ],

autogen/agentchat/contrib/swarm_agent.py

@@ -505,10 +505,14 @@ def _determine_next_agent(
 
     after_work_next_agent_selection_msg = None
 
-    # Resolve after_work condition (agent-level overrides global)
-    after_work_condition = (
-        last_swarm_speaker._swarm_after_work if last_swarm_speaker._swarm_after_work is not None else swarm_after_work
-    )
+    if after_work_condition is None:
+        # Resolve after_work condition if one hasn't been passed in (agent-level overrides global)
+        after_work_condition = (
+            last_swarm_speaker._swarm_after_work
+            if last_swarm_speaker._swarm_after_work is not None
+            else swarm_after_work
+        )
+
     if isinstance(after_work_condition, AfterWork):
         after_work_next_agent_selection_msg = after_work_condition.next_agent_selection_msg
         after_work_condition = after_work_condition.agent
@@ -835,12 +839,12 @@ class SwarmResult(BaseModel):
 
     Args:
         values (str): The result values as a string.
-        agent (ConversableAgent, str): The agent instance or agent name as a string, if applicable.
+        agent (ConversableAgent, AfterWorkOption, str): The agent instance, AfterWorkOption, or agent name as a string, if applicable.
         context_variables (dict): A dictionary of context variables.
     """
 
     values: str = ""
-    agent: Optional[Union[ConversableAgent, str]] = None
+    agent: Optional[Union[ConversableAgent, AfterWorkOption, str]] = None
     context_variables: dict[str, Any] = {}
 
     @field_serializer("agent", when_used="json")

test/agentchat/contrib/test_swarm.py

@@ -65,6 +65,10 @@ def test_swarm_result():
     result = SwarmResult(values="test", agent=agent)
     assert result.agent == agent
 
+    # Test AfterWorkOption
+    result = SwarmResult(agent=AfterWorkOption.TERMINATE)
+    assert isinstance(result.agent, AfterWorkOption)
+
 
 def test_swarm_result_serialization():
     agent = ConversableAgent(name="test_agent", human_input_mode="NEVER")
@@ -1287,5 +1291,77 @@ def test_compress_message_func():
     )
 
 
+def test_swarmresult_afterworkoption_tool_swarmresult():
+    """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_from_tool_execution(
+        last_speaker_agent: ConversableAgent,
+        tool_execution_swarm_result: Union[ConversableAgent, AfterWorkOption, str],
+        next_agent_afterworkoption: AfterWorkOption,
+        swarm_afterworkoption: AfterWorkOption,
+    ) -> Optional[Agent]:
+        another_agent = ConversableAgent(name="another_agent")
+        tool_executor, _ = _prepare_swarm_agents(last_speaker_agent, [last_speaker_agent, another_agent])
+        tool_executor._swarm_next_agent = tool_execution_swarm_result
+        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"},
+                {"role": "assistant", "content": "Test message 3", "name": "dummy_1"},
+            ],
+        )
+
+        last_speaker_agent._swarm_after_work = 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,
+        )
+
+    dummy_agent = ConversableAgent("dummy_1")
+    next_speaker = call_determine_next_agent_from_tool_execution(
+        dummy_agent, AfterWorkOption.TERMINATE, AfterWorkOption.STAY, AfterWorkOption.STAY
+    )
+    assert next_speaker is None, "Expected None as the next speaker for AfterWorkOption.TERMINATE"
+
+    dummy_agent = ConversableAgent("dummy_1")
+    next_speaker = call_determine_next_agent_from_tool_execution(
+        dummy_agent, AfterWorkOption.STAY, AfterWorkOption.TERMINATE, AfterWorkOption.TERMINATE
+    )
+    assert next_speaker.name == "dummy_1", "Expected the last speaker as the next speaker for AfterWorkOption.TERMINATE"
+
+    dummy_agent = ConversableAgent("dummy_1")
+    next_speaker = call_determine_next_agent_from_tool_execution(
+        dummy_agent, AfterWorkOption.REVERT_TO_USER, AfterWorkOption.TERMINATE, AfterWorkOption.TERMINATE
+    )
+    assert next_speaker.name == "User", "Expected the user agent as the next speaker for AfterWorkOption.REVERT_TO_USER"
+
+    dummy_agent = ConversableAgent("dummy_1")
+    next_speaker = call_determine_next_agent_from_tool_execution(
+        dummy_agent, AfterWorkOption.SWARM_MANAGER, AfterWorkOption.TERMINATE, AfterWorkOption.TERMINATE
+    )
+    assert next_speaker == "auto", "Expected the auto speaker selection mode for AfterWorkOption.SWARM_MANAGER"
+
+    dummy_agent = ConversableAgent("dummy_1")
+    next_speaker = call_determine_next_agent_from_tool_execution(
+        dummy_agent, "dummy_1", AfterWorkOption.TERMINATE, AfterWorkOption.TERMINATE
+    )
+    assert next_speaker == dummy_agent, "Expected the auto speaker selection mode for AfterWorkOption.SWARM_MANAGER"
+
+    dummy_agent = ConversableAgent("dummy_1")
+    next_speaker = call_determine_next_agent_from_tool_execution(
+        dummy_agent, dummy_agent, AfterWorkOption.TERMINATE, AfterWorkOption.TERMINATE
+    )
+    assert next_speaker == dummy_agent, "Expected the auto speaker selection mode for AfterWorkOption.SWARM_MANAGER"
+
+
 if __name__ == "__main__":
     pytest.main([__file__])