Skip to content

Commit

Permalink
Added tests
Browse files Browse the repository at this point in the history
  • Loading branch information
marklysze committed May 19, 2024
1 parent aea58fa commit 9de72c0
Show file tree
Hide file tree
Showing 2 changed files with 45 additions and 3 deletions.
6 changes: 3 additions & 3 deletions autogen/agentchat/groupchat.py
Original file line number Diff line number Diff line change
Expand Up @@ -256,7 +256,7 @@ def __post_init__(self):
if isinstance(self.select_speaker_auto_message_transforms, transform_messages.TransformMessages):
self._auto_message_transforms = self.select_speaker_auto_message_transforms
else:
raise ValueError("select_speaker_auto_message_transforms must be None or MessageTransforms")
raise ValueError("select_speaker_auto_message_transforms must be None or MessageTransforms.")
else:
self._auto_message_transforms = None

Expand Down Expand Up @@ -738,8 +738,8 @@ def validate_speaker_name(recipient, messages, sender, config) -> Tuple[bool, Un
)

# Add the message transforms, if any, to the speaker selection agent
if self._auto_message_transforms_context_handling is not None:
self._auto_message_transforms_context_handling.add_to_agent(speaker_selection_agent)
if self._auto_message_transforms is not None:
self._auto_message_transforms.add_to_agent(speaker_selection_agent)

# Run the speaker selection chat
result = await checking_agent.a_initiate_chat(
Expand Down
42 changes: 42 additions & 0 deletions test/agentchat/test_groupchat.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down Expand Up @@ -2007,6 +2008,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_auto_message_transforms=test_add_transforms)

# Ensure the transform have been added to the GroupChat
assert groupchat._auto_message_transforms == test_add_transforms

# Attempt to add a non MessageTransforms object, such as a list of transforms
with pytest.raises(ValueError, match="select_speaker_auto_message_transforms must be None or MessageTransforms."):
groupchat = GroupChat(
messages=[],
agents=[coder],
select_speaker_auto_message_transforms=List[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._auto_message_transforms is None

# Ensure we can pass in None
groupchat_none = GroupChat(
messages=[],
agents=[coder],
select_speaker_auto_message_transforms=None,
)

assert groupchat_none._auto_message_transforms is None


if __name__ == "__main__":
# test_func_call_groupchat()
# test_broadcast()
Expand Down Expand Up @@ -2037,4 +2078,5 @@ def test_manager_resume_messages():
# test_manager_resume_functions()
# test_manager_resume_returns()
# test_manager_resume_messages()
test_select_speaker_transform_messages()
pass

0 comments on commit 9de72c0

Please sign in to comment.