Ability to add MessageTransforms to the GroupChat's Select Speaker ne…

…sted chat (speaker_selection_method='auto') (#2719)

* Initial commit with ability to add transforms to GroupChat

* Added tests

* Tidy up

* Tidy up of variable names and commented out test

* Tidy up comment

* Update import to relative

* Added documentation topic for transform messages for speaker selection.

* Added handling of non-core module, transforms, in groupchat

* Adjusted parameter if module not loaded.

* Updated groupchat test which failed during CI/CD

---------

Co-authored-by: Li Jiang <[email protected]>

autogen/agentchat/contrib/capabilities/transform_messages.py

@@ -1,9 +1,8 @@
 import copy
 from typing import Dict, List
 
-from autogen import ConversableAgent
-
 from ....formatting_utils import colored
+from ...conversable_agent import ConversableAgent
 from .transforms import MessageTransform
 
 

autogen/agentchat/groupchat.py

@@ -5,7 +5,7 @@
 import re
 import sys
 from dataclasses import dataclass, field
-from typing import Callable, Dict, List, Literal, Optional, Tuple, Union
+from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Union
 
 from ..code_utils import content_str
 from ..exception_utils import AgentNameConflict, NoEligibleSpeaker, UndefinedNextAgent
@@ -17,6 +17,12 @@
 from .chat import ChatResult
 from .conversable_agent import ConversableAgent
 
+try:
+    # Non-core module
+    from .contrib.capabilities import transform_messages
+except ImportError:
+    transform_messages = None
+
 logger = logging.getLogger(__name__)
 
 
@@ -76,6 +82,8 @@ def custom_speaker_selection_func(
         of times until a single agent is returned or it exhausts the maximum attempts.
         Applies only to "auto" speaker selection method.
         Default is 2.
+    - select_speaker_transform_messages: (optional) the message transformations to apply to the nested select speaker agent-to-agent chat messages.
+        Takes a TransformMessages object, defaults to None and is only utilised when the speaker selection method is "auto".
     - select_speaker_auto_verbose: whether to output the select speaker responses and selections
         If set to True, the outputs from the two agents in the nested select speaker chat will be output, along with
         whether the responses were successful, or not, in selecting an agent
@@ -132,6 +140,7 @@ def custom_speaker_selection_func(
     The names are case-sensitive and should not be abbreviated or changed.
     The only names that are accepted are {agentlist}.
     Respond with ONLY the name of the speaker and DO NOT provide a reason."""
+    select_speaker_transform_messages: Optional[Any] = None
     select_speaker_auto_verbose: Optional[bool] = False
     role_for_select_speaker_messages: Optional[str] = "system"
 
@@ -249,6 +258,20 @@ def __post_init__(self):
         elif self.max_retries_for_selecting_speaker < 0:
             raise ValueError("max_retries_for_selecting_speaker must be greater than or equal to zero")
 
+        # Load message transforms here (load once for the Group Chat so we don't have to re-initiate it and it maintains the cache across subsequent select speaker calls)
+        self._speaker_selection_transforms = None
+        if self.select_speaker_transform_messages is not None:
+            if transform_messages is not None:
+                if isinstance(self.select_speaker_transform_messages, transform_messages.TransformMessages):
+                    self._speaker_selection_transforms = self.select_speaker_transform_messages
+                else:
+                    raise ValueError("select_speaker_transform_messages must be None or MessageTransforms.")
+            else:
+                logger.warning(
+                    "TransformMessages could not be loaded, the 'select_speaker_transform_messages' transform"
+                    "will not apply."
+                )
+
         # Validate select_speaker_auto_verbose
         if self.select_speaker_auto_verbose is None or not isinstance(self.select_speaker_auto_verbose, bool):
             raise ValueError("select_speaker_auto_verbose cannot be None or non-bool")
@@ -655,6 +678,10 @@ def validate_speaker_name(recipient, messages, sender, config) -> Tuple[bool, Un
         else:
             start_message = messages[-1]
 
+        # Add the message transforms, if any, to the speaker selection agent
+        if self._speaker_selection_transforms is not None:
+            self._speaker_selection_transforms.add_to_agent(speaker_selection_agent)
+
         # Run the speaker selection chat
         result = checking_agent.initiate_chat(
             speaker_selection_agent,
@@ -749,6 +776,10 @@ def validate_speaker_name(recipient, messages, sender, config) -> Tuple[bool, Un
         else:
             start_message = messages[-1]
 
+        # Add the message transforms, if any, to the speaker selection agent
+        if self._speaker_selection_transforms is not None:
+            self._speaker_selection_transforms.add_to_agent(speaker_selection_agent)
+
         # Run the speaker selection chat
         result = await checking_agent.a_initiate_chat(
             speaker_selection_agent,

test/agentchat/test_groupchat.py

@@ -12,6 +12,7 @@
 
 import autogen
 from autogen import Agent, AssistantAgent, GroupChat, GroupChatManager
+from autogen.agentchat.contrib.capabilities import transform_messages, transforms
 from autogen.exception_utils import AgentNameConflict, UndefinedNextAgent
 
 
@@ -2061,6 +2062,46 @@ def test_manager_resume_messages():
         return_agent, return_message = manager.resume(messages="Let's get this conversation started.")
 
 
+def test_select_speaker_transform_messages():
+    """Tests adding transform messages to a GroupChat for speaker selection when in 'auto' mode"""
+
+    # Test adding a TransformMessages to a group chat
+    test_add_transforms = transform_messages.TransformMessages(
+        transforms=[
+            transforms.MessageHistoryLimiter(max_messages=10),
+            transforms.MessageTokenLimiter(max_tokens=3000, max_tokens_per_message=500, min_tokens=300),
+        ]
+    )
+
+    coder = AssistantAgent(name="Coder", llm_config=None)
+    groupchat = GroupChat(messages=[], agents=[coder], select_speaker_transform_messages=test_add_transforms)
+
+    # Ensure the transform have been added to the GroupChat
+    assert groupchat._speaker_selection_transforms == test_add_transforms
+
+    # Attempt to add a non MessageTransforms object, such as a list of transforms
+    with pytest.raises(ValueError, match="select_speaker_transform_messages must be None or MessageTransforms."):
+        groupchat = GroupChat(
+            messages=[],
+            agents=[coder],
+            select_speaker_transform_messages=[transforms.MessageHistoryLimiter(max_messages=10)],
+        )
+
+    # Ensure if we don't pass any transforms in, none are on the GroupChat
+    groupchat_missing = GroupChat(messages=[], agents=[coder])
+
+    assert groupchat_missing._speaker_selection_transforms is None
+
+    # Ensure we can pass in None
+    groupchat_none = GroupChat(
+        messages=[],
+        agents=[coder],
+        select_speaker_transform_messages=None,
+    )
+
+    assert groupchat_none._speaker_selection_transforms is None
+
+
 if __name__ == "__main__":
     # test_func_call_groupchat()
     # test_broadcast()
@@ -2089,4 +2130,5 @@ def test_manager_resume_messages():
     test_manager_resume_functions()
     # test_manager_resume_returns()
     # test_manager_resume_messages()
+    # test_select_speaker_transform_messages()
     pass