[Bug Fix] Always include tool calls in output of trim_messages (#11517)

* Check content and order of trimmed messages

* Assert tool calls are preserved if below max_tokens

* Unreverse order of tool calls

* Return tool calls alongside other messages

* Write test for trimming untokenizable field

* Return original messages in case of exception

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

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