test(utils): cover trim_messages return_response_tokens when trimming fails

Also say in the docstring which limit response_tokens is measured against.
This commit is contained in:
taltal-beep 2026-10-11 09:57:31 +03:00
parent bdfc1930df
commit 811768a38d
2 changed files with 18 additions and 2 deletions

View file

@ -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

View file

@ -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 = [