diff --git a/docs/my-website/docs/completion/message_trimming.md b/docs/my-website/docs/completion/message_trimming.md index 2da9ade6a..e3f0204d8 100644 --- a/docs/my-website/docs/completion/message_trimming.md +++ b/docs/my-website/docs/completion/message_trimming.md @@ -1,25 +1,25 @@ # Trimming Input Messages -**Use litellm.safe_messages() to ensure messages does not exceed a model's token limit or specified `max_tokens`** +**Use litellm.trim_messages() to ensure messages does not exceed a model's token limit or specified `max_tokens`** ## Usage ```python from litellm import completion -from litellm.utils import safe_messages +from litellm.utils import trim_messages response = completion( model=model, - messages=safe_messages(messages, model) # safe_messages ensures tokens(messages) < max_tokens(model) + messages=trim_messages(messages, model) # trim_messages ensures tokens(messages) < max_tokens(model) ) ``` ## Usage - set max_tokens ```python from litellm import completion -from litellm.utils import safe_messages +from litellm.utils import trim_messages response = completion( model=model, - messages=safe_messages(messages, model, max_tokens=10), # safe_messages ensures tokens(messages) < max_tokens + messages=trim_messages(messages, model, max_tokens=10), # trim_messages ensures tokens(messages) < max_tokens ) ``` diff --git a/litellm/tests/test_utils.py b/litellm/tests/test_utils.py index cd96f2848..4b93fac02 100644 --- a/litellm/tests/test_utils.py +++ b/litellm/tests/test_utils.py @@ -10,14 +10,14 @@ sys.path.insert( ) # Adds the parent directory to the system path import pytest import litellm -from litellm.utils import safe_messages, get_token_count +from litellm.utils import trim_messages, get_token_count # Assuming your trim_messages, shorten_message_to_fit_limit, and get_token_count functions are all in a module named 'message_utils' # Test 1: Check trimming of normal message def test_basic_trimming(): messages = [{"role": "user", "content": "This is a long message that definitely exceeds the token limit."}] - trimmed_messages = safe_messages(messages, model="claude-2", max_tokens=8) + trimmed_messages = trim_messages(messages, model="claude-2", max_tokens=8) print("trimmed messages") print(trimmed_messages) # print(get_token_count(messages=trimmed_messages, model="claude-2")) @@ -26,7 +26,7 @@ test_basic_trimming() def test_basic_trimming_no_max_tokens_specified(): messages = [{"role": "user", "content": "This is a long message that is definitely under the token limit."}] - trimmed_messages = safe_messages(messages, model="gpt-4") + trimmed_messages = trim_messages(messages, model="gpt-4") print("trimmed messages for gpt-4") print(trimmed_messages) # print(get_token_count(messages=trimmed_messages, model="claude-2")) @@ -38,7 +38,7 @@ def test_multiple_messages_trimming(): {"role": "user", "content": "This is a long message that will exceed the token limit."}, {"role": "user", "content": "This is another long message that will also exceed the limit."} ] - trimmed_messages = safe_messages(messages=messages, model="gpt-3.5-turbo", max_tokens=20) + trimmed_messages = trim_messages(messages=messages, model="gpt-3.5-turbo", max_tokens=20) print("Trimmed messages") print(trimmed_messages) # print(get_token_count(messages=trimmed_messages, model="gpt-3.5-turbo")) @@ -50,7 +50,7 @@ def test_multiple_messages_no_trimming(): {"role": "user", "content": "This is a long message that will exceed the token limit."}, {"role": "user", "content": "This is another long message that will also exceed the limit."} ] - trimmed_messages = safe_messages(messages=messages, model="gpt-3.5-turbo", max_tokens=100) + trimmed_messages = trim_messages(messages=messages, model="gpt-3.5-turbo", max_tokens=100) print("Trimmed messages") print(trimmed_messages) assert(messages==trimmed_messages) @@ -60,7 +60,7 @@ test_multiple_messages_no_trimming() def test_large_trimming(): messages = [{"role": "user", "content": "This is a singlelongwordthatexceedsthelimit."}, {"role": "user", "content": "This is a singlelongwordthatexceedsthelimit."},{"role": "user", "content": "This is a singlelongwordthatexceedsthelimit."},{"role": "user", "content": "This is a singlelongwordthatexceedsthelimit."},{"role": "user", "content": "This is a singlelongwordthatexceedsthelimit."}] - trimmed_messages = safe_messages(messages, max_tokens=20, model="random") + trimmed_messages = trim_messages(messages, max_tokens=20, model="random") print("trimmed messages") print(trimmed_messages) assert(get_token_count(messages=trimmed_messages, model="random")) <= 20 diff --git a/litellm/utils.py b/litellm/utils.py index 5fc70f673..1a36fd918 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2496,7 +2496,7 @@ def shorten_message_to_fit_limit( # LiteLLM token trimmer # this code is borrowed from https://github.com/KillianLucas/tokentrim/blob/main/tokentrim/tokentrim.py # Credits for this code go to Killian Lucas -def safe_messages( +def trim_messages( messages, model = None, system_message = None, # str of user system message