phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

(test) add test token counter

Browse source
This commit is contained in:
ishaan-jaff 2023-11-13 07:42:08 -08:00
parent 62013520aa
commit e5ec4a92fe
1 changed files with 35 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

36
litellm/tests/test_utils.py
View file

@ -10,7 +10,7 @@ sys.path.insert(
) # Adds the parent directory to the system path
import pytest
import litellm
from litellm.utils import trim_messages, get_token_count, get_valid_models, check_valid_key, validate_environment, function_to_dict
from litellm.utils import trim_messages, get_token_count, get_valid_models, check_valid_key, validate_environment, function_to_dict, token_counter
# Assuming your trim_messages, shorten_message_to_fit_limit, and get_token_count functions are all in a module named 'message_utils'
@ -179,3 +179,37 @@ def test_function_to_dict():
print("passed")
# test_function_to_dict()
def test_token_counter():
messages = [
{
"role": "user",
"content": "hi how are you what time is it"
}
]
tokens = token_counter(
model = "gpt-3.5-turbo",
messages=messages
)
print(tokens)
tokens = token_counter(
model = "claude-2",
messages=messages
)
print("claude-2")
print(tokens)
tokens = token_counter(
model = "command-nightly",
messages=messages
)
print("command-nightly")
print(tokens)
# test_token_counter()