From bdfc1930df1cd5b752ad2817b8aa280732686ba9 Mon Sep 17 00:00:00 2001 From: taltal-beep Date: Sat, 10 Oct 2026 23:18:52 +0300 Subject: [PATCH 1/2] fix(utils): make trim_messages return_response_tokens always return a tuple with a correct count With return_response_tokens=True, trim_messages returned a (messages, tokens) tuple only when it actually trimmed. When the messages already fit, when the system message used the whole limit, when the model was unknown, or when trimming raised, it returned the bare message list, so `messages, tokens = trim_messages(...)` raised or, for a two-message conversation, silently unpacked the two messages. When it did trim a conversation with a system message, the count was measured against a limit that already had the system message deducted, so the system message was subtracted twice and the result could be negative. Every path now returns the tuple when asked, measured against the caller's limit; response_tokens is None when there is no limit to measure against. --- litellm/utils.py | 26 ++++++++++++++----- tests/unit/test_utils.py | 56 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 75 insertions(+), 7 deletions(-) diff --git a/litellm/utils.py b/litellm/utils.py index e7a171f56c6..97590fafe64 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -7686,7 +7686,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. @@ -7699,7 +7699,9 @@ 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 None when there is no token limit to + measure against (unknown model and no `max_tokens`). """ # Initialize max_tokens # if users pass in max tokens, trim to this amount @@ -7717,7 +7719,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: @@ -7742,7 +7746,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( @@ -7755,7 +7766,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 @@ -7774,12 +7786,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 diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 7a192d17347..29f8180fb5e 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -7570,6 +7570,62 @@ 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 + + @pytest.mark.parametrize("model", ["gpt-5.4-mini", "claude-sonnet-4-6"]) def test_trimming_with_model_cost_max_input_tokens(model): messages = [ From 811768a38d08d0e7c852ebff7557d4cf745c5fc1 Mon Sep 17 00:00:00 2001 From: taltal-beep Date: Sun, 11 Oct 2026 09:57:31 +0300 Subject: [PATCH 2/2] test(utils): cover trim_messages return_response_tokens when trimming fails Also say in the docstring which limit response_tokens is measured against. --- litellm/utils.py | 5 +++-- tests/unit/test_utils.py | 15 +++++++++++++++ 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/litellm/utils.py b/litellm/utils.py index 97590fafe64..899f1cf80f3 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -7700,8 +7700,9 @@ def trim_messages( Returns: 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 None when there is no token limit to - measure against (unknown model and no `max_tokens`). + 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 diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 29f8180fb5e..e4c9849ac4e 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -7626,6 +7626,21 @@ def test_trim_messages_return_response_tokens_is_none_when_no_token_limit_is_kno 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 = [