From a34e4cc515ae85c877a5e1ea59cc0bacf1320cd9 Mon Sep 17 00:00:00 2001 From: dkirsche Date: Tue, 20 Feb 2024 17:53:05 -0500 Subject: [PATCH] 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 --- .../contrib/capabilities/context_handling.py | 85 +-- ...hat_capability_long_context_handling.ipynb | 558 ++++++++++-------- .../capabilities/test_context_handling.py | 123 ++++ 3 files changed, 487 insertions(+), 279 deletions(-) diff --git a/autogen/agentchat/contrib/capabilities/context_handling.py b/autogen/agentchat/contrib/capabilities/context_handling.py index d33ef45bc..d12af6051 100644 --- a/autogen/agentchat/contrib/capabilities/context_handling.py +++ b/autogen/agentchat/contrib/capabilities/context_handling.py @@ -3,6 +3,7 @@ from termcolor import colored from typing import Dict, Optional, List from autogen import ConversableAgent from autogen import token_count_utils +import tiktoken class TransformChatHistory: @@ -53,56 +54,70 @@ class TransformChatHistory: 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 - - # 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:] - + system_message = None processed_messages_tokens = 0 - for msg in messages: - msg["content"] = truncate_str_to_tokens(msg["content"], self.max_tokens_per_message) - # iterate through rest of the messages and append them to the processed messages - for msg in rest_messages[-self.max_messages :]: + if messages[0]["role"] == "system": + system_message = messages[0].copy() + temp_messages.pop(0) + + total_tokens = sum( + token_count_utils.count_token(msg["content"]) for msg in temp_messages + ) # Calculate tokens for all 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 diff --git a/notebook/agentchat_capability_long_context_handling.ipynb b/notebook/agentchat_capability_long_context_handling.ipynb index 222e2e473..edeb8216b 100644 --- a/notebook/agentchat_capability_long_context_handling.ipynb +++ b/notebook/agentchat_capability_long_context_handling.ipynb @@ -25,10 +25,22 @@ "metadata": {}, "outputs": [], "source": [ + "import os\n", "import autogen\n", "from autogen.agentchat.contrib.capabilities import context_handling" ] }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [], + "source": [ + "llm_config = {\n", + " \"config_list\": [{\"model\": \"gpt-3.5-turbo\", \"api_key\": os.environ.get(\"OPENAI_API_KEY\")}],\n", + "}" + ] + }, { "cell_type": "markdown", "metadata": {}, @@ -38,7 +50,7 @@ }, { "cell_type": "code", - "execution_count": 3, + "execution_count": 4, "metadata": {}, "outputs": [ { @@ -49,16 +61,10 @@ "\n", "plot and save a graph of x^2 from -10 to 10\n", "\n", - "--------------------------------------------------------------------------------\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ + "--------------------------------------------------------------------------------\n", "\u001b[33massistant\u001b[0m (to user_proxy):\n", "\n", - "Here's a Python code snippet to plot and save a graph of x^2 from -10 to 10:\n", + "To plot and save a graph of the function x^2 from -10 to 10, you can use the matplotlib library in Python. Here is the code:\n", "\n", "```python\n", "import matplotlib.pyplot as plt\n", @@ -67,7 +73,143 @@ "# Generate x values from -10 to 10\n", "x = np.linspace(-10, 10, 100)\n", "\n", - "# Evaluate y values using x^2\n", + "# Calculate corresponding y values (x^2)\n", + "y = x**2\n", + "\n", + "# Create the plot\n", + "plt.plot(x, y)\n", + "\n", + "# Add labels and title\n", + "plt.xlabel('x')\n", + "plt.ylabel('y')\n", + "plt.title('Plot of x^2')\n", + "\n", + "# Save the plot as a file\n", + "plt.savefig('x_squared_plot.png')\n", + "\n", + "# Show the plot\n", + "plt.show()\n", + "```\n", + "\n", + "This code will create a plot of the function x^2 and save it as \"x_squared_plot.png\" in the current directory. Make sure you have the matplotlib library installed before running this code.\n", + "\n", + "--------------------------------------------------------------------------------\n", + "\u001b[31m\n", + ">>>>>>>> EXECUTING CODE BLOCK 0 (inferred language is python)...\u001b[0m\n", + "\u001b[33muser_proxy\u001b[0m (to assistant):\n", + "\n", + "exitcode: 0 (execution succeeded)\n", + "Code output: \n", + "Figure(640x480)\n", + "\n", + "\n", + "--------------------------------------------------------------------------------\n", + "\u001b[33mTruncated 0 messages. Reduced from 3 to 3.\u001b[0m\n", + "\u001b[33mTruncated 139 tokens. Tokens reduced from 223 to 84\u001b[0m\n", + "\u001b[33massistant\u001b[0m (to user_proxy):\n", + "\n", + "Great! The code executed successfully and generated a plot of the function x^2 from -10 to 10. The plot was displayed in a figure with size 640x480. \n", + "\n", + "To save the graph as an image file, you can modify the code as follows:\n", + "\n", + "```python\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "\n", + "# Generate x values from -10 to 10\n", + "x = np.linspace(-10, 10, 100)\n", + "\n", + "# Generate y values by squaring x\n", + "y = x ** 2\n", + "\n", + "# Plot the graph\n", + "plt.plot(x, y)\n", + "plt.xlabel('x')\n", + "plt.ylabel('x^2')\n", + "plt.title('Graph of x^2')\n", + "plt.grid(True)\n", + "\n", + "# Save the graph as an image file, for example as 'graph.png'\n", + "plt.savefig('graph.png')\n", + "```\n", + "\n", + "By executing this updated code, the graph will be saved as an image file named 'graph.png' in the same directory as your Python script.\n", + "\n", + "Please let me know if you need any further assistance.\n", + "\n", + "--------------------------------------------------------------------------------\n", + "\u001b[31m\n", + ">>>>>>>> EXECUTING CODE BLOCK 0 (inferred language is python)...\u001b[0m\n", + "\u001b[33muser_proxy\u001b[0m (to assistant):\n", + "\n", + "exitcode: 0 (execution succeeded)\n", + "Code output: \n", + "\n", + "\n", + "--------------------------------------------------------------------------------\n", + "\u001b[33mTruncated 0 messages. Reduced from 5 to 5.\u001b[0m\n", + "\u001b[33mTruncated 159 tokens. Tokens reduced from 306 to 147\u001b[0m\n", + "\u001b[33massistant\u001b[0m (to user_proxy):\n", + "\n", + "Great! The code executed successfully and generated a plot of the function x^2 from -10 to 10. The plot was displayed in a figure with size 640x480.\n", + "\n", + "To save the graph as an image file, you can modify the code as follows:\n", + "\n", + "```python\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "\n", + "# Generate x values from -10 to 10\n", + "x = np.linspace(-10, 10, 100)\n", + "\n", + "# Calculate y values (x^2)\n", + "y = x**2\n", + "\n", + "# Plot the graph\n", + "plt.plot(x, y)\n", + "\n", + "# Add labels and title\n", + "plt.xlabel('x')\n", + "plt.ylabel('y')\n", + "plt.title('Graph of x^2')\n", + "\n", + "# Save the graph as an image file\n", + "plt.savefig('graph.png')\n", + "\n", + "# Close the plot\n", + "plt.close()\n", + "```\n", + "\n", + "This code will save the plot as an image file named \"graph.png\" in the current directory. You can change the filename and path if needed.\n", + "\n", + "Please let me know if you need any further assistance.\n", + "\n", + "--------------------------------------------------------------------------------\n", + "\u001b[31m\n", + ">>>>>>>> EXECUTING CODE BLOCK 0 (inferred language is python)...\u001b[0m\n", + "\u001b[33muser_proxy\u001b[0m (to assistant):\n", + "\n", + "exitcode: 0 (execution succeeded)\n", + "Code output: \n", + "\n", + "\n", + "--------------------------------------------------------------------------------\n", + "\u001b[33mTruncated 0 messages. Reduced from 7 to 7.\u001b[0m\n", + "\u001b[33mTruncated 159 tokens. Tokens reduced from 369 to 210\u001b[0m\n", + "\u001b[33massistant\u001b[0m (to user_proxy):\n", + "\n", + "Great! The code executed successfully and generated a plot of the function x^2 from -10 to 10. The plot was displayed in a figure with size 640x480.\n", + "\n", + "To save the graph as an image file, you can modify the code as follows:\n", + "\n", + "```python\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "\n", + "# Generate x values\n", + "x = np.linspace(-10, 10, 100)\n", + "\n", + "# Generate y values\n", "y = x**2\n", "\n", "# Plot the graph\n", @@ -76,90 +218,35 @@ "# Set labels and title\n", "plt.xlabel('x')\n", "plt.ylabel('y')\n", - "plt.title('Graph of y = x^2')\n", + "plt.title('Graph of x^2')\n", "\n", - "# Save the graph as an image file (e.g., PNG)\n", - "plt.savefig('x_squared_graph.png')\n", - "\n", - "# Show the graph\n", - "plt.show()\n", - "```\n", - "\n", - "Please make sure to have the `matplotlib` library installed in your Python environment. After executing the code, the graph will be saved as \"x_squared_graph.png\" in the current directory. You can change the filename if desired.\n", - "\n", - "--------------------------------------------------------------------------------\n", - "\u001b[31m\n", - ">>>>>>>> EXECUTING CODE BLOCK 0 (inferred language is python)...\u001b[0m\n", - "\u001b[33muser_proxy\u001b[0m (to assistant):\n", - "\n", - "exitcode: 0 (execution succeeded)\n", - "Code output: \n", - "Figure(640x480)\n", - "\n", - "\n", - "--------------------------------------------------------------------------------\n", - "\u001b[33massistant\u001b[0m (to user_proxy):\n", - "\n", - "Great! The code executed successfully and generated a graph of x^2 from -10 to 10. You can save the graph by adding the following code snippet:\n", - "\n", - "```python\n", - "# Save the graph as a file\n", + "# Save the graph as an image file\n", "plt.savefig('graph.png')\n", "```\n", "\n", - "This will save the graph as a PNG file named \"graph.png\" in your current working directory.\n", + "This code will save the graph as a PNG image file named 'graph.png' in the same directory where you are running the script. You can change the filename and format (e.g., 'graph.jpg') if desired.\n", "\n", - "Now, you can check the saved graph in your current directory. Let me know if you need any further assistance.\n", + "Please note that you will need to have the matplotlib library installed to run this code. If you don't have it installed, you can install it by running `pip install matplotlib` in your terminal or command prompt.\n", + "\n", + "Let me know if you need any further assistance!\n", "\n", "--------------------------------------------------------------------------------\n", "\u001b[31m\n", ">>>>>>>> EXECUTING CODE BLOCK 0 (inferred language is python)...\u001b[0m\n", "\u001b[33muser_proxy\u001b[0m (to assistant):\n", "\n", - "exitcode: 1 (execution failed)\n", - "Code output: \n", - "Traceback (most recent call last):\n", - " File \"\", line 2, in \n", - " plt.savefig('graph.png')\n", - "NameError: name 'plt' is not defined\n", - "\n", - "\n", - "--------------------------------------------------------------------------------\n", - "\u001b[33massistant\u001b[0m (to user_proxy):\n", - "\n", - "I apologize for the mistake. It seems that the `plt` module was not imported correctly. Please make sure that you have the `matplotlib` library installed. You can install it using the following command:\n", - "\n", - "```sh\n", - "pip install matplotlib\n", - "```\n", - "\n", - "Once you have the library installed, please try running the code again. If you are still facing any issues, please let me know.\n", - "\n", - "--------------------------------------------------------------------------------\n", - "\u001b[31m\n", - ">>>>>>>> EXECUTING CODE BLOCK 0 (inferred language is sh)...\u001b[0m\n", - "\u001b[33muser_proxy\u001b[0m (to assistant):\n", - "\n", "exitcode: 0 (execution succeeded)\n", "Code output: \n", - "Defaulting to user installation because normal site-packages is not writeable\n", - "Requirement already satisfied: matplotlib in /home/vscode/.local/lib/python3.10/site-packages (3.8.2)\n", - "Requirement already satisfied: pyparsing>=2.3.1 in /home/vscode/.local/lib/python3.10/site-packages (from matplotlib) (3.1.1)\n", - "Requirement already satisfied: fonttools>=4.22.0 in /home/vscode/.local/lib/python3.10/site-packages (from matplotlib) (4.47.2)\n", - "Requirement already satisfied: python-dateutil>=2.7 in /home/vscode/.local/lib/python3.10/site-packages (from matplotlib) (2.8.2)\n", - "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.10/site-packages (from matplotlib) (23.2)\n", - "Requirement already satisfied: pillow>=8 in /home/vscode/.local/lib/python3.10/site-packages (from matplotlib) (10.2.0)\n", - "Requirement already satisfied: cycler>=0.10 in /home/vscode/.local/lib/python3.10/site-packages (from matplotlib) (0.12.1)\n", - "Requirement already satisfied: kiwisolver>=1.3.1 in /home/vscode/.local/lib/python3.10/site-packages (from matplotlib) (1.4.5)\n", - "Requirement already satisfied: numpy<2,>=1.21 in /home/vscode/.local/lib/python3.10/site-packages (from matplotlib) (1.26.3)\n", - "Requirement already satisfied: contourpy>=1.0.1 in /home/vscode/.local/lib/python3.10/site-packages (from matplotlib) (1.2.0)\n", - "Requirement already satisfied: six>=1.5 in /home/vscode/.local/lib/python3.10/site-packages (from python-dateutil>=2.7->matplotlib) (1.16.0)\n", "\n", "\n", "--------------------------------------------------------------------------------\n", + "\u001b[33mTruncated 0 messages. Reduced from 9 to 9.\u001b[0m\n", + "\u001b[33mTruncated 198 tokens. Tokens reduced from 471 to 273\u001b[0m\n", "\u001b[33massistant\u001b[0m (to user_proxy):\n", "\n", - "Thank you for installing the `matplotlib` library. Let's try running the code again:\n", + "Great! The code executed successfully and generated a plot of the function x^2 from -10 to 10. The plot was displayed in a figure with size 640x480.\n", + "\n", + "To save the graph as an image file, you can modify the code as follows:\n", "\n", "```python\n", "import matplotlib.pyplot as plt\n", @@ -168,20 +255,25 @@ "# Generate x values from -10 to 10\n", "x = np.linspace(-10, 10, 100)\n", "\n", - "# Compute the y values (x^2)\n", + "# Generate y values by squaring x\n", "y = x**2\n", "\n", "# Plot the graph\n", "plt.plot(x, y)\n", "\n", - "# Save the graph as a file\n", - "plt.savefig('graph.png')\n", + "# Add labels and title\n", + "plt.xlabel('x')\n", + "plt.ylabel('x^2')\n", + "plt.title('Graph of x^2')\n", "\n", - "# Show the graph\n", + "# Save the graph as an image file\n", + "plt.savefig('x_squared_plot.png')\n", + "\n", + "# Display the graph\n", "plt.show()\n", "```\n", "\n", - "Please give it a try and let me know if it works for you.\n", + "This code will save the graph as a PNG image file named \"x_squared_plot.png\" in the current working directory. You can customize the filename and file format according to your needs.\n", "\n", "--------------------------------------------------------------------------------\n", "\u001b[31m\n", @@ -194,123 +286,39 @@ "\n", "\n", "--------------------------------------------------------------------------------\n", + "\u001b[33mTruncated 1 messages. Reduced from 11 to 10.\u001b[0m\n", + "\u001b[33mTruncated 174 tokens. Tokens reduced from 501 to 327\u001b[0m\n", "\u001b[33massistant\u001b[0m (to user_proxy):\n", "\n", - "It seems that the code executed successfully and generated a figure of x^2 from -10 to 10. To save the graph as an image file, you can add the following code snippet:\n", + "Great! The code executed successfully and generated a plot of the function x^2 from -10 to 10. The plot was displayed in a figure with size 640x480.\n", "\n", - "```python\n", - "# Save the graph as a file\n", - "plt.savefig('graph.png')\n", - "```\n", - "\n", - "This will save the graph as a PNG file named \"graph.png\" in the current directory. You can change the file name or modify the file format as needed.\n", - "Remember to replace `plt.savefig('graph.png')` with your desired file name if necessary.\n", - "\n", - "--------------------------------------------------------------------------------\n", - "\u001b[31m\n", - ">>>>>>>> EXECUTING CODE BLOCK 0 (inferred language is python)...\u001b[0m\n", - "\u001b[33muser_proxy\u001b[0m (to assistant):\n", - "\n", - "exitcode: 1 (execution failed)\n", - "Code output: \n", - "Traceback (most recent call last):\n", - " File \"\", line 2, in \n", - " plt.savefig('graph.png')\n", - "NameError: name 'plt' is not defined\n", - "\n", - "\n", - "--------------------------------------------------------------------------------\n", - "\u001b[33mTruncated 1 messages.\u001b[0m\n", - "\u001b[33mTruncated 15 tokens.\u001b[0m\n", - "\u001b[33massistant\u001b[0m (to user_proxy):\n", - "\n", - "My apologies for the mistake again. It seems that the `plt` module was not imported correctly. Please make sure that you have imported the `matplotlib.pyplot` module.\n", - "\n", - "Here is the correct code snippet:\n", + "To save the graph as an image file, you can modify the code as follows:\n", "\n", "```python\n", "import matplotlib.pyplot as plt\n", "import numpy as np\n", "\n", - "# Generate x values from -10 to 10\n", - "x = np.linspace(-10, 10, 100)\n", - "y = x ** 2\n", - "\n", - "# Create a plot of x^2\n", - "plt.plot(x, y)\n", - "\n", - "# Save the graph as a file\n", - "plt.savefig('graph.png')\n", - "```\n", - "\n", - "Please try running this code. It should generate the graph and save it as `graph.png` in the current working directory.\n", - "\n", - "--------------------------------------------------------------------------------\n", - "\u001b[31m\n", - ">>>>>>>> EXECUTING CODE BLOCK 0 (inferred language is python)...\u001b[0m\n", - "\u001b[33muser_proxy\u001b[0m (to assistant):\n", - "\n", - "exitcode: 0 (execution succeeded)\n", - "Code output: \n", - "\n", - "\n", - "--------------------------------------------------------------------------------\n", - "\u001b[33mTruncated 3 messages.\u001b[0m\n", - "\u001b[33mTruncated 84 tokens.\u001b[0m\n", - "\u001b[33massistant\u001b[0m (to user_proxy):\n", - "\n", - "Great! The code executed successfully and generated a graph of x^2 from -10 to 10. You can save the graph by adding the following code snippet:\n", - "\n", - "```python\n", - "# Save the graph as a file\n", - "plt.savefig('graph.png')\n", - "```\n", - "\n", - "This will save the graph as a PNG image file named \"graph.png\" in the same directory as your Python script. You can change the filename and format (e.g., 'graph.jpg', 'graph.pdf') as needed.\n", - "\n", - "--------------------------------------------------------------------------------\n", - "\u001b[31m\n", - ">>>>>>>> EXECUTING CODE BLOCK 0 (inferred language is python)...\u001b[0m\n", - "\u001b[33muser_proxy\u001b[0m (to assistant):\n", - "\n", - "exitcode: 1 (execution failed)\n", - "Code output: \n", - "Traceback (most recent call last):\n", - " File \"\", line 2, in \n", - " plt.savefig('graph.png')\n", - "NameError: name 'plt' is not defined\n", - "\n", - "\n", - "--------------------------------------------------------------------------------\n", - "\u001b[33mTruncated 5 messages.\u001b[0m\n", - "\u001b[33mTruncated 184 tokens.\u001b[0m\n", - "\u001b[33massistant\u001b[0m (to user_proxy):\n", - "\n", - "I apologize for the confusion. It looks like the code didn't execute as expected. \n", - "\n", - "To save the graph in a file, try the following code:\n", - "\n", - "```python\n", - "import matplotlib.pyplot as plt\n", - "import numpy as np\n", - "\n", - "# Generate x values from -10 to 10\n", + "# Generate x values\n", "x = np.linspace(-10, 10, 100)\n", "\n", - "# Calculate the y values (square of x)\n", + "# Generate y values\n", "y = x ** 2\n", "\n", - "# Plot the graph\n", + "# Create the plot\n", "plt.plot(x, y)\n", + "plt.xlabel('x')\n", + "plt.ylabel('y')\n", + "plt.title('Graph of x^2')\n", "\n", - "# Save the graph as a file\n", - "plt.savefig('graph.png')\n", + "# Save the plot as an image file\n", + "plt.savefig('x_squared.png')\n", "\n", - "# Show the graph\n", "plt.show()\n", "```\n", "\n", - "Make sure you have the `matplotlib` library installed. Let me know if you encounter any further issues.\n", + "The code will save the plot as a PNG image file named \"x_squared.png\" in the current directory. You can change the filename or file extension as needed.\n", + "\n", + "Please let me know if you need any further assistance!\n", "\n", "--------------------------------------------------------------------------------\n", "\u001b[31m\n", @@ -323,35 +331,30 @@ "\n", "\n", "--------------------------------------------------------------------------------\n", - "\u001b[33mTruncated 7 messages.\u001b[0m\n", - "\u001b[33mTruncated 284 tokens.\u001b[0m\n", + "\u001b[33mTruncated 3 messages. Reduced from 13 to 10.\u001b[0m\n", + "\u001b[33mTruncated 227 tokens. Tokens reduced from 554 to 327\u001b[0m\n", "\u001b[33massistant\u001b[0m (to user_proxy):\n", "\n", - "Great! The code executed successfully and generated a graph of x^2 from -10 to 10. You can save the graph by adding the following code snippet:\n", + "To save the graph as an image file, you can modify the code as follows:\n", "\n", "```python\n", - "# Save the graph as a file\n", - "plt.savefig('graph.png')\n", + "import matplotlib.pyplot as plt\n", "\n", - "--------------------------------------------------------------------------------\n", - "\u001b[33muser_proxy\u001b[0m (to assistant):\n", + "# Create the plot\n", + "plt.plot(x, y)\n", "\n", + "# Set the plot title and axes labels\n", + "plt.title(\"Plot of x^2\")\n", + "plt.xlabel(\"x\")\n", + "plt.ylabel(\"y\")\n", "\n", - "\n", - "--------------------------------------------------------------------------------\n", - "\u001b[33mTruncated 9 messages.\u001b[0m\n", - "\u001b[33mTruncated 353 tokens.\u001b[0m\n", - "\u001b[33massistant\u001b[0m (to user_proxy):\n", - "\n", - "Sorry for the mistake again. The code snippet provided was incomplete. To correctly save the graph as an image file, you need to add a line of code to close the plot after saving. Here's the updated code:\n", - "\n", - "```python\n", - "# Save the graph as a file\n", - "plt.savefig('graph.png')\n", - "plt.close()\n", + "# Save the plot as an image file\n", + "plt.savefig(\"plot.png\")\n", "```\n", "\n", - "Please try running the complete code again and it should save the graph as \"graph.png\" in the current directory. Let me know if you encounter any further issues.\n", + "This code will save the plot as a PNG file with the name \"plot.png\" in the same directory as your Python script. You can change the file name and format (e.g., \"plot.jpg\") if desired.\n", + "\n", + "Once you run this modified code, the plot will be saved as an image file. You can then use this file for further analysis or to share the plot with others.\n", "\n", "--------------------------------------------------------------------------------\n", "\u001b[31m\n", @@ -361,43 +364,110 @@ "exitcode: 1 (execution failed)\n", "Code output: \n", "Traceback (most recent call last):\n", - " File \"\", line 2, in \n", - " plt.savefig('graph.png')\n", - "NameError: name 'plt' is not defined\n", + " File \"\", line 4, in \n", + " plt.plot(x, y)\n", + "NameError: name 'x' is not defined\n", "\n", "\n", "--------------------------------------------------------------------------------\n", - "\u001b[33mTruncated 11 messages.\u001b[0m\n", - "\u001b[33mTruncated 453 tokens.\u001b[0m\n", + "\u001b[33mTruncated 5 messages. Reduced from 15 to 10.\u001b[0m\n", + "\u001b[33mTruncated 258 tokens. Tokens reduced from 622 to 364\u001b[0m\n", "\u001b[33massistant\u001b[0m (to user_proxy):\n", "\n", - "I apologize for the mistake once again. It seems that the matplotlib module is not installed on your system. To install it, you can use pip, the package manager for Python.\n", + "It seems that the code encountered an error because the variables `x` and `y` were not defined.\n", "\n", - "Here are the steps to install matplotlib:\n", + "Let's modify the code to define the variables `x` and `y` and then save the graph as an image file:\n", "\n", - "1. Open a command prompt or terminal.\n", - "2. Type the following command and press Enter:\n", + "```python\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", "\n", - "```sh\n", - "pip install matplotlib\n", + "# Define the x values as an array from -10 to 10\n", + "x = np.linspace(-10, 10, 100)\n", + "\n", + "# Define the y values as the square of x\n", + "y = x ** 2\n", + "\n", + "# Create the plot\n", + "plt.plot(x, y)\n", + "\n", + "# Set the plot title and axes labels\n", + "plt.title(\"Plot of x^2\")\n", + "plt.xlabel(\"x\")\n", + "plt.ylabel(\"y\")\n", + "\n", + "# Save the plot as an image file\n", + "plt.savefig('plot.png')\n", "```\n", + "This code uses the `numpy` library to generate an array of x values from -10 to 10 and calculates the corresponding y values as the square of x. The plot is then saved as an image file named `plot.png`.\n", "\n", - "3. Wait for the installation to complete.\n", + "Please make sure you have the `numpy` and `matplotlib` libraries installed before running this code. Let me know if you need help with that.\n", "\n", - "Once you have installed matplotlib, you can try running the code again to save the graph as an image file.\n", + "--------------------------------------------------------------------------------\n", + "\u001b[31m\n", + ">>>>>>>> EXECUTING CODE BLOCK 0 (inferred language is python)...\u001b[0m\n", + "\u001b[33muser_proxy\u001b[0m (to assistant):\n", + "\n", + "exitcode: 0 (execution succeeded)\n", + "Code output: \n", + "\n", + "\n", + "--------------------------------------------------------------------------------\n", + "\u001b[33mTruncated 7 messages. Reduced from 17 to 10.\u001b[0m\n", + "\u001b[33mTruncated 401 tokens. Tokens reduced from 765 to 364\u001b[0m\n", + "\u001b[33massistant\u001b[0m (to user_proxy):\n", + "\n", + "Great! The code executed successfully. The graph was saved as an image file with the name \"plot.png\" in the current working directory.\n", + "\n", + "Now you can share the \"plot.png\" file or view it using an image viewer.\n", + "\n", + "Let me know if there's anything else I can help with!\n", + "\n", + "--------------------------------------------------------------------------------\n", + "\u001b[33muser_proxy\u001b[0m (to assistant):\n", + "\n", + "\n", + "\n", + "--------------------------------------------------------------------------------\n", + "\u001b[33mTruncated 9 messages. Reduced from 19 to 10.\u001b[0m\n", + "\u001b[33mTruncated 282 tokens. Tokens reduced from 633 to 351\u001b[0m\n", + "\u001b[33massistant\u001b[0m (to user_proxy):\n", + "\n", + "Is there anything else I can help you with?\n", + "\n", + "--------------------------------------------------------------------------------\n", + "\u001b[33muser_proxy\u001b[0m (to assistant):\n", + "\n", + "\n", + "\n", + "--------------------------------------------------------------------------------\n", + "\u001b[33mTruncated 11 messages. Reduced from 21 to 10.\u001b[0m\n", + "\u001b[33mTruncated 342 tokens. Tokens reduced from 634 to 292\u001b[0m\n", + "\u001b[33massistant\u001b[0m (to user_proxy):\n", + "\n", + "If you need any further assistance, feel free to ask. I'm here to help!\n", "\n", "--------------------------------------------------------------------------------\n" ] + }, + { + "data": { + "text/plain": [ + "ChatResult(chat_history=[{'content': 'plot and save a graph of x^2 from -10 to 10', 'role': 'assistant'}, {'content': 'To plot and save a graph of the function x^2 from -10 to 10, you can use the matplotlib library in Python. Here is the code:\\n\\n```python\\nimport matplotlib.pyplot as plt\\nimport numpy as np\\n\\n# Generate x values from -10 to 10\\nx = np.linspace(-10, 10, 100)\\n\\n# Calculate corresponding y values (x^2)\\ny = x**2\\n\\n# Create the plot\\nplt.plot(x, y)\\n\\n# Add labels and title\\nplt.xlabel(\\'x\\')\\nplt.ylabel(\\'y\\')\\nplt.title(\\'Plot of x^2\\')\\n\\n# Save the plot as a file\\nplt.savefig(\\'x_squared_plot.png\\')\\n\\n# Show the plot\\nplt.show()\\n```\\n\\nThis code will create a plot of the function x^2 and save it as \"x_squared_plot.png\" in the current directory. Make sure you have the matplotlib library installed before running this code.', 'role': 'user'}, {'content': 'exitcode: 0 (execution succeeded)\\nCode output: \\nFigure(640x480)\\n', 'role': 'assistant'}, {'content': \"Great! The code executed successfully and generated a plot of the function x^2 from -10 to 10. The plot was displayed in a figure with size 640x480. \\n\\nTo save the graph as an image file, you can modify the code as follows:\\n\\n```python\\nimport matplotlib.pyplot as plt\\nimport numpy as np\\n\\n# Generate x values from -10 to 10\\nx = np.linspace(-10, 10, 100)\\n\\n# Generate y values by squaring x\\ny = x ** 2\\n\\n# Plot the graph\\nplt.plot(x, y)\\nplt.xlabel('x')\\nplt.ylabel('x^2')\\nplt.title('Graph of x^2')\\nplt.grid(True)\\n\\n# Save the graph as an image file, for example as 'graph.png'\\nplt.savefig('graph.png')\\n```\\n\\nBy executing this updated code, the graph will be saved as an image file named 'graph.png' in the same directory as your Python script.\\n\\nPlease let me know if you need any further assistance.\", 'role': 'user'}, {'content': 'exitcode: 0 (execution succeeded)\\nCode output: \\n', 'role': 'assistant'}, {'content': 'Great! The code executed successfully and generated a plot of the function x^2 from -10 to 10. The plot was displayed in a figure with size 640x480.\\n\\nTo save the graph as an image file, you can modify the code as follows:\\n\\n```python\\nimport matplotlib.pyplot as plt\\nimport numpy as np\\n\\n# Generate x values from -10 to 10\\nx = np.linspace(-10, 10, 100)\\n\\n# Calculate y values (x^2)\\ny = x**2\\n\\n# Plot the graph\\nplt.plot(x, y)\\n\\n# Add labels and title\\nplt.xlabel(\\'x\\')\\nplt.ylabel(\\'y\\')\\nplt.title(\\'Graph of x^2\\')\\n\\n# Save the graph as an image file\\nplt.savefig(\\'graph.png\\')\\n\\n# Close the plot\\nplt.close()\\n```\\n\\nThis code will save the plot as an image file named \"graph.png\" in the current directory. You can change the filename and path if needed.\\n\\nPlease let me know if you need any further assistance.', 'role': 'user'}, {'content': 'exitcode: 0 (execution succeeded)\\nCode output: \\n', 'role': 'assistant'}, {'content': \"Great! The code executed successfully and generated a plot of the function x^2 from -10 to 10. The plot was displayed in a figure with size 640x480.\\n\\nTo save the graph as an image file, you can modify the code as follows:\\n\\n```python\\nimport matplotlib.pyplot as plt\\nimport numpy as np\\n\\n# Generate x values\\nx = np.linspace(-10, 10, 100)\\n\\n# Generate y values\\ny = x**2\\n\\n# Plot the graph\\nplt.plot(x, y)\\n\\n# Set labels and title\\nplt.xlabel('x')\\nplt.ylabel('y')\\nplt.title('Graph of x^2')\\n\\n# Save the graph as an image file\\nplt.savefig('graph.png')\\n```\\n\\nThis code will save the graph as a PNG image file named 'graph.png' in the same directory where you are running the script. You can change the filename and format (e.g., 'graph.jpg') if desired.\\n\\nPlease note that you will need to have the matplotlib library installed to run this code. If you don't have it installed, you can install it by running `pip install matplotlib` in your terminal or command prompt.\\n\\nLet me know if you need any further assistance!\", 'role': 'user'}, {'content': 'exitcode: 0 (execution succeeded)\\nCode output: \\n', 'role': 'assistant'}, {'content': 'Great! The code executed successfully and generated a plot of the function x^2 from -10 to 10. The plot was displayed in a figure with size 640x480.\\n\\nTo save the graph as an image file, you can modify the code as follows:\\n\\n```python\\nimport matplotlib.pyplot as plt\\nimport numpy as np\\n\\n# Generate x values from -10 to 10\\nx = np.linspace(-10, 10, 100)\\n\\n# Generate y values by squaring x\\ny = x**2\\n\\n# Plot the graph\\nplt.plot(x, y)\\n\\n# Add labels and title\\nplt.xlabel(\\'x\\')\\nplt.ylabel(\\'x^2\\')\\nplt.title(\\'Graph of x^2\\')\\n\\n# Save the graph as an image file\\nplt.savefig(\\'x_squared_plot.png\\')\\n\\n# Display the graph\\nplt.show()\\n```\\n\\nThis code will save the graph as a PNG image file named \"x_squared_plot.png\" in the current working directory. You can customize the filename and file format according to your needs.', 'role': 'user'}, {'content': 'exitcode: 0 (execution succeeded)\\nCode output: \\nFigure(640x480)\\n', 'role': 'assistant'}, {'content': 'Great! The code executed successfully and generated a plot of the function x^2 from -10 to 10. The plot was displayed in a figure with size 640x480.\\n\\nTo save the graph as an image file, you can modify the code as follows:\\n\\n```python\\nimport matplotlib.pyplot as plt\\nimport numpy as np\\n\\n# Generate x values\\nx = np.linspace(-10, 10, 100)\\n\\n# Generate y values\\ny = x ** 2\\n\\n# Create the plot\\nplt.plot(x, y)\\nplt.xlabel(\\'x\\')\\nplt.ylabel(\\'y\\')\\nplt.title(\\'Graph of x^2\\')\\n\\n# Save the plot as an image file\\nplt.savefig(\\'x_squared.png\\')\\n\\nplt.show()\\n```\\n\\nThe code will save the plot as a PNG image file named \"x_squared.png\" in the current directory. You can change the filename or file extension as needed.\\n\\nPlease let me know if you need any further assistance!', 'role': 'user'}, {'content': 'exitcode: 0 (execution succeeded)\\nCode output: \\nFigure(640x480)\\n', 'role': 'assistant'}, {'content': 'To save the graph as an image file, you can modify the code as follows:\\n\\n```python\\nimport matplotlib.pyplot as plt\\n\\n# Create the plot\\nplt.plot(x, y)\\n\\n# Set the plot title and axes labels\\nplt.title(\"Plot of x^2\")\\nplt.xlabel(\"x\")\\nplt.ylabel(\"y\")\\n\\n# Save the plot as an image file\\nplt.savefig(\"plot.png\")\\n```\\n\\nThis code will save the plot as a PNG file with the name \"plot.png\" in the same directory as your Python script. You can change the file name and format (e.g., \"plot.jpg\") if desired.\\n\\nOnce you run this modified code, the plot will be saved as an image file. You can then use this file for further analysis or to share the plot with others.', 'role': 'user'}, {'content': 'exitcode: 1 (execution failed)\\nCode output: \\nTraceback (most recent call last):\\n File \"\", line 4, in \\n plt.plot(x, y)\\nNameError: name \\'x\\' is not defined\\n', 'role': 'assistant'}, {'content': 'It seems that the code encountered an error because the variables `x` and `y` were not defined.\\n\\nLet\\'s modify the code to define the variables `x` and `y` and then save the graph as an image file:\\n\\n```python\\nimport matplotlib.pyplot as plt\\nimport numpy as np\\n\\n# Define the x values as an array from -10 to 10\\nx = np.linspace(-10, 10, 100)\\n\\n# Define the y values as the square of x\\ny = x ** 2\\n\\n# Create the plot\\nplt.plot(x, y)\\n\\n# Set the plot title and axes labels\\nplt.title(\"Plot of x^2\")\\nplt.xlabel(\"x\")\\nplt.ylabel(\"y\")\\n\\n# Save the plot as an image file\\nplt.savefig(\\'plot.png\\')\\n```\\nThis code uses the `numpy` library to generate an array of x values from -10 to 10 and calculates the corresponding y values as the square of x. The plot is then saved as an image file named `plot.png`.\\n\\nPlease make sure you have the `numpy` and `matplotlib` libraries installed before running this code. Let me know if you need help with that.', 'role': 'user'}, {'content': 'exitcode: 0 (execution succeeded)\\nCode output: \\n', 'role': 'assistant'}, {'content': 'Great! The code executed successfully. The graph was saved as an image file with the name \"plot.png\" in the current working directory.\\n\\nNow you can share the \"plot.png\" file or view it using an image viewer.\\n\\nLet me know if there\\'s anything else I can help with!', 'role': 'user'}, {'content': '', 'role': 'assistant'}, {'content': 'Is there anything else I can help you with?', 'role': 'user'}, {'content': '', 'role': 'assistant'}, {'content': \"If you need any further assistance, feel free to ask. I'm here to help!\", 'role': 'user'}], summary=\"If you need any further assistance, feel free to ask. I'm here to help!\", cost=({'total_cost': 0.015855, 'gpt-3.5-turbo-0613': {'cost': 0.015855, 'prompt_tokens': 8242, 'completion_tokens': 1746, 'total_tokens': 9988}}, {'total_cost': 0.0147465, 'gpt-3.5-turbo-0613': {'cost': 0.0147465, 'prompt_tokens': 7755, 'completion_tokens': 1557, 'total_tokens': 9312}}), human_input=[])" + ] + }, + "execution_count": 4, + "metadata": {}, + "output_type": "execute_result" } ], "source": [ "assistant = autogen.AssistantAgent(\n", " \"assistant\",\n", - " llm_config={\n", - " \"config_list\": [{\"model\": \"gpt-3.5-turbo\", \"api_key\": \"YOUR_API_KEY\"}],\n", - " },\n", + " llm_config=llm_config,\n", ")\n", "\n", + "\n", "# Instantiate the capability to manage chat history\n", "manage_chat_history = context_handling.TransformChatHistory(max_tokens_per_message=50, max_messages=10, max_tokens=1000)\n", "# Add the capability to the assistant\n", @@ -430,7 +500,7 @@ }, { "cell_type": "code", - "execution_count": 4, + "execution_count": 5, "metadata": {}, "outputs": [ { @@ -452,40 +522,42 @@ "plot and save a graph of x^2 from -10 to 10\n", "\n", "--------------------------------------------------------------------------------\n", - "\u001b[33mTruncated 1991 messages.\u001b[0m\n", - "\u001b[33mTruncated 49800 tokens.\u001b[0m\n", + "\u001b[33mTruncated 1991 messages. Reduced from 2001 to 10.\u001b[0m\n", + "\u001b[33mTruncated 1000800 tokens. Tokens reduced from 1001015 to 215\u001b[0m\n", "\u001b[33massistant\u001b[0m (to user_proxy):\n", "\n", - "Here is the Python code to plot and save a graph of x^2 from -10 to 10:\n", + "Here's the Python code to plot and save a graph of x^2 from -10 to 10:\n", "\n", "```python\n", + "# filename: plot_graph.py\n", + "\n", "import matplotlib.pyplot as plt\n", "import numpy as np\n", "\n", - "# Generate values for x from -10 to 10\n", + "# Generate x values from -10 to 10\n", "x = np.linspace(-10, 10, 100)\n", "\n", - "# Calculate y values for x^2\n", - "y = x ** 2\n", + "# Calculate y values as x^2\n", + "y = x**2\n", "\n", - "# Plot the graph\n", + "# Create plot\n", "plt.plot(x, y)\n", + "\n", + "# Add labels and title\n", "plt.xlabel('x')\n", - "plt.ylabel('y = x^2')\n", + "plt.ylabel('y')\n", "plt.title('Graph of y = x^2')\n", "\n", - "# Save the graph as a PNG file\n", + "# Save the plot as a PNG image\n", "plt.savefig('graph.png')\n", "\n", - "# Close the plot\n", - "plt.close()\n", - "\n", - "print('Graph saved as graph.png')\n", + "# Show the plot\n", + "plt.show()\n", "```\n", "\n", - "Please make sure you have the `matplotlib` library installed before running this code. You can install it by running `pip install matplotlib` in your terminal.\n", + "To execute this code, save it to a file called `plot_graph.py` and run it using Python. This will generate a file called `graph.png` in the same directory, which will contain the graph of x^2 from -10 to 10.\n", "\n", - "After executing the code, a graph of y = x^2 will be saved as `graph.png` in your current working directory.\n", + "Note: Make sure you have the matplotlib library installed. You can install it by running `pip install matplotlib` in your terminal or command prompt.\n", "\n", "Let me know if you need any further assistance!\n", "\n", @@ -496,15 +568,17 @@ "\n", "exitcode: 0 (execution succeeded)\n", "Code output: \n", - "Graph saved as graph.png\n", + "Figure(640x480)\n", "\n", "\n", "--------------------------------------------------------------------------------\n", - "\u001b[33mTruncated 1993 messages.\u001b[0m\n", - "\u001b[33mTruncated 49850 tokens.\u001b[0m\n", + "\u001b[33mTruncated 1993 messages. Reduced from 2003 to 10.\u001b[0m\n", + "\u001b[33mTruncated 997232 tokens. Tokens reduced from 997466 to 234\u001b[0m\n", "\u001b[33massistant\u001b[0m (to user_proxy):\n", "\n", - "Great! The code executed successfully and the graph has been saved as \"graph.png\". You can now view the graph to see the plot of x^2 from -10 to 10. If you have any more questions or need further assistance, feel free to ask.\n", + "Great! The graph of x^2 from -10 to 10 has been plotted and saved successfully. You can find the saved graph as an image file on your computer. \n", + "\n", + "Is there anything else I can help you with?\n", "\n", "--------------------------------------------------------------------------------\n", "\u001b[33muser_proxy\u001b[0m (to assistant):\n", @@ -512,8 +586,8 @@ "\n", "\n", "--------------------------------------------------------------------------------\n", - "\u001b[33mTruncated 1995 messages.\u001b[0m\n", - "\u001b[33mTruncated 49900 tokens.\u001b[0m\n", + "\u001b[33mTruncated 1995 messages. Reduced from 2005 to 10.\u001b[0m\n", + "\u001b[33mTruncated 997096 tokens. Tokens reduced from 997326 to 230\u001b[0m\n", "\u001b[33massistant\u001b[0m (to user_proxy):\n", "\n", "TERMINATE\n", @@ -525,16 +599,12 @@ "source": [ "assistant_base = autogen.AssistantAgent(\n", " \"assistant\",\n", - " llm_config={\n", - " \"config_list\": [{\"model\": \"gpt-3.5-turbo\", \"api_key\": \"YOUR_API_KEY\"}],\n", - " },\n", + " llm_config=llm_config,\n", ")\n", "\n", "assistant_with_context_handling = autogen.AssistantAgent(\n", " \"assistant\",\n", - " llm_config={\n", - " \"config_list\": [{\"model\": \"gpt-3.5-turbo\", \"api_key\": \"YOUR_API_KEY\"}],\n", - " },\n", + " llm_config=llm_config,\n", ")\n", "# suppose this capability is not available\n", "manage_chat_history = context_handling.TransformChatHistory(max_tokens_per_message=50, max_messages=10, max_tokens=1000)\n", diff --git a/test/agentchat/contrib/capabilities/test_context_handling.py b/test/agentchat/contrib/capabilities/test_context_handling.py index f2ef4ea81..cc224d550 100644 --- a/test/agentchat/contrib/capabilities/test_context_handling.py +++ b/test/agentchat/contrib/capabilities/test_context_handling.py @@ -102,6 +102,129 @@ def test_transform_chat_history_with_agents(): assert False, f"Chat initiation failed with error {str(e)}" +def test_transform_messages(): + """ + Test transform_messages_retain_order() + """ + # Test case 1: Test that the order of messages is retained after transformation and Test that the messages are properly truncating. + messages = [ + {"role": "system", "content": "System message"}, + {"role": "user", "content": "Hi"}, + {"role": "user", "content": "user sending the 2nd test message"}, + {"role": "assistant", "content": "assistant sending the 3rd test message"}, + {"role": "assistant", "content": "assistant sending the 4th test message"}, + ] + + transform_chat_history = TransformChatHistory(max_messages=3, max_tokens_per_message=10, max_tokens=100) + transformed_messages = transform_chat_history._transform_messages(messages) + + assert transformed_messages[0]["role"] == "system" + assert transformed_messages[0]["content"] == "System message" + assert transformed_messages[1]["role"] == "user" + assert transformed_messages[1]["content"] == "user sending the 2nd test message" + assert transformed_messages[2]["role"] == "assistant" + assert transformed_messages[2]["content"] == "assistant sending the 3rd test message" + assert transformed_messages[3]["role"] == "assistant" + assert transformed_messages[3]["content"] == "assistant sending the 4th test message" + + # Test case 2: Test when no system message + messages = [ + {"role": "user", "content": "Hi"}, + {"role": "user", "content": "user sending the 2nd test message"}, + {"role": "assistant", "content": "assistant sending the 3rd test message"}, + {"role": "assistant", "content": "assistant sending the 4th test message"}, + ] + + transform_chat_history = TransformChatHistory(max_messages=3, max_tokens_per_message=10, max_tokens=100) + transformed_messages = transform_chat_history._transform_messages(messages) + + assert transformed_messages[0]["role"] == "user" + assert transformed_messages[0]["content"] == "user sending the 2nd test message" + assert transformed_messages[1]["role"] == "assistant" + assert transformed_messages[1]["content"] == "assistant sending the 3rd test message" + assert transformed_messages[2]["role"] == "assistant" + assert transformed_messages[2]["content"] == "assistant sending the 4th test message" + + messages = [ + {"role": "user", "content": "Out of max messages"}, + {"role": "assistant", "content": "first second third fourth"}, + {"role": "user", "content": "a"}, + ] + print(f"----Messages (N={len(messages)})----") + orignal_tokens = 0 + for i, msg in enumerate(messages): + print(f"[{msg['role']}-{i}]: {msg['content']}") + tokens = token_count_utils.count_token(msg["content"]) + print("Number of tokens: ", tokens) + orignal_tokens += tokens + print("-----Total tokens: ", orignal_tokens, "-----") + + allowed_max_tokens = 2 + transform_chat_history = TransformChatHistory(max_messages=2, max_tokens=allowed_max_tokens) + transformed_messages = transform_chat_history._transform_messages(messages) + + print("Max allowed tokens: ", allowed_max_tokens) + + print("Transformed contents") + for msg in transformed_messages: + print(msg["content"]) + print("Number of tokens: ", token_count_utils.count_token(msg["content"])) + assert len(transformed_messages) == 1 + assert transformed_messages[0]["role"] == "user" + + +def test_truncate_str_to_tokens(): + """ + Test the truncate_str_to_tokens function. + """ + from autogen.agentchat.contrib.capabilities.context_handling import truncate_str_to_tokens + + # Test case 1: Truncate string with fewer tokens than max_tokens + text = "This is a test" + max_tokens = 5 + truncated_text = truncate_str_to_tokens(text, max_tokens) + assert truncated_text == text + + # Test case 2: Truncate string with more tokens than max_tokens + text = "This is a test" + max_tokens = 3 + truncated_text = truncate_str_to_tokens(text, max_tokens) + assert truncated_text == "This is a" + + # Test case 3: Truncate empty string + text = "" + max_tokens = 5 + truncated_text = truncate_str_to_tokens(text, max_tokens) + assert truncated_text == "" + + # Test case 4: Truncate string with exact number of tokens as max_tokens + text = "This is a test" + max_tokens = 4 + truncated_text = truncate_str_to_tokens(text, max_tokens) + assert truncated_text == "This is a test" + + # Test case 5: Truncate string with no tokens found + text = "This is a test" + max_tokens = 0 + truncated_text = truncate_str_to_tokens(text, max_tokens) + assert truncated_text == "" + + # Test case 6: Truncate string when actual tokens are more than max_tokens + text = "This is a test with a looooooonngggg word" + max_tokens = 8 + truncated_text = truncate_str_to_tokens(text, max_tokens) + word_count = len(truncated_text.split()) + assert word_count <= max_tokens + + # Test case 7: Truncate string with exact number of tokens as max_tokens + text = "This\nis\na test" + max_tokens = 4 + truncated_text = truncate_str_to_tokens(text, max_tokens) + assert "This\nis" in truncated_text + + if __name__ == "__main__": test_transform_chat_history() test_transform_chat_history_with_agents() + test_truncate_str_to_tokens() + test_transform_messages()