From b36bf27bee7427f61d448f188480491223ea643e Mon Sep 17 00:00:00 2001 From: Wael Karkoub Date: Fri, 12 Apr 2024 06:03:32 +0100 Subject: [PATCH] adds fix (#2364) --- .../contrib/capabilities/transforms.py | 9 +++++++- .../capabilities/test_transform_messages.py | 23 +++++++++++++++++++ 2 files changed, 31 insertions(+), 1 deletion(-) diff --git a/autogen/agentchat/contrib/capabilities/transforms.py b/autogen/agentchat/contrib/capabilities/transforms.py index f2ba6719118..cc4faace3f1 100644 --- a/autogen/agentchat/contrib/capabilities/transforms.py +++ b/autogen/agentchat/contrib/capabilities/transforms.py @@ -126,9 +126,16 @@ def apply_transform(self, messages: List[Dict]) -> List[Dict]: processed_messages_tokens = 0 # calculate tokens for all messages - total_tokens = sum(_count_tokens(msg["content"]) for msg in temp_messages) + total_tokens = sum( + _count_tokens(msg["content"]) for msg in temp_messages if isinstance(msg.get("content"), (str, list)) + ) for msg in reversed(temp_messages): + # Some messages may not have content. + if not isinstance(msg.get("content"), (str, list)): + processed_messages.insert(0, msg) + continue + expected_tokens_remained = self._max_tokens - processed_messages_tokens - self._max_tokens_per_message # If adding this message would exceed the token limit, truncate the last message to meet the total token diff --git a/test/agentchat/contrib/capabilities/test_transform_messages.py b/test/agentchat/contrib/capabilities/test_transform_messages.py index c5b7c1dcf2d..ac0cdf58755 100644 --- a/test/agentchat/contrib/capabilities/test_transform_messages.py +++ b/test/agentchat/contrib/capabilities/test_transform_messages.py @@ -76,6 +76,29 @@ def test_limit_token_transform(): assert len(transformed_messages) <= len(messages) +def test_limit_token_transform_without_content(): + """Test the TokenLimitTransform with messages that don't have content.""" + + messages = [{"role": "user", "function_call": "example"}, {"role": "assistant", "content": None}] + + # check if token limit per message works nicely with total token limit. + token_limit_transform = MessageTokenLimiter(max_tokens=10, max_tokens_per_message=5) + + transformed_messages = token_limit_transform.apply_transform(copy.deepcopy(messages)) + + assert len(transformed_messages) == len(messages) + + +def test_limit_token_transform_total_token_count(): + """Tests if the TokenLimitTransform truncates without dropping messages.""" + messages = [{"role": "very very very very very"}] + + token_limit_transform = MessageTokenLimiter(max_tokens=1) + transformed_messages = token_limit_transform.apply_transform(copy.deepcopy(messages)) + + assert len(transformed_messages) == 1 + + def test_max_message_history_length_transform(): """ Test the MessageHistoryLimiter capability to limit the number of messages.