phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-26 19:24:27 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Improved trimming logic and OpenAI token counter

Browse source
This commit is contained in:
Duc Pham 2023-11-10 01:26:13 +07:00
parent 880768f83d
commit c74e6f8cdd
2 changed files with 122 additions and 24 deletions
Download patch file Download diff file Expand all files Collapse all files

40
litellm/tests/test_utils.py
View file

@ -1,6 +1,6 @@
import sys, os
import traceback
from dotenv import load_dotenv
import copy
load_dotenv()
import os
@ -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 = trim_messages(messages=messages, model="gpt-3.5-turbo", max_tokens=20)
trimmed_messages = trim_messages(messages_copy=messages, model="gpt-3.5-turbo", max_tokens=20)
# print(get_token_count(messages=trimmed_messages, model="gpt-3.5-turbo"))
assert(get_token_count(messages=trimmed_messages, model="gpt-3.5-turbo")) <= 20
# test_multiple_messages_trimming()
@ -48,7 +48,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 = trim_messages(messages=messages, model="gpt-3.5-turbo", max_tokens=100)
trimmed_messages = trim_messages(messages_copy=messages, model="gpt-3.5-turbo", max_tokens=100)
print("Trimmed messages")
print(trimmed_messages)
assert(messages==trimmed_messages)
@ -56,14 +56,42 @@ def test_multiple_messages_no_trimming():
# test_multiple_messages_no_trimming()
def test_large_trimming():
def test_large_trimming_multiple_messages():
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 = trim_messages(messages, max_tokens=20, model="random")
trimmed_messages = trim_messages(messages, max_tokens=20, model="gpt-4-0613")
print("trimmed messages")
print(trimmed_messages)
assert(get_token_count(messages=trimmed_messages, model="random")) <= 20
assert(get_token_count(messages=trimmed_messages, model="gpt-4-0613")) <= 20
# test_large_trimming()
def test_large_trimming_single_message():
messages = [{"role": "user", "content": "This is a singlelongwordthatexceedsthelimit."}]
trimmed_messages = trim_messages(messages, max_tokens=5, model="gpt-4-0613")
assert(get_token_count(messages=trimmed_messages, model="gpt-4-0613")) <= 5
assert(get_token_count(messages=trimmed_messages, model="gpt-4-0613")) > 0
def test_trimming_with_system_message_within_max_tokens():
# This message is 33 tokens long
messages = [{"role": "system", "content": "This is a short system message"}, {"role": "user", "content": "This is a medium normal message, let's say litellm is awesome."}]
trimmed_messages = trim_messages(messages, max_tokens=30, model="gpt-4-0613") # The system message should fit within the token limit
assert len(trimmed_messages) == 2
assert trimmed_messages[0]["content"] == "This is a short system message"
def test_trimming_with_system_message_exceeding_max_tokens():
# This message is 33 tokens long. The system message is 13 tokens long.
messages = [{"role": "system", "content": "This is a short system message"}, {"role": "user", "content": "This is a medium normal message, let's say litellm is awesome."}]
trimmed_messages = trim_messages(messages, max_tokens=12, model="gpt-4-0613")
assert len(trimmed_messages) == 1
assert '..' in trimmed_messages[0]["content"]
def test_trimming_should_not_change_original_messages():
messages = [{"role": "system", "content": "This is a short system message"}, {"role": "user", "content": "This is a medium normal message, let's say litellm is awesome."}]
messages_copy = copy.deepcopy(messages)
trimmed_messages = trim_messages(messages, max_tokens=12, model="gpt-4-0613")
assert(messages==messages_copy)
def test_get_valid_models():
old_environ = os.environ
os.environ = {'OPENAI_API_KEY': 'temp'} # mock set only openai key in environ