mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-14 04:09:09 +00:00
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:
parent
bdfc1930df
commit
811768a38d
2 changed files with 18 additions and 2 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 = [
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue