Swarm: Allow functions to update agent's state, including system message, before replying

autogen/agentchat/contrib/swarm_agent.py

@@ -110,7 +110,17 @@ def custom_afterwork_func(last_speaker: SwarmAgent, messages: List[Dict[str, Any
     INIT_AGENT_USED = False
 
     def swarm_transition(last_speaker: SwarmAgent, groupchat: GroupChat):
-        """Swarm transition function to determine the next agent in the conversation"""
+        """Swarm transition function to determine and prepare the next agent in the conversation"""
+        next_agent = determine_next_agent(last_speaker, groupchat)
+
+        if next_agent and isinstance(next_agent, SwarmAgent):
+            # Update their state
+            next_agent.update_state(context_variables, groupchat.messages)
+
+        return next_agent
+
+    def determine_next_agent(last_speaker: SwarmAgent, groupchat: GroupChat):
+        """Determine the next agent in the conversation"""
         nonlocal INIT_AGENT_USED
         if not INIT_AGENT_USED:
             INIT_AGENT_USED = True
@@ -257,6 +267,7 @@ def __init__(
         human_input_mode: Literal["ALWAYS", "NEVER", "TERMINATE"] = "NEVER",
         description: Optional[str] = None,
         code_execution_config=False,
+        system_message_func: Optional[Callable] = None,
         **kwargs,
     ) -> None:
         super().__init__(
@@ -286,6 +297,8 @@ def __init__(
         self._context_variables = {}
         self._next_agent = None
 
+        self._system_message_func = system_message_func
+
     def _set_to_tool_execution(self, context_variables: Optional[Dict[str, Any]] = None):
         """Set to a special instance of SwarmAgent that is responsible for executing tool calls from other swarm agents.
         This agent will be used internally and should not be visible to the user.
@@ -458,6 +471,11 @@ def add_functions(self, func_list: List[Callable]):
         for func in func_list:
             self.add_single_function(func)
 
+    def update_state(self, context_variables: Optional[Dict[str, Any]], messages: List[Dict[str, Any]]):
+        """Updates the state of the agent, system message so far. This is called when they're selected and just before they speak."""
+        if self._system_message_func:
+            self.update_system_message(self._system_message_func(context_variables, messages))
+
 
 # Forward references for SwarmAgent in SwarmResult
 SwarmResult.update_forward_refs()

test/agentchat/contrib/test_swarm.py

@@ -460,6 +460,50 @@ def test_initialization():
             initial_agent=agent1, messages=TEST_MESSAGES, agents=[agent1, agent2], max_rounds=3
         )
 
+    def test_sys_message_func():
+        """Tests a custom system message function"""
+
+        # This test will use context variables and the messages to construct a custom system message
+        # This will be available at the point of reply (we use register a reply to capture it at that point)
+
+        # To store the system message
+        class MessageContainer:
+            def __init__(self):
+                self.final_sys_message = ""
+
+        message_container = MessageContainer()
+
+        def my_sys_message(context_variables, messages) -> str:
+            return f"This is a custom system message with {context_variables['sample_name']} and a total of {len(messages)} message(s)."
+
+        agent1 = SwarmAgent("agent1", system_message_func=my_sys_message)
+        agent2 = SwarmAgent("agent2")
+
+        test_context_variables = {"sample_name": "Bob"}
+
+        # Mock a reply to be able to capture the system message
+        def mock_generate_oai_reply(*args, **kwargs):
+            message_container.final_sys_message = args[0]._oai_system_message[0][
+                "content"
+            ]  # The SwarmAgent's system message
+            return True, "This is a mock response from the agent."
+
+        agent1.register_reply([ConversableAgent, None], mock_generate_oai_reply)
+
+        chat_result, context_vars, last_speaker = initiate_swarm_chat(
+            initial_agent=agent1,
+            messages=TEST_MESSAGES,
+            agents=[agent1, agent2],
+            context_variables=test_context_variables,
+            max_rounds=4,
+        )
+
+        # The system message should be the custom message
+        assert (
+            message_container.final_sys_message
+            == "This is a custom system message with Bob and a total of 1 message(s)."
+        )
+
 
 if __name__ == "__main__":
     pytest.main([__file__])