Merge 811768a38d into c0ca4a51f7

litellm/utils.py

@@ -7692,7 +7692,7 @@ def trim_messages(
     model: str | None = None,
     trim_ratio: float = DEFAULT_TRIM_RATIO,
     return_response_tokens: bool = False,
-    max_tokens=None,
+    max_tokens: int | None = None,
 ):
     """
     Trim a list of messages to fit within a model's token limit.
@@ -7705,7 +7705,10 @@ def trim_messages(
         max_tokens: Instead of specifying a model or trim_ratio, you can specify this directly.
 
     Returns:
-        Trimmed messages and optionally the number of tokens available for response.
+        Trimmed messages. If return_response_tokens is True, a `(messages, response_tokens)` tuple instead,
+        whether or not anything had to be trimmed. `response_tokens` is what is left of the limit the messages
+        were trimmed to (`max_tokens`, or the model's max input tokens times `trim_ratio`). It is None when
+        there is no token limit to measure against (unknown model and no `max_tokens`) or trimming failed.
     """
     # Initialize max_tokens
     # if users pass in max tokens, trim to this amount
@@ -7723,7 +7726,9 @@ def trim_messages(
                 # if user did not specify max (input) tokens
                 # or passed an llm litellm does not know
                 # do nothing, just return messages
-                return messages
+                return (messages, None) if return_response_tokens else messages
+        # the caller's limit: `max_tokens` has the system message deducted from it further down
+        token_limit: Final = max_tokens
 
         system_message = ""
         for message in messages:
@@ -7748,7 +7753,14 @@ def trim_messages(
 
         # Do nothing if current tokens under messages
         if current_tokens < max_tokens:
-            return messages + tool_messages
+            untrimmed_messages: Final[list[AllMessageValues]] = (  # mutable-ok: handed back to the caller as a list
+                messages + tool_messages
+            )
+            return (
+                (untrimmed_messages, token_limit - get_token_count(untrimmed_messages, model))
+                if return_response_tokens
+                else untrimmed_messages
+            )
 
         #### Trimming messages if current_tokens > max_tokens
         print_verbose(
@@ -7761,7 +7773,8 @@ def trim_messages(
             )
 
             if max_tokens == 0:  # the system messages are too long
-                return [system_message_event]
+                # the system message alone uses the whole limit, so nothing is left for the response
+                return ([system_message_event], 0) if return_response_tokens else [system_message_event]
 
             # Since all system messages are combined and trimmed to fit the max_tokens,
             # we remove all system messages from the messages list
@@ -7780,12 +7793,12 @@ def trim_messages(
 
         verbose_logger.debug("Final messages: %s, return_response_tokens: %s", final_messages, return_response_tokens)
         if return_response_tokens:  # if user wants token count with new trimmed messages
-            response_tokens: Final = max_tokens - get_token_count(final_messages, model)
+            response_tokens: Final = token_limit - get_token_count(final_messages, model)
             return final_messages, response_tokens
         return final_messages
     except Exception as e:  # [NON-Blocking, if error occurs just return final_messages
         verbose_logger.exception("Got exception while token trimming - %s", e)
-        return original_messages
+        return (original_messages, None) if return_response_tokens else original_messages
 
 
 from litellm.caching.in_memory_cache import InMemoryCache

tests/unit/test_utils.py

@@ -7683,6 +7683,77 @@ def test_trimming_should_not_change_original_messages():
     assert messages == messages_copy
 
 
+def test_trim_messages_return_response_tokens_returns_a_tuple_when_nothing_is_trimmed():
+    messages = [
+        {"role": "system", "content": "This is a short system message"},
+        {"role": "user", "content": "hi"},
+    ]
+
+    result = trim_messages(messages, model="gpt-4-0613", max_tokens=1000, return_response_tokens=True)
+
+    assert isinstance(result, tuple)
+    trimmed_messages, response_tokens = result
+    assert trimmed_messages == messages
+    assert response_tokens == 1000 - get_token_count(messages, model="gpt-4-0613")
+
+
+def test_trim_messages_return_response_tokens_counts_the_system_message_once():
+    max_tokens = 100
+    messages = [
+        {"role": "system", "content": "This is a short system message"},
+        {"role": "user", "content": "This is a long message that will not fit. " * 100},
+        {"role": "user", "content": "This is the last question"},
+    ]
+
+    trimmed_messages, response_tokens = trim_messages(
+        messages, model="gpt-4-0613", max_tokens=max_tokens, return_response_tokens=True
+    )
+
+    assert trimmed_messages[0]["role"] == "system"
+    assert response_tokens == max_tokens - get_token_count(trimmed_messages, model="gpt-4-0613")
+    assert response_tokens >= 0
+
+
+def test_trim_messages_return_response_tokens_returns_a_tuple_when_the_system_message_uses_the_whole_limit():
+    max_tokens = 30
+    messages = [
+        {"role": "system", "content": "This is a system message that is far too long. " * 50},
+        {"role": "user", "content": "hi"},
+    ]
+
+    result = trim_messages(messages, model="gpt-4-0613", max_tokens=max_tokens, return_response_tokens=True)
+
+    assert isinstance(result, tuple)
+    trimmed_messages, response_tokens = result
+    assert response_tokens == max_tokens - get_token_count(trimmed_messages, model="gpt-4-0613")
+
+
+def test_trim_messages_return_response_tokens_is_none_when_no_token_limit_is_known():
+    messages = [{"role": "user", "content": "hi"}]
+
+    trimmed_messages, response_tokens = trim_messages(
+        messages, model="a-model-litellm-does-not-know", return_response_tokens=True
+    )
+
+    assert trimmed_messages == messages
+    assert response_tokens is None
+
+
+def test_trim_messages_return_response_tokens_returns_a_tuple_when_trimming_fails(monkeypatch):
+    def failing_token_count(*args, **kwargs):
+        raise RuntimeError("tokenizer unavailable")
+
+    monkeypatch.setattr(litellm.utils, "get_token_count", failing_token_count)
+    messages = [{"role": "user", "content": "hi"}]
+
+    assert trim_messages(messages, max_tokens=1000) == messages
+
+    trimmed_messages, response_tokens = trim_messages(messages, max_tokens=1000, return_response_tokens=True)
+
+    assert trimmed_messages == messages
+    assert response_tokens is None
+
+
 @pytest.mark.parametrize("model", ["gpt-5.4-mini", "claude-sonnet-4-6"])
 def test_trimming_with_model_cost_max_input_tokens(model):
     messages = [