From 6170a6d4c84116b3a2588874e2dab08b3c6bd766 Mon Sep 17 00:00:00 2001 From: Ignacio Date: Wed, 26 Feb 2025 14:50:02 -0800 Subject: [PATCH] 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 <66362098+marklysze@users.noreply.github.com> --- .secrets.baseline | 2 +- autogen/agentchat/contrib/swarm_agent.py | 16 +++-- test/agentchat/contrib/test_swarm.py | 76 ++++++++++++++++++++++++ 3 files changed, 87 insertions(+), 7 deletions(-) diff --git a/.secrets.baseline b/.secrets.baseline index bacf7cb3aa..bcbfb16124 100644 --- a/.secrets.baseline +++ b/.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 } ], diff --git a/autogen/agentchat/contrib/swarm_agent.py b/autogen/agentchat/contrib/swarm_agent.py index 26f1f01c46..c957265345 100644 --- a/autogen/agentchat/contrib/swarm_agent.py +++ b/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") diff --git a/test/agentchat/contrib/test_swarm.py b/test/agentchat/contrib/test_swarm.py index 69b87ae9e0..b019f815e3 100644 --- a/test/agentchat/contrib/test_swarm.py +++ b/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__])