diff --git a/litellm/utils.py b/litellm/utils.py index 44dc75f45d7..d7dd553f97a 100644 --- a/litellm/utils.py +++ b/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 diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 466f3c988e4..6fa263e6a06 100644 --- a/tests/unit/test_utils.py +++ b/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 = [