diff --git a/litellm/utils.py b/litellm/utils.py index 3a26c3f071..c9c03ab271 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -6007,6 +6007,7 @@ def trim_messages( """ # Initialize max_tokens # if users pass in max tokens, trim to this amount + original_messages = messages messages = copy.deepcopy(messages) try: if max_tokens is None: @@ -6035,6 +6036,7 @@ def trim_messages( if message["role"] != "tool": break tool_messages.append(message) + tool_messages.reverse() # # Remove the collected tool messages from the original list if len(tool_messages): messages = messages[: -len(tool_messages)] @@ -6044,7 +6046,7 @@ def trim_messages( # Do nothing if current tokens under messages if current_tokens < max_tokens: - return messages + return messages + tool_messages #### Trimming messages if current_tokens > max_tokens print_verbose( @@ -6089,7 +6091,7 @@ def trim_messages( verbose_logger.exception( "Got exception while token trimming - {}".format(str(e)) ) - return messages + return original_messages from litellm.caching.in_memory_cache import InMemoryCache diff --git a/tests/litellm_utils_tests/test_utils.py b/tests/litellm_utils_tests/test_utils.py index a30c05c417..065c27eea9 100644 --- a/tests/litellm_utils_tests/test_utils.py +++ b/tests/litellm_utils_tests/test_utils.py @@ -1,4 +1,5 @@ import copy +import logging import sys import time from datetime import datetime @@ -238,11 +239,19 @@ def test_trimming_with_tool_calls(): "content": '{"location": "Paris", "temperature": "22", "unit": "celsius"}', }, ] - result = trim_messages(messages=messages, max_tokens=1, return_response_tokens=True) + num_tool_calls = 3 + + result = trim_messages(messages=messages, max_tokens=1) print(result) - assert len(result[0]) == 3 # final 3 messages are tool calls + # only trailing tool calls are returned + assert len(result) == num_tool_calls + assert result == messages[-num_tool_calls:] + + result = trim_messages(messages=messages, max_tokens=999) + # message length is below max_tokens, so output should match input + assert messages == result def test_trimming_should_not_change_original_messages(): @@ -274,6 +283,53 @@ def test_trimming_with_model_cost_max_input_tokens(model): ) +def test_trimming_with_untokenizable_field(caplog: pytest.LogCaptureFixture) -> None: + from litellm.types.utils import ChatCompletionMessageToolCall, Function, Message + + messages = [ + { + "role": "system", + "content": "You are a helpful assistant.", + }, + { + "role": "user", + "content": "What's the weather like in San Francisco?", + # non-string values will cause the tokenizer to raise an exception + "user_id": 123, + }, + Message( + content=None, + role="assistant", + tool_calls=[ + ChatCompletionMessageToolCall( + function=Function( + arguments='{"location": "San Francisco, CA", "unit": "celsius"}', + name="get_current_weather", + ), + id="call_G11shFcS024xEKjiAOSt6Tc9", + type="function", + ), + ], + function_call=None, + ), + { + "tool_call_id": "call_G11shFcS024xEKjiAOSt6Tc9", + "role": "tool", + "name": "get_current_weather", + "content": '{"location": "San Francisco", "temperature": "72", "unit": "fahrenheit"}', + }, + ] + + # trim_messages() catches the exception raised by the tokenizer and logs an error + with caplog.at_level(level=logging.ERROR, logger="LiteLLM"): + trimmed_messages = trim_messages(messages, max_tokens=999) + + assert trimmed_messages == messages + + assert len(caplog.records) >= 1 + assert "Got exception while token trimming" in caplog.text + + def test_aget_valid_models(): old_environ = os.environ os.environ = {"OPENAI_API_KEY": "temp"} # mock set only openai key in environ