tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-21 06:27:00 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix and rewrite of token_counter (#10409)

Browse Source 
* added tests

messages_with_counts: Made tolerance explicit for each test. But they match the new implementation(which beats the old)

* new token counter impl

* compare old and new implementation in test

* delete old token counter

* moved tests to /tests/litellm/litellm_core_utils

* use existing types

* docstrings

* warn about using default params on unknown model.

* created type for the token_counter_function

* check key == "content"

* throw error on invalid detail-type, ignore type-warning.

* fix imports
This commit is contained in:
Carlos Freund
2025-05-01 23:34:37 -07:00
committed by GitHub
parent 9cc39af131
commit cb177dbd7a
7 changed files with 743 additions and 381 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/litellm_core_utils/token_counter.py
+368 -2
View File
@@ -3,7 +3,9 @@
import base64
import io
import struct
from typing import Literal, Optional, Tuple, Union
from typing import Callable, List, Literal, Optional, Tuple, Union
import tiktoken
import litellm
from litellm import verbose_logger
@@ -16,13 +18,21 @@ from litellm.constants import (
MAX_TILE_HEIGHT,
MAX_TILE_WIDTH,
)
from litellm.litellm_core_utils.default_encoding import encoding as default_encoding
from litellm.llms.custom_httpx.http_handler import _get_httpx_client
from litellm.types.llms.openai import (
AllMessageValues,
ChatCompletionNamedToolChoiceParam,
ChatCompletionToolParam,
OpenAIMessageContent,
)
from litellm.types.utils import SelectTokenizerResponse
def get_modified_max_tokens(
model: str,
base_model: str,
messages: Optional[list],
messages: Optional[List[AllMessageValues]],
user_max_tokens: Optional[int],
buffer_perc: Optional[float],
buffer_num: Optional[float],
@@ -281,3 +291,359 @@ def calculate_img_tokens(
tile_tokens = (base_tokens * 2) * tiles_needed_high_res
total_tokens = base_tokens + tile_tokens
return total_tokens
TokenCounterFunction = Callable[[str], int]
"""
Type for a function that counts tokens in a string.
"""
class _MessageCountParams:
"""
A class to hold the parameters for counting tokens in messages.
"""
def __init__(
self,
model: str,
custom_tokenizer: Optional[Union[dict, SelectTokenizerResponse]],
):
from litellm.utils import print_verbose
actual_model = _fix_model_name(model)
if actual_model == "gpt-3.5-turbo-0301":
self.tokens_per_message = (
4 # every message follows <|start|>{role/name}\n{content}<|end|>\n
)
self.tokens_per_name = -1 # if there's a name, the role is omitted
elif actual_model in litellm.open_ai_chat_completion_models:
self.tokens_per_message = 3
self.tokens_per_name = 1
elif actual_model in litellm.azure_llms:
self.tokens_per_message = 3
self.tokens_per_name = 1
else:
print_verbose(f"Warning: unknown model {model}. Using default token params.")
self.tokens_per_message = 3
self.tokens_per_name = 1
self.count_function = _get_count_function(model, custom_tokenizer)
def token_counter(
model="",
custom_tokenizer: Optional[Union[dict, SelectTokenizerResponse]] = None,
text: Optional[Union[str, List[str]]] = None,
messages: Optional[List[AllMessageValues]] = None,
count_response_tokens: Optional[bool] = False,
tools: Optional[List[ChatCompletionToolParam]] = None,
tool_choice: Optional[ChatCompletionNamedToolChoiceParam] = None,
use_default_image_token_count: Optional[bool] = False,
default_token_count: Optional[int] = None,
) -> int:
"""
Count the number of tokens in a given text using a specified model.
Args:
model (str): The name of the model to use for tokenization. Default is an empty string.
custom_tokenizer (Optional[dict]): A custom tokenizer created with the `create_pretrained_tokenizer` or `create_tokenizer` method. Must be a dictionary with a string value for `type` and Tokenizer for `tokenizer`. Default is None.
text (str): The raw text string to be passed to the model. Default is None.
messages (Optional[List[AllMessageValues]]): Alternative to passing in text. A list of dictionaries representing messages with "role" and "content" keys. Default is None.
count_response_tokens (Optional[bool]): set to True to indicate we are processing a stream response.
tools (Optional[List[ChatCompletionToolParam]]): The available tools. Default is None.
tool_choice (Optional[ChatCompletionNamedToolChoiceParam]): The tool choice. Default is None.
use_default_image_token_count (Optional[bool]): When True, will NOT make a GET request to the image URL and instead return the default image dimensions. Default is False.
default_token_count (Optional[int]): The default number of tokens to return for a message block, if an error occurs. Default is None.
Returns:
int: The number of tokens in the text.
"""
if text is not None and messages is not None:
raise ValueError("text and messages cannot both be set")
if use_default_image_token_count is None:
use_default_image_token_count = False
if text is not None:
if tools or tool_choice:
raise ValueError("tools or tool_choice cannot be set if using text")
if isinstance(text, List):
text_to_count = "".join(t for t in text if isinstance(t, str))
elif isinstance(text, str):
text_to_count = text
count_function = _get_count_function(model, custom_tokenizer)
num_tokens = count_function(text_to_count)
elif messages is not None:
params = _MessageCountParams(model, custom_tokenizer)
num_tokens = _count_messages(
params, messages, use_default_image_token_count, default_token_count
)
if count_response_tokens is False:
includes_system_message = any(
[message.get("role", None) == "system" for message in messages]
)
num_tokens += _count_extra(
params.count_function, tools, tool_choice, includes_system_message
)
else:
raise ValueError("Either text or messages must be provided")
return num_tokens
def _count_messages(
params: _MessageCountParams,
messages: List[AllMessageValues],
use_default_image_token_count: bool,
default_token_count: Optional[int],
) -> int:
"""
Count the number of tokens in a list of messages.
Args:
params (_MessageCountParams): The parameters for counting tokens.
messages (List[AllMessageValues]): The list of messages to count tokens in.
use_default_image_token_count (bool): When True, will NOT make a GET request to the image URL and instead return the default image dimensions.
default_token_count (Optional[int]): The default number of tokens to return for a message block, if an error occurs.
"""
num_tokens = 0
for message in messages:
num_tokens += params.tokens_per_message
for key, value in message.items():
if value is None:
pass
elif key == "tool_calls":
if isinstance(value, List):
for tool_call in value:
if "function" in tool_call:
function_arguments = tool_call["function"].get(
"arguments", []
)
num_tokens += params.count_function(str(function_arguments))
else:
raise ValueError(
f"Unsupported tool call {tool_call} must contain a function key"
)
else:
raise ValueError(
f"Unsupported type {type(value)} for key tool_calls in message {message}"
)
elif isinstance(value, str):
num_tokens += params.count_function(value)
if key == "name":
num_tokens += params.tokens_per_name
elif key == 'content' and isinstance(value, List):
num_tokens += _count_content_list(
params.count_function,
value,
use_default_image_token_count,
default_token_count,
)
else:
raise ValueError(
f"Unsupported type {type(value)} for key {key} in message {message}"
)
return num_tokens
def _count_extra(
count_function: TokenCounterFunction,
tools: Optional[List[ChatCompletionToolParam]],
tool_choice: Optional[ChatCompletionNamedToolChoiceParam],
includes_system_message: bool,
) -> int:
"""Count extra tokens for function definitions and tool choices.
Args:
count_function (TokenCounterFunction): The function to count tokens.
tools (Optional[List[ChatCompletionToolParam]]): The available tools.
tool_choice (Optional[ChatCompletionNamedToolChoiceParam]): The tool choice.
includes_system_message (bool): Whether the messages include a system message.
"""
num_tokens = 3 # every reply is primed with <|start|>assistant<|message|>
if tools:
num_tokens += count_function(_format_function_definitions(tools))
num_tokens += 9 # Additional tokens for function definition of tools
# If there's a system message and tools are present, subtract four tokens
if tools and includes_system_message:
num_tokens -= 4
# If tool_choice is 'none', add one token.
# If it's an object, add 4 + the number of tokens in the function name.
# If it's undefined or 'auto', don't add anything.
if tool_choice == "none":
num_tokens += 1
elif isinstance(tool_choice, dict):
num_tokens += 7
num_tokens += count_function(str(tool_choice["function"]["name"]))
return num_tokens
def _get_count_function(
model: Optional[str],
custom_tokenizer: Optional[Union[dict, SelectTokenizerResponse]] = None,
) -> TokenCounterFunction:
"""
Get the function to count tokens based on the model and custom tokenizer."""
from litellm.utils import _select_tokenizer, print_verbose
if model is not None or custom_tokenizer is not None:
tokenizer_json = custom_tokenizer or _select_tokenizer(model) # type: ignore
if tokenizer_json["type"] == "huggingface_tokenizer":
def count_tokens(text: str) -> int:
enc = tokenizer_json["tokenizer"].encode(text)
return len(enc.ids)
elif tokenizer_json["type"] == "openai_tokenizer":
model_to_use = _fix_model_name(model) # type: ignore
try:
if "gpt-4o" in model_to_use:
encoding = tiktoken.get_encoding("o200k_base")
else:
encoding = tiktoken.encoding_for_model(model_to_use)
except KeyError:
print_verbose("Warning: model not found. Using cl100k_base encoding.")
encoding = tiktoken.get_encoding("cl100k_base")
def count_tokens(text: str) -> int:
return len(encoding.encode(text))
else:
raise ValueError("Unsupported tokenizer type")
else:
def count_tokens(text: str) -> int:
return len(default_encoding.encode(text, disallowed_special=()))
return count_tokens
def _fix_model_name(model: str) -> str:
"""We normalize some model names to others"""
if model in litellm.azure_llms:
# azure llms use gpt-35-turbo instead of gpt-3.5-turbo 🙃
return model.replace("-35", "-3.5")
elif model in litellm.open_ai_chat_completion_models:
return model # type: ignore
else:
return "gpt-3.5-turbo"
def _count_content_list(
count_function: TokenCounterFunction,
content_list: OpenAIMessageContent,
use_default_image_token_count: bool,
default_token_count: Optional[int],
) -> int:
"""
Get the number of tokens from a list of content.
"""
try:
num_tokens = 0
for c in content_list:
if isinstance(c, str):
num_tokens += count_function(c)
elif c["type"] == "text":
num_tokens += count_function(c["text"])
elif c["type"] == "image_url":
if isinstance(c["image_url"], dict):
image_url_dict = c["image_url"]
detail = image_url_dict.get("detail", "auto")
if detail not in ["low", "high", "auto"]:
raise ValueError(
f"Invalid detail value: {detail}. Expected 'low', 'high', or 'auto'."
)
url = image_url_dict.get("url")
num_tokens += calculate_img_tokens(
data=url,
mode=detail, # type: ignore
use_default_image_token_count=use_default_image_token_count,
)
elif isinstance(c["image_url"], str):
image_url_str = c["image_url"]
num_tokens += calculate_img_tokens(
data=image_url_str,
mode="auto",
use_default_image_token_count=use_default_image_token_count,
)
else:
raise ValueError(
f"Invalid image_url type: {type(c['image_url'])}. Expected str or dict."
)
else:
raise ValueError(
f"Invalid content type: {type(c)}. Expected str or dict."
)
return num_tokens
except Exception as e:
if default_token_count is not None:
return default_token_count
raise ValueError(
f"Error getting number of tokens from content list: {e}, default_token_count={default_token_count}"
)
def _format_function_definitions(tools):
"""Formats tool definitions in the format that OpenAI appears to use.
Based on https://github.com/forestwanglin/openai-java/blob/main/jtokkit/src/main/java/xyz/felh/openai/jtokkit/utils/TikTokenUtils.java
"""
lines = []
lines.append("namespace functions {")
lines.append("")
for tool in tools:
function = tool.get("function")
if function_description := function.get("description"):
lines.append(f"// {function_description}")
function_name = function.get("name")
parameters = function.get("parameters", {})
properties = parameters.get("properties")
if properties and properties.keys():
lines.append(f"type {function_name} = (_: {{")
lines.append(_format_object_parameters(parameters, 0))
lines.append("}) => any;")
else:
lines.append(f"type {function_name} = () => any;")
lines.append("")
lines.append("} // namespace functions")
return "\n".join(lines)
def _format_object_parameters(parameters, indent):
properties = parameters.get("properties")
if not properties:
return ""
required_params = parameters.get("required", [])
lines = []
for key, props in properties.items():
description = props.get("description")
if description:
lines.append(f"// {description}")
question = "?"
if required_params and key in required_params:
question = ""
lines.append(f"{key}{question}: {_format_type(props, indent)},")
return "\n".join([" " * max(0, indent) + line for line in lines])
def _format_type(props, indent):
type = props.get("type")
if type == "string":
if "enum" in props:
return " | ".join([f'"{item}"' for item in props["enum"]])
return "string"
elif type == "array":
# items is required, OpenAI throws an error if it's missing
return f"{_format_type(props['items'], indent)}[]"
elif type == "object":
return f"{{\n{_format_object_parameters(props, indent + 2)}\n}}"
elif type in ["integer", "number"]:
if "enum" in props:
return " | ".join([f'"{item}"' for item in props["enum"]])
return "number"
elif type == "boolean":
return "boolean"
elif type == "null":
return "null"
else:
# This is a guess, as an empty string doesn't yield the expected token count
return "any"
litellm/utils.py
+7 -299
View File
@@ -122,7 +122,6 @@ from litellm.litellm_core_utils.redact_messages import (
from litellm.litellm_core_utils.rules import Rules
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
from litellm.litellm_core_utils.token_counter import (
calculate_img_tokens,
get_modified_max_tokens,
)
from litellm.llms.bedrock.common_utils import BedrockModelInfo
@@ -266,6 +265,8 @@ from .types.llms.openai import (
ChatCompletionToolCallFunctionChunk,
)
from .types.router import LiteLLM_Params
from litellm.litellm_core_utils.token_counter import token_counter as token_counter_new
####### ENVIRONMENT VARIABLES ####################
# Adjust to your specific application needs / system capabilities.
@@ -759,7 +760,7 @@ def function_setup( # noqa: PLR0915
messages = "default-message-value"
stream = True if "stream" in kwargs and kwargs["stream"] is True else False
logging_obj = LiteLLMLogging(
model=model,
model=model, # type: ignore
messages=messages,
stream=stream,
litellm_call_id=kwargs["litellm_call_id"],
@@ -1607,98 +1608,6 @@ def decode(model="", tokens: List[int] = [], custom_tokenizer: Optional[dict] =
dec = tokenizer_json["tokenizer"].decode(tokens)
return dec
def openai_token_counter( # noqa: PLR0915
messages: Optional[list] = None,
model="gpt-3.5-turbo-0613",
text: Optional[str] = None,
is_tool_call: Optional[bool] = False,
tools: Optional[List[ChatCompletionToolParam]] = None,
tool_choice: Optional[ChatCompletionNamedToolChoiceParam] = None,
count_response_tokens: Optional[
bool
] = False, # Flag passed from litellm.stream_chunk_builder, to indicate counting tokens for LLM Response. We need this because for LLM input we add +3 tokens per message - based on OpenAI's token counter
use_default_image_token_count: Optional[bool] = False,
default_token_count: Optional[int] = None,
):
"""
Return the number of tokens used by a list of messages.
Borrowed from https://github.com/openai/openai-cookbook/blob/main/examples/How_to_count_tokens_with_tiktoken.ipynb.
"""
print_verbose(f"LiteLLM: Utils - Counting tokens for OpenAI model={model}")
try:
if "gpt-4o" in model:
encoding = tiktoken.get_encoding("o200k_base")
else:
encoding = tiktoken.encoding_for_model(model)
except KeyError:
print_verbose("Warning: model not found. Using cl100k_base encoding.")
encoding = tiktoken.get_encoding("cl100k_base")
if model == "gpt-3.5-turbo-0301":
tokens_per_message = (
4 # every message follows <|start|>{role/name}\n{content}<|end|>\n
)
tokens_per_name = -1 # if there's a name, the role is omitted
elif model in litellm.open_ai_chat_completion_models:
tokens_per_message = 3
tokens_per_name = 1
elif model in litellm.azure_llms:
tokens_per_message = 3
tokens_per_name = 1
else:
raise NotImplementedError(
f"""num_tokens_from_messages() is not implemented for model {model}. See https://github.com/openai/openai-python/blob/main/chatml.md for information on how messages are converted to tokens."""
)
num_tokens = 0
includes_system_message = False
if is_tool_call and text is not None:
# if it's a tool call we assembled 'text' in token_counter()
num_tokens = len(encoding.encode(text, disallowed_special=()))
elif messages is not None:
for message in messages:
num_tokens += tokens_per_message
if message.get("role", None) == "system":
includes_system_message = True
for key, value in message.items():
if isinstance(value, str):
num_tokens += len(encoding.encode(value, disallowed_special=()))
if key == "name":
num_tokens += tokens_per_name
elif isinstance(value, List):
text, num_tokens_from_list = _get_num_tokens_from_content_list(
content_list=value,
use_default_image_token_count=use_default_image_token_count,
default_token_count=default_token_count,
)
num_tokens += num_tokens_from_list
elif text is not None and count_response_tokens is True:
# This is the case where we need to count tokens for a streamed response. We should NOT add +3 tokens per message in this branch
num_tokens = len(encoding.encode(text, disallowed_special=()))
return num_tokens
elif text is not None:
num_tokens = len(encoding.encode(text, disallowed_special=()))
num_tokens += 3 # every reply is primed with <|start|>assistant<|message|>
if tools:
num_tokens += len(encoding.encode(_format_function_definitions(tools)))
num_tokens += 9 # Additional tokens for function definition of tools
# If there's a system message and tools are present, subtract four tokens
if tools and includes_system_message:
num_tokens -= 4
# If tool_choice is 'none', add one token.
# If it's an object, add 4 + the number of tokens in the function name.
# If it's undefined or 'auto', don't add anything.
if tool_choice == "none":
num_tokens += 1
elif isinstance(tool_choice, dict):
num_tokens += 7
num_tokens += len(encoding.encode(tool_choice["function"]["name"]))
return num_tokens
def create_pretrained_tokenizer(
identifier: str, revision="main", auth_token: Optional[str] = None
):
@@ -1741,118 +1650,6 @@ def create_tokenizer(json: str):
return {"type": "huggingface_tokenizer", "tokenizer": tokenizer}
def _format_function_definitions(tools):
"""Formats tool definitions in the format that OpenAI appears to use.
Based on https://github.com/forestwanglin/openai-java/blob/main/jtokkit/src/main/java/xyz/felh/openai/jtokkit/utils/TikTokenUtils.java
"""
lines = []
lines.append("namespace functions {")
lines.append("")
for tool in tools:
function = tool.get("function")
if function_description := function.get("description"):
lines.append(f"// {function_description}")
function_name = function.get("name")
parameters = function.get("parameters", {})
properties = parameters.get("properties")
if properties and properties.keys():
lines.append(f"type {function_name} = (_: {{")
lines.append(_format_object_parameters(parameters, 0))
lines.append("}) => any;")
else:
lines.append(f"type {function_name} = () => any;")
lines.append("")
lines.append("} // namespace functions")
return "\n".join(lines)
def _format_object_parameters(parameters, indent):
properties = parameters.get("properties")
if not properties:
return ""
required_params = parameters.get("required", [])
lines = []
for key, props in properties.items():
description = props.get("description")
if description:
lines.append(f"// {description}")
question = "?"
if required_params and key in required_params:
question = ""
lines.append(f"{key}{question}: {_format_type(props, indent)},")
return "\n".join([" " * max(0, indent) + line for line in lines])
def _format_type(props, indent):
type = props.get("type")
if type == "string":
if "enum" in props:
return " | ".join([f'"{item}"' for item in props["enum"]])
return "string"
elif type == "array":
# items is required, OpenAI throws an error if it's missing
return f"{_format_type(props['items'], indent)}[]"
elif type == "object":
return f"{{\n{_format_object_parameters(props, indent + 2)}\n}}"
elif type in ["integer", "number"]:
if "enum" in props:
return " | ".join([f'"{item}"' for item in props["enum"]])
return "number"
elif type == "boolean":
return "boolean"
elif type == "null":
return "null"
else:
# This is a guess, as an empty string doesn't yield the expected token count
return "any"
def _get_num_tokens_from_content_list(
content_list: List[Dict[str, Any]],
use_default_image_token_count: Optional[bool] = False,
default_token_count: Optional[int] = None,
) -> Tuple[str, int]:
"""
Get the number of tokens from a list of content.
Returns:
Tuple[str, int]: A tuple containing the text and the number of tokens.
"""
try:
num_tokens = 0
text = ""
for c in content_list:
if c["type"] == "text":
text += c["text"]
num_tokens += len(encoding.encode(c["text"], disallowed_special=()))
elif c["type"] == "image_url":
if isinstance(c["image_url"], dict):
image_url_dict = c["image_url"]
detail = image_url_dict.get("detail", "auto")
url = image_url_dict.get("url")
num_tokens += calculate_img_tokens(
data=url,
mode=detail,
use_default_image_token_count=use_default_image_token_count
or False,
)
elif isinstance(c["image_url"], str):
image_url_str = c["image_url"]
num_tokens += calculate_img_tokens(
data=image_url_str,
mode="auto",
use_default_image_token_count=use_default_image_token_count
or False,
)
return text, num_tokens
except Exception as e:
if default_token_count is not None:
return "", default_token_count
raise ValueError(
f"Error getting number of tokens from content list: {e}, default_token_count={default_token_count}"
)
def token_counter(
model="",
custom_tokenizer: Optional[Union[dict, SelectTokenizerResponse]] = None,
@@ -1865,100 +1662,11 @@ def token_counter(
default_token_count: Optional[int] = None,
) -> int:
"""
Count the number of tokens in a given text using a specified model.
Args:
model (str): The name of the model to use for tokenization. Default is an empty string.
custom_tokenizer (Optional[dict]): A custom tokenizer created with the `create_pretrained_tokenizer` or `create_tokenizer` method. Must be a dictionary with a string value for `type` and Tokenizer for `tokenizer`. Default is None.
text (str): The raw text string to be passed to the model. Default is None.
messages (Optional[List[Dict[str, str]]]): Alternative to passing in text. A list of dictionaries representing messages with "role" and "content" keys. Default is None.
default_token_count (Optional[int]): The default number of tokens to return for a message block, if an error occurs. Default is None.
Returns:
int: The number of tokens in the text.
The same as `litellm.litellm_core_utils.token_counter`.
Kept for backwards compatibility.
"""
# use tiktoken, anthropic, cohere, llama2, or llama3's tokenizer depending on the model
is_tool_call = False
num_tokens = 0
if text is None:
if messages is not None:
print_verbose(f"token_counter messages received: {messages}")
text = ""
for message in messages:
if message.get("content", None) is not None:
content = message.get("content")
if isinstance(content, str):
text += message["content"]
elif isinstance(content, List):
text, num_tokens = _get_num_tokens_from_content_list(
content_list=content,
use_default_image_token_count=use_default_image_token_count,
default_token_count=default_token_count,
)
if message.get("tool_calls"):
is_tool_call = True
for tool_call in message["tool_calls"]:
if "function" in tool_call:
function_arguments = tool_call["function"]["arguments"]
text = (
text if isinstance(text, str) else "".join(text or [])
) + (str(function_arguments) if function_arguments else "")
else:
raise ValueError("text and messages cannot both be None")
elif isinstance(text, List):
text = "".join(t for t in text if isinstance(t, str))
elif isinstance(text, str):
count_response_tokens = True # user just trying to count tokens for a text. don't add the chat_ml +3 tokens to this
if model is not None or custom_tokenizer is not None:
tokenizer_json = custom_tokenizer or _select_tokenizer(model=model)
if tokenizer_json["type"] == "huggingface_tokenizer":
enc = tokenizer_json["tokenizer"].encode(text)
num_tokens = len(enc.ids)
elif tokenizer_json["type"] == "openai_tokenizer":
if (
model in litellm.open_ai_chat_completion_models
or model in litellm.azure_llms
):
if model in litellm.azure_llms:
# azure llms use gpt-35-turbo instead of gpt-3.5-turbo 🙃
model = model.replace("-35", "-3.5")
print_verbose(
f"Token Counter - using OpenAI token counter, for model={model}"
)
num_tokens = openai_token_counter(
text=text, # type: ignore
model=model,
messages=messages,
is_tool_call=is_tool_call,
count_response_tokens=count_response_tokens,
tools=tools,
tool_choice=tool_choice,
use_default_image_token_count=use_default_image_token_count
or False,
default_token_count=default_token_count,
)
else:
print_verbose(
f"Token Counter - using generic token counter, for model={model}"
)
num_tokens = openai_token_counter(
text=text, # type: ignore
model="gpt-3.5-turbo",
messages=messages,
is_tool_call=is_tool_call,
count_response_tokens=count_response_tokens,
tools=tools,
tool_choice=tool_choice,
use_default_image_token_count=use_default_image_token_count
or False,
default_token_count=default_token_count,
)
else:
num_tokens = len(encoding.encode(text, disallowed_special=())) # type: ignore
return num_tokens
return token_counter_new(model, custom_tokenizer, text, messages, count_response_tokens, tools, tool_choice, use_default_image_token_count, default_token_count)
def supports_httpx_timeout(custom_llm_provider: str) -> bool:
tests/local_testing/messages_with_counts.py → tests/litellm/litellm_core_utils/messages_with_counts.py
+13 -5
View File
@@ -227,6 +227,7 @@ integer_enum = {
],
"tool_choice": "none",
"count": 54,
"count-tolerate": 56, # over by 2
}
@@ -255,6 +256,7 @@ integer_enum_tool_choice_name = {
"function": {"name": "data_demonstration"},
}, # 4 tokens for "data_demonstration"
"count": 64,
"count-tolerate": 66, # over by 2
}
no_parameters = {
@@ -316,6 +318,7 @@ no_parameter_description_or_required = {
],
"tool_choice": "auto",
"count": 49,
"count-tolerate": 50, # over by 1
}
no_parameter_description = {
@@ -339,6 +342,7 @@ no_parameter_description = {
],
"tool_choice": "auto",
"count": 49,
"count-tolerate": 50, # over by 1
}
string_enum = {
@@ -406,7 +410,8 @@ inner_object = {
}
],
"tool_choice": "none",
"count": 65, # counted 67, over by 2
"count": 65,
"count-tolerate" : 67 #over by 2
}
"""
namespace functions {
@@ -453,7 +458,8 @@ inner_object_with_enum_only = {
}
],
"tool_choice": "none",
"count": 73, # counted 74, over by 1
"count": 73,
"count-tolerate" : 74 #over by 1
}
"""
namespace functions {
@@ -504,7 +510,8 @@ inner_object_with_enum = {
}
],
"tool_choice": "none",
"count": 89, # counted 92, over by 3
"count": 89,
"count-tolerate" : 92, #over by 3
}
"""
namespace functions {
@@ -561,7 +568,8 @@ inner_object_and_string = {
}
],
"tool_choice": "none",
"count": 103, # counted 106, over by 3
"count": 103,
"count-tolerate" : 106, #over by 3
}
"""
namespace functions {
@@ -606,7 +614,7 @@ boolean = {
}
],
"tool_choice": "none",
"count": 89, # over by 3
"count": 89,
}
array = {
tests/local_testing/test_token_counter.py → tests/litellm/litellm_core_utils/test_token_counter.py
+105 -39
View File
@@ -1,5 +1,5 @@
#### What this tests ####
# This tests litellm.token_counter() function
# This tests litellm.token_counter.token_counter() function
import traceback
import os
import sys
@@ -8,6 +8,7 @@ from unittest.mock import MagicMock
import pytest
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
@@ -19,9 +20,10 @@ from litellm import (
decode,
encode,
get_modified_max_tokens,
token_counter,
token_counter as token_counter_old,
)
from large_text import text
from litellm.litellm_core_utils.token_counter import token_counter as token_counter_new
from tests.large_text import text
from messages_with_counts import (
MESSAGES_TEXT,
MESSAGES_WITH_IMAGES,
@@ -29,38 +31,48 @@ from messages_with_counts import (
)
def token_counter_both_assert_same(**args):
new = token_counter_new(**args)
old = token_counter_old(**args)
assert new == old, f"New token counter {new} does not match old token counter {old}"
return new
## Choose which token_counter the test will use.
#token_counter = token_counter_new
#token_counter = token_counter_old
token_counter = token_counter_both_assert_same
def test_token_counter_normal_plus_function_calling():
try:
messages = [
{"role": "system", "content": "System prompt"},
{"role": "user", "content": "content1"},
{"role": "assistant", "content": "content2"},
{"role": "user", "content": "conten3"},
{
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": "call_E0lOb1h6qtmflUyok4L06TgY",
"function": {
"arguments": '{"query":"search query","domain":"google.ca","gl":"ca","hl":"en"}',
"name": "SearchInternet",
},
"type": "function",
}
],
},
{
"tool_call_id": "call_E0lOb1h6qtmflUyok4L06TgY",
"role": "tool",
"name": "SearchInternet",
"content": "tool content",
},
]
tokens = token_counter(model="gpt-3.5-turbo", messages=messages)
print(f"tokens: {tokens}")
except Exception as e:
pytest.fail(f"An exception occurred - {str(e)}")
messages = [
{"role": "system", "content": "System prompt"},
{"role": "user", "content": "content1"},
{"role": "assistant", "content": "content2"},
{"role": "user", "content": "conten3"},
{
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": "call_E0lOb1h6qtmflUyok4L06TgY",
"function": {
"arguments": '{"query":"search query","domain":"google.ca","gl":"ca","hl":"en"}',
"name": "SearchInternet",
},
"type": "function",
}
],
},
{
"tool_call_id": "call_E0lOb1h6qtmflUyok4L06TgY",
"role": "tool",
"name": "SearchInternet",
"content": "tool content",
},
]
tokens = token_counter(model="gpt-3.5-turbo", messages=messages)
assert tokens == 80
# test_token_counter_normal_plus_function_calling()
@@ -76,6 +88,18 @@ def test_token_counter_textonly(message_count_pair):
)
assert counted_tokens == message_count_pair["count"]
@pytest.mark.parametrize(
"message_count_pair",
MESSAGES_TEXT,
)
def test_token_counter_count_response_tokens(message_count_pair):
counted_tokens = token_counter(
model="gpt-35-turbo", messages=[message_count_pair["message"]], count_response_tokens=True
)
# 3 tokens are not added because of count_response_tokens=True
expected = message_count_pair["count"] - 3
assert counted_tokens == expected
@pytest.mark.parametrize(
"message_count_pair",
@@ -100,11 +124,23 @@ def test_token_counter_with_tools(message_count_pair):
tool_choice=message_count_pair["tool_choice"],
)
expected_tokens = message_count_pair["count"]
diff = counted_tokens - expected_tokens
assert (
diff >= 0 and diff <= 3
), f"Expected {expected_tokens} tokens, got {counted_tokens}. Counted tokens is only allowed to be off by 3 in the over-counting direction."
actual_diff = counted_tokens - expected_tokens
if "count-tolerate" in message_count_pair:
if message_count_pair["count-tolerate"] == counted_tokens:
pass # expected
else:
tolerated_diff = message_count_pair["count-tolerate"] - expected_tokens
assert actual_diff <= tolerated_diff, f"Expected {expected_tokens} tokens, got {counted_tokens}. Counted tokens is only allowed to be off by {tolerated_diff} in the over-counting direction."
if actual_diff != tolerated_diff:
raise NeedsToleranceUpdateError(f"SOMETHING BROKEN GOT FIXED! THIS is good! Adjust 'count-tolerate' from {message_count_pair['count-tolerate']} to {counted_tokens}")
else:
assert expected_tokens == counted_tokens, f"Expected {expected_tokens} tokens, got {counted_tokens}."
class NeedsToleranceUpdateError(Exception):
"""Custom exception to mark tests that have improved"""
pass
def test_tokenizers():
try:
@@ -184,7 +220,7 @@ def test_encoding_and_decoding():
# llama2 encoding + decoding
llama2_tokens = encode(model="meta-llama/Llama-2-7b-chat", text=sample_text)
llama2_text = decode(
model="meta-llama/Llama-2-7b-chat", tokens=llama2_tokens.ids
model="meta-llama/Llama-2-7b-chat", tokens=llama2_tokens.ids # type: ignore
)
assert llama2_text == sample_text
@@ -350,6 +386,7 @@ def test_empty_tools():
print(result)
@pytest.mark.skip(reason="Skipping this test temporarily because it relies on a function being called that I am removing.")
def test_gpt_4o_token_counter():
with patch.object(
litellm.utils, "openai_token_counter", new=MagicMock()
@@ -383,6 +420,35 @@ def test_img_url_token_counter(img_url):
def test_token_encode_disallowed_special():
encode(model="gpt-3.5-turbo", text="Hello, world! <|endoftext|>")
def test_token_counter():
try:
messages = [{"role": "user", "content": "hi how are you what time is it"}]
tokens = token_counter(model="gpt-3.5-turbo", messages=messages)
print("gpt-35-turbo")
print(tokens)
assert tokens > 0
tokens = token_counter(model="claude-2", messages=messages)
print("claude-2")
print(tokens)
assert tokens > 0
tokens = token_counter(model="gemini/chat-bison", messages=messages)
print("gemini/chat-bison")
print(tokens)
assert tokens > 0
tokens = token_counter(model="ollama/llama2", messages=messages)
print("ollama/llama2")
print(tokens)
assert tokens > 0
tokens = token_counter(model="anthropic.claude-instant-v1", messages=messages)
print("anthropic.claude-instant-v1")
print(tokens)
assert tokens > 0
except Exception as e:
pytest.fail(f"Error occurred: {e}")
import unittest
from unittest.mock import patch, MagicMock
tests/litellm/litellm_core_utils/test_token_counter_tool.py
+84
View File
@@ -0,0 +1,84 @@
#### What this tests ####
# This tests litellm.token_counter() function
import os
import sys
import pytest
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
#Use the same token_counter as the main test.
from test_token_counter import token_counter
from test_token_counter_tool_data import *
@pytest.mark.parametrize(
"messages",
[
SHORT,
CONTENT_AND_TOOL_CALL,
SYSTEM_LONG,
TOOL_CALL_CONTENT_ARRAY,
],
ids=[
"SHORT",
"CONTENT_AND_TOOL_CALL",
"SYSTEM_LONG",
"TOOL_CALL_CONTENT_ARRAY",
],
)
def test_token_counter_tool_increases(messages):
conversation = []
prev_tokens = 0
for message in messages:
conversation.append(message)
tokens = token_counter(model="gpt-3.5-turbo", messages=conversation, tools=TOOLS) # type: ignore
print(f"tokens: {tokens}")
assert (
tokens > prev_tokens
), f"Token did not increase: {tokens} <= {prev_tokens}"
prev_tokens = tokens
# Reuse in multiple tests
@pytest.mark.parametrize("usermessage", USER_MESSAGES, ids=USER_MESSAGES_IDS)
@pytest.mark.parametrize("tool_call", TOOL_CALL_MESSAGES, ids=TOOL_CALL_MESSAGE_ids)
def test_grow(usermessage, tool_call):
assertGrow(usermessage, tool_call, False)
@pytest.mark.parametrize("usermessage", USER_MESSAGES, ids=USER_MESSAGES_IDS)
@pytest.mark.parametrize("tool_call", TOOL_CALL_MESSAGES, ids=TOOL_CALL_MESSAGE_ids)
def test_sum(usermessage, tool_call):
assertGrow(usermessage, tool_call, True)
def assertGrow(usermessage, tool_call, assertBiggerThanBoth=True):
print(usermessage, tool_call)
tokens_usermessage = token_counter(
model="anthropic.claude-instant-v1",
messages=[usermessage],
tools=TOOLS, # type: ignore
)
assert tokens_usermessage > 0
tokens_tool_call = token_counter(
model="anthropic.claude-instant-v1",
messages=[tool_call],
tools=TOOLS, # type: ignore
)
assert tokens_tool_call > 0
tokens_both = token_counter(
model="anthropic.claude-instant-v1",
messages=[usermessage, tool_call],
tools=TOOLS, # type: ignore
)
assert tokens_both > tokens_usermessage
assert tokens_both > tokens_tool_call
if assertBiggerThanBoth:
assert abs(tokens_usermessage + tokens_tool_call - tokens_both) <= 61, (
f"tokens_usermessage: {tokens_usermessage}, tokens_tool_call: {tokens_tool_call}, "
+ f"tokens_both: {tokens_both} diff: {tokens_usermessage + tokens_tool_call - tokens_both}"
)
tests/litellm/litellm_core_utils/test_token_counter_tool_data.py
+166
View File
@@ -0,0 +1,166 @@
SHORT = [
{"role": "system", "content": "System prompt"},
{
"tool_call_id": "call_E0lOb1h6qtmflUyok4L06TgY",
"role": "tool",
"name": "SearchInternet",
"content": "tool content",
},
]
CONTENT_AND_TOOL_CALL = [
{"role": "system", "content": "System prompt"},
{
"role": "assistant",
"content": "..",
"tool_calls": [
{
"id": "toolu_01NZjJ3e6fzhhbYYcb5k1k4d",
"type": "function",
"function": {
"name": "str_replace_editor",
"arguments": '{"command": "view","path": "/workspace"}',
},
}
],
},
]
_OPENHANDS_SYSTEM_MESSAGE = {
"content": [
{
"type": "text",
"text": "You are OpenHands agent, a helpful AI assistant that can "
"interact with a computer to solve tasks.\n\n<ROLE>\nYour primary "
"role is to assist users by executing commands, modifying code, and "
"solving technical problems effectively. You should be thorough, "
"methodical, and prioritize quality over speed.\n* If the user asks a "
"question, like 'why is X happening', don’t try to fix the problem. ",
}
],
"role": "system",
}
SYSTEM_LONG = [
_OPENHANDS_SYSTEM_MESSAGE,
{
"role": "assistant",
"tool_calls": [
{
"id": "toolu_01NZjJ3e6fzhhbYYcb5k1k4d",
"type": "function",
"function": {
"name": "str_replace_editor",
"arguments": '{"command": "view","path": "/workspace"}',
},
}
],
},
]
TOOL_CALL_CONTENT_ARRAY = [
_OPENHANDS_SYSTEM_MESSAGE,
{
"role": "assistant",
"content": [{"type": "text", "text": "very short"}],
"tool_calls": [
{
"id": "toolu_01NZjJ3e6fzhhbYYcb5k1k4d",
"type": "function",
"function": {
"name": "str_replace_editor",
"arguments": '{"command": "view","path": "/workspace"}',
},
}
],
},
]
USER_MESSAGES = [
{"role": "system", "content": "System prompt"},
{"content": [{"type": "text", "text": "2025-04-25 (UTC)"}], "role": "user"},
{
"role": "system",
"content": "System prompt long " + ("long" * 100),
},
{
"role": "system",
"content": ["System prompt long " + ("long" * 100)],
},
{
"role": "system",
"content": [{"type": "text", "text": "System prompt long " + ("long" * 100)}],
},
]
USER_MESSAGES_IDS = ["system", "user-arr", "long", "long-arr", "long-arr-obj"]
TOOL_CALL_MESSAGES = [
{
"role": "assistant",
"tool_calls": [
{
"id": "x",
"type": "function",
"function": {
"arguments": '{"arg":"value"}',
"name": "foo",
},
}
],
},
{
"role": "assistant",
"content": None,
"tool_calls": [
{
"id": "call_E0lOb1h6qtmflUyok4L06TgY",
"function": {
"arguments": '{"query":"search query","domain":"google.ca","gl":"ca","hl":"en"}',
"name": "SearchInternet",
},
"type": "function",
}
],
},
{"role": "user", "content": "Hi"},
{
"role": "assistant",
"content": [{"type": "text", "text": "very short"}],
"tool_calls": [
{
"id": "toolu_01NZjJ3e6fzhhbYYcb5k1k4d",
"type": "function",
"function": {
"name": "str_replace_editor",
"arguments": '{"command": "view","path": "/workspace"}',
},
}
],
},
]
TOOL_CALL_MESSAGE_ids = ["t_short", "t_normal", "usermsg", "t_content_array"]
TOOLS = [
{
"type": "function",
"function": {
"name": "data_demonstration",
"description": "This is the main function description",
"parameters": {
"type": "object",
"properties": {
"object_1": {
"type": "object",
"description": "The object data type as a property",
"properties": {
"string1": {"type": "string"},
},
}
},
"required": ["object_1"],
},
},
}
]
tests/litellm_utils_tests/test_utils.py
-36
View File
@@ -33,7 +33,6 @@ from litellm.utils import (
get_supported_openai_params,
get_token_count,
get_valid_models,
token_counter,
trim_messages,
validate_environment,
)
@@ -445,41 +444,6 @@ def test_function_to_dict():
# test_function_to_dict()
def test_token_counter():
try:
messages = [{"role": "user", "content": "hi how are you what time is it"}]
tokens = token_counter(model="gpt-3.5-turbo", messages=messages)
print("gpt-35-turbo")
print(tokens)
assert tokens > 0
tokens = token_counter(model="claude-2", messages=messages)
print("claude-2")
print(tokens)
assert tokens > 0
tokens = token_counter(model="gemini/chat-bison", messages=messages)
print("gemini/chat-bison")
print(tokens)
assert tokens > 0
tokens = token_counter(model="ollama/llama2", messages=messages)
print("ollama/llama2")
print(tokens)
assert tokens > 0
tokens = token_counter(model="anthropic.claude-instant-v1", messages=messages)
print("anthropic.claude-instant-v1")
print(tokens)
assert tokens > 0
except Exception as e:
pytest.fail(f"Error occurred: {e}")
# test_token_counter()
@pytest.mark.parametrize(
"model, expected_bool",
[