Refactor transform_messages (#1631)

* refactored code to simplify

* optimize function. Instead of iterating over each character, guess at size and then iterate by token.

* adding tests

* Add missing tests

* minor test fix

* simplified token truncation by using tiktoken to encode and decode

* updated truncated notification message

* Fix llm_config spec to use os.environ

* Add test case and fix bug in loop

---------

Co-authored-by: gagb <[email protected]>

Author: dkirsche

autogen/agentchat/contrib/capabilities/context_handling.py

@@ -3,6 +3,7 @@
 from typing import Dict, Optional, List
 from autogen import ConversableAgent
 from autogen import token_count_utils
+import tiktoken
 
 
 class TransformChatHistory:
@@ -53,56 +54,70 @@ def _transform_messages(self, messages: List[Dict]) -> List[Dict]:
             messages: List of messages to process.
 
         Returns:
-            List of messages with the first system message and the last max_messages messages.
+            List of messages with the first system message and the last max_messages messages,
+            ensuring each message does not exceed max_tokens_per_message.
         """
+        temp_messages = messages.copy()
         processed_messages = []
-        messages = messages.copy()
-        rest_messages = messages
+        system_message = None
+        processed_messages_tokens = 0
 
-        # check if the first message is a system message and append it to the processed messages
-        if len(messages) > 0:
-            if messages[0]["role"] == "system":
-                msg = messages[0]
-                processed_messages.append(msg)
-                rest_messages = messages[1:]
+        if messages[0]["role"] == "system":
+            system_message = messages[0].copy()
+            temp_messages.pop(0)
 
-        processed_messages_tokens = 0
-        for msg in messages:
-            msg["content"] = truncate_str_to_tokens(msg["content"], self.max_tokens_per_message)
+        total_tokens = sum(
+            token_count_utils.count_token(msg["content"]) for msg in temp_messages
+        )  # Calculate tokens for all messages
 
-        # iterate through rest of the messages and append them to the processed messages
-        for msg in rest_messages[-self.max_messages :]:
+        # Truncate each message's content to a maximum token limit of each message
+
+        # Process recent messages first
+        for msg in reversed(temp_messages[-self.max_messages :]):
+            msg["content"] = truncate_str_to_tokens(msg["content"], self.max_tokens_per_message)
             msg_tokens = token_count_utils.count_token(msg["content"])
             if processed_messages_tokens + msg_tokens > self.max_tokens:
                 break
-            processed_messages.append(msg)
+            # append the message to the beginning of the list to preserve order
+            processed_messages = [msg] + processed_messages
             processed_messages_tokens += msg_tokens
-
-        total_tokens = 0
-        for msg in messages:
-            total_tokens += token_count_utils.count_token(msg["content"])
-
+        if system_message:
+            processed_messages.insert(0, system_message)
+        # Optionally, log the number of truncated messages and tokens if needed
         num_truncated = len(messages) - len(processed_messages)
+
         if num_truncated > 0 or total_tokens > processed_messages_tokens:
-            print(colored(f"Truncated {len(messages) - len(processed_messages)} messages.", "yellow"))
-            print(colored(f"Truncated {total_tokens - processed_messages_tokens} tokens.", "yellow"))
+            print(
+                colored(
+                    f"Truncated {num_truncated} messages. Reduced from {len(messages)} to {len(processed_messages)}.",
+                    "yellow",
+                )
+            )
+            print(
+                colored(
+                    f"Truncated {total_tokens - processed_messages_tokens} tokens. Tokens reduced from {total_tokens} to {processed_messages_tokens}",
+                    "yellow",
+                )
+            )
         return processed_messages
 
 
-def truncate_str_to_tokens(text: str, max_tokens: int) -> str:
-    """
-    Truncate a string so that number of tokens in less than max_tokens.
+def truncate_str_to_tokens(text: str, max_tokens: int, model: str = "gpt-3.5-turbo-0613") -> str:
+    """Truncate a string so that the number of tokens is less than or equal to max_tokens using tiktoken.
 
     Args:
-        content: String to process.
-        max_tokens: Maximum number of tokens to keep.
+        text: The string to truncate.
+        max_tokens: The maximum number of tokens to keep.
+        model: The target OpenAI model for tokenization alignment.
 
     Returns:
-        Truncated string.
+        The truncated string.
     """
-    truncated_string = ""
-    for char in text:
-        truncated_string += char
-        if token_count_utils.count_token(truncated_string) == max_tokens:
-            break
-    return truncated_string
+
+    encoding = tiktoken.encoding_for_model(model)  # Get the appropriate tokenizer
+
+    encoded_tokens = encoding.encode(text)
+    truncated_tokens = encoded_tokens[:max_tokens]
+    truncated_text = encoding.decode(truncated_tokens)  # Decode back to text
+
+    return truncated_text